text:
###############################################################################
#
# PEinvestment.jl solves the Benchmark PE Firm Investment Model
#
# Shane McMiken
# November 2021
# Boston College
#
################################################################################

#--------------------------------#
#         PREAMBLE
#--------------------------------#
using Plots # plots
using Distributions # cdf(normal())
using Parameters  # macros: @with_kw & @unpack
using DelimitedFiles # writedlm

#--------------------------------#
#         PARAMETERIZATIONS
#--------------------------------#

@with_kw struct NoAC
  α::Float64    = 0.7
  p::Float64    = 1.2
  δ::Float64    = 0.1
  r::Float64    = 0.04
  ρ::Float64    = 0.8
  σ::Float64    = 0.2
  γ_c::Float64  = 0.0
  γ_f::Float64  = 0.0
end

@with_kw struct ConvexAC
  α::Float64    = 0.7
  p::Float64    = 1.2
  δ::Float64    = 0.1
  r::Float64    = 0.04
  ρ::Float64    = 0.8
  σ::Float64    = 0.2
  γ_c::Float64  = 2.5
  γ_f::Float64  = 0.0
end

@with_kw struct FixedAC
  α::Float64    = 0.7
  p::Float64    = 1.2
  δ::Float64    = 0.1
  r::Float64    = 0.04
  ρ::Float64    = 0.8
  σ::Float64    = 0.2
  γ_c::Float64  = 0.0
  γ_f::Float64  = 0.05
end

#--------------------------------#
#  DESCRETIZE AR(1) PROCESS
#--------------------------------#

# Tauchen (1986)
function MyTauchen(σ_a::Float64,ρ_a::Float64,na::Int64,m::Float64)

  # Discrete markov process approximation for AR(1)
  # y = ρ_a * y_  + σ_a * ϵ
  # m: number of standard deviations
  # σ_a: standard deviations
  # ρ_a: AR(1) parameter
  # na: number of grid points

  vecsize = na
  σ_y = σ_a / sqrt(1-(ρ_a^2)) # Stationary distribution variance
  step = 2 * m * σ_y / (vecsize-1)
  agrid  = [-m*σ_y + (i-1)*step for i∈1:na]

  # Calculate transition probabilities, AR(1) - Tauchen
  P = zeros(na, na)

  mm = agrid[2] - agrid[1]; # re-compute step
  for j = 1:na
    for k = 1:na
      if(k == 1) # first column
        P[j, k] = cdf(Normal(), (agrid[k] - ρ_a*agrid[j] + (mm/2))/σ_a)
      elseif(k == na) # second-last column
        P[j, k] = 1 - cdf(Normal(), (agrid[k] - ρ_a*agrid[j] - (mm/2))/σ_a)
      else
        P[j, k] = cdf(Normal(), (agrid[k] - ρ_a*agrid[j] + (mm/2))/σ_a) - cdf(Normal(), (agrid[k] - ρ_a*agrid[j] - (mm/2))/σ_a)
      end
    end #k
  end #j

  # agrid = exp.(agrid) # Exponentiate grid

  return agrid, P #return grid and Markov transition matrix
end

#--------------------------------#
#  VALUE FUNCTION ITERATION
#--------------------------------#
# RECURSIVE PROBLEM:
#
# V(z,k) = max_k' {R(z,k,k') + (1/(1+r))* E[V(z',k')|z] }
#
# WHERE R(z,k,k') = zk^α - pk + p(1-δ)k - AC(z,k,k')

struct state
  ind::Int64
  nz::Int64 # number of prod grid points
  nk::Int64 # number of capital grid points
  P::Array
  zgrid::Vector{Float64} # technology grid
  kgrid::Vector{Float64} # capital grid
  r::Float64 # interest rate rate
  δ::Float64 # Depreication rate
  α::Float64 # Capital share of income
  p::Float64 # price of investment goods
  γ_c::Float64 # variable adjustment cost
  γ_f::Float64 # fixed adjustment cost
  V::Array
end

function value(currentstate::state)

  ind     = currentstate.ind
  nz      = currentstate.nz
  nk      = currentstate.nk
  P       = currentstate.P
  zgrid   = currentstate.zgrid
  kgrid   = currentstate.kgrid
  r       = currentstate.r
  δ       = currentstate.δ
  α       = currentstate.α
  p       = currentstate.p
  γ_c     = currentstate.γ_c
  γ_f     = currentstate.γ_f
  VV      = currentstate.V

  # get states
  kk      = convert(Int, floor((ind-0.05)/nz))+1
  zz      = convert(Int, floor(mod(ind-0.05,nz))+1)
  kc      = kgrid[kk] # current capital
  zc      = zgrid[zz] # current technology

  # compute continuation value
  cont    = P * VV
  cont    = cont[zz,:]

  # compute return function
  R       = zc*kc^α .- p.*kgrid .+ p*(1.0-δ)*kc .- γ_c/2.0 .* ((kgrid .- (1.0-δ)*kc)./kc).^2 .*kc - γ_f.*zc.*kc^α.*(kgrid.>(1-δ)*kc)

  # compute new value funciton - if return is less than zero than value is -Inf
  VVp     = R .+ (1.0+r).^(-1) .* cont - Inf*(R.<0)

  # find maximum
  Vnew, tempidx = findmax(vec(VVp))

  return Vnew, tempidx
end

function vfi(Vbar::Float64, zgrid::Array, kgrid::Array, P::Array, γ_c, γ_f, α, r, δ, p ; crit = Inf, maxiter = 10000, iter = 0, tol = 1.0e-6, nfix = 10)

  nz = length(zgrid)
  nk = length(kgrid)
  V_init = fill(Vbar, (nz,nk))
  V_next = Array{Float64,1}
  index = ones(nz*nk,)

  # useful grids
  kkgr = vec(repeat(kgrid',nz))
  zzgr = repeat(zgrid,nk,1)

  while (crit > tol && iter < maxiter)

    V_next = copy(V_init)

    if mod(iter,nfix) == 0 # Howard acceleration

      for ind = 1:(nz*nk)

        modelstate = state(ind,nz,nk,P,zgrid,kgrid,r,δ,α,p,γ_c,γ_f,V_next)
        tempVal, tempIndex  = value(modelstate)

        V_next[ind] = tempVal
        index[ind] = tempIndex
      end

    else # on mod(iter,nfix) != 0, update the value function only.

      # Re-use conjectured index
      idx = convert(Array{Int64,1}, index)

      # E(V) condition on A(t) and each possible K(t+1)
      EVp = P*V_init
      evp_k = zeros(nk*nz,)
      for ind = 1:(nz*nk)
        kk = convert(Int, floor((ind-0.05)/nz))+1
        zz = convert(Int, floor(mod(ind-0.05,nz))+1)
        evp_k[ind] = EVp[zz,idx[ind]]
      end

      # Compute return function using conjectured policy - idx
      R = zzgr.*kkgr.^α .- p.*kgrid[idx].+ p.*(1.0-δ).*kkgr .- γ_c/2.0 .* ((kgrid[idx] .- (1.0-δ)*kkgr)./kkgr).^2 .*kkgr - γ_f.*zzgr.*kkgr.^α.*(kgrid[idx].>(1-δ)*kkgr)

      # Re-compute value function under conjectured optimal policy
      V_next = R + (1.0+r).^(-1) .* evp_k - Inf*(R.<0)

    end

    # check tolerance
    V_next = reshape(V_next,nz,nk)
    crit = maximum(abs.(V_next .- V_init))
    V_init = V_next
    iter += 1

    #print("iteration: ",iter, ", critical val: ", crit*1.0e5," times 10^5", "\n")
  end

  index = convert(Vector{Int64}, index)
  return V_next, index, crit
end

#------------------------------------------#
#  VFI FAST - VECTORIZED
#------------------------------------------#

function vfivec(Vbar::Float64, zgrid::Array, kgrid::Array, P::Array, γ_c, γ_f, α, r, δ, p ; crit = Inf, maxiter = 10000, iter = 0, tol = 1.0e-8, nfix = 10)

  nz = length(zgrid)
  nk = length(kgrid)
  V_init = fill(Vbar, (nz,nk))
  V_next = Array{Float64,1}
  index = Array{CartesianIndex{2},2}

  # useful grids
  kkgr = vec(repeat(kgrid',nz))
  zzgr = repeat(zgrid,nk,1)

  # compute model objects
  R = zzgr.*kkgr.^α .- p*kgrid' .+ p*(1.0-δ)*kkgr .- γ_c/2.0 .* ((kgrid' .- (1.0-δ)*kkgr)./kkgr).^2 .*kkgr - γ_f.*zzgr.*kkgr.^α.*(kgrid'.>(1-δ)*kkgr)

  while (crit > tol && iter < maxiter)

    if mod(iter,nfix) == 0 # Howard acceleration

      # compute (nz*nk*nk)-dimensional value function
      obj = R + (1.0/(1.0+r)).*repeat(P*V_init,nk) - Inf*(R.<0)

      # Get maximum
      V_next, index = findmax(obj, dims = 2)

    else # on mod(iter,nfix) != 0, update the value function only.

      # Re-use conjectured index
      idx = [index[i][2] for i∈1:(nz*nk)]

      # E(V) condition on A(t) and each possible K(t+1)
      EVp = P*V_init
      evp_k = zeros(nk*nz,)
      for ind = 1:(nz*nk)
        zz = convert(Int, floor(mod(ind-0.05,nz))+1)
        evp_k[ind] = EVp[zz,idx[ind]]
      end

      # Compute return function using conjectured policy - idx
      R_idx = zzgr.*kkgr.^α .- p.*kgrid[idx].+ p.*(1.0-δ).*kkgr .- γ_c/2.0 .* ((kgrid[idx] .- (1.0-δ)*kkgr)./kkgr).^2 .*kkgr - γ_f.*zzgr.*kkgr.^α.*(kgrid[idx].>(1-δ)*kkgr)

      # Re-compute value function under conjectured optimal policy
      V_next = R_idx + (1.0+r).^(-1) .* evp_k - Inf*(R_idx.<0)

    end

    # check tolerance
    V_next = reshape(V_next,nz,nk)
    crit = maximum(abs.(V_next .- V_init))
    V_init = copy(V_next)
    iter += 1

    #print("iteration: ",iter, ", critical val: ", crit*1.0e5," times 10^5", "\n")
  end

  index = [index[i][2] for i∈1:(nz*nk)]
  return V_init, index, crit

end

#------------------------------------------#
#  ERGODIC DISTRIBUTION - Young (2010)
#------------------------------------------#

function young(index::Array, P::Array, nz::Int64, nk::Int64 ; crit = Inf, maxiter = 10000, iter = 0, tol = 1.0e-6)

  # initialise mass
  μ_init = fill((1/(nz*nk)), (nz,nk))

  while (crit > tol && iter < maxiter)

    # intialise
    μ_next = zeros(nz,nk)

    # use policies from vfi to infer masses next period
    for ind = 1:(nz*nk)
      kk = convert(Int, floor((ind-0.05)/nz))+1
      zz = convert(Int, floor(mod(ind-0.05,nz))+1)
      μ_next[zz,index[ind]] =  μ_next[zz,index[ind]] + P[:,zz]'*μ_init[:,kk]
    end

    # check tolerance
    crit = maximum(abs.(μ_init .- μ_next))
    μ_init = μ_next
    iter += 1

    #print("iteration: ",iter, ", critical val: ", crit*1.0e5," times 10^5", "\n")

  end

  return μ_init, crit, iter
end

#------------------------------------------#
#  EXTRACT AGGREGATES FROM THE MODEL
#------------------------------------------#

function moments(μ::Matrix,zgrid::Vector,kgrid::Vector,index::Vector, valfunc::Array, α, δ)

  nz = length(zgrid)
  nk = length(kgrid)

  #plot policy function
  pols = reshape(index,(nz,nk)) # policies

  # Get useful grids
  kkgr = log.(repeat(kgrid',nz))
  zzgr = log.(repeat(zgrid,1,nk))
  yygr = log.(exp.(zzgr).*exp.(kkgr).^α)
  iigr = (kgrid[pols]-(1-δ).*exp.(kkgr))./exp.(kkgr)

  # mean
  X = [kkgr, zzgr, yygr, iigr]
  EX = zeros(size(X),)
  for (ii, x) in enumerate(X)
    EX[ii] = sum(x.*μ)
  end

  #Std deviation
  STDX = zeros(size(X),)
  for (ii, x) in enumerate(X)
    STDX[ii] = sqrt.(sum(x.^2 .*μ) - EX[ii].^2)
  end

  # correlation matrix
  CORRX = zeros(size(X)[1],size(X)[1])
  for (ii,xi) in enumerate(X)
    for (jj,xj) in enumerate(X)
      CORRX[ii,jj] = (sum(xi.*xj.*μ) - EX[ii]*EX[jj])/(STDX[ii]*STDX[jj])
    end
  end

  return EX, CORRX

end

#------------------------------------------#
#  OBTAIN PLOTS
#------------------------------------------#

# solve model
para = [NoAC, ConvexAC, FixedAC]
strpara = ["NoAC", "ConvexAC","FixedAC"]

for (ii,str) in enumerate(para)

  @unpack α, p, δ, r, ρ, σ, γ_c, γ_f= str()

  #steady state
  kbar = (α/(p*(r+δ)))^(1.0/(1.0-α))
  Vbar = ((1.0+r)/r)*(kbar^α - δ*kbar)

  # technology grid
  cover = 5
  nz    = 2*cover+1
  m     = 2.5
  zgrid, P = MyTauchen(σ,ρ,nz,m)
  zgrid = exp.(zgrid) # exponentiate grid

  # capital grid
  nk = 100
  kmax = ((maximum(zgrid)*α)/(p*(r+δ)))^(1.0/(1.0-α))
  kgrid = [kmax*(1-δ)^(nk-i-1) for i∈0:(nk-1)]

  #@time valfunc, index, crit = vfi(Vbar,zgrid,kgrid,P) #slower
  @time valfunc, index, crit = vfivec(Vbar, zgrid, kgrid, P, γ_c, γ_f, α, r, δ, p)
  @time μ_ergodic, crit, iterfin = young(index,P,nz,nk)

  ## (a)
  #plot ergodic distribution
  stationarydist = heatmap(kgrid,zgrid,μ_ergodic, title = "Stationary distribution - $(strpara[ii])", ylabel = "z", xlabel = "k")
  distribution = plot(stationarydist, xlabel = "K", ylabel = "Z")
  savefig(distribution,"./src/stationarydistribution_$(strpara[ii]).png")
  display(distribution)

  #plot policy function
  pols = reshape(index,(nz,nk)) # policies
  zlow = pols[cover-2,:] # fix productivity at low
  zmid = pols[cover+1,:] # fix productivity at steady state
  zhigh = pols[cover+4,:] # fix productivity high

  p1 = plot(kgrid,kgrid[zlow], ylabel = "low productivity", title = "k policy - $(strpara[ii])")
  p2 = plot(kgrid,kgrid[zmid], ylabel = "mid productivity")
  p3 = plot(kgrid,kgrid[zhigh], ylabel = "high productivity")
  p4 = heatmap(kgrid,zgrid,pols, ylabel = "Z", xlabel = "K")

  policyfunctions = plot(p1,p2,p3,p4, layout = (2,2), legend = false)
  savefig(policyfunctions,"./src/policyfunctions_$(strpara[ii]).png")
  display(policyfunctions)

  #plot value function
  p5 = plot(kgrid,valfunc[cover-2,zlow], ylabel = "low productivity", title = "Value Function - $(strpara[ii])")
  p6 = plot(kgrid,valfunc[cover+1,zlow], ylabel = "mid productivity")
  p7 = plot(kgrid,valfunc[cover+4,zlow], ylabel = "high productivity")
  p8 = heatmap(kgrid,zgrid,valfunc, ylabel = "Z", xlabel = "K")

  valuefunction = plot(p5,p6,p7,p8, layout = (2,2), legend = false)
  savefig(valuefunction,"./src/valuefunction_$(strpara[ii]).png")
  display(valuefunction)

  ## (b) marginal Distributions
  marginal_k = sum(μ_ergodic, dims = 1)
  marginal_z = sum(μ_ergodic, dims = 2)
  p9 = plot(zgrid, marginal_z, ylabel = "probability density", xlabel = "productivity" )
  p10 = plot(kgrid, marginal_k', xlabel = "capital" )

  marginaldistributions = plot(p9,p10, layout = (1,2), title = "$(strpara[ii])" , legend = false)
  savefig(marginaldistributions,"./src/marginaldistributions_$(strpara[ii]).png")
  display(marginaldistributions)

  ## (c) moments

  EX, CORRX = moments(μ_ergodic, zgrid, kgrid, index, valfunc, α, δ)
  writedlm( "./src/correlationmatrix_$(strpara[ii]).csv", round.([EX' ; CORRX]; digits = 2), ',')

  ## (d) aggregates
  kkgr  = vec(repeat(kgrid',nz))
  zzgr  = repeat(zgrid,nk,1)
  μ_vec = vec(μ_ergodic)
  Kagg  = sum(μ_vec.*kkgr)
  Zagg  = sum(μ_vec.*zzgr)
  Yagg  = sum(μ_vec.*zzgr.*kkgr.^α)
  Iagg  = sum(μ_vec.*(kgrid[index]-(1-δ).*kkgr))
  Vagg  = sum(μ_vec.*vec(valfunc))
  TFP   = Yagg/Kagg^α
  writedlm("./src/aggregates_$(strpara[ii]).csv", round.([Yagg Iagg Kagg Vagg Zagg TFP]; digits = 2), ',')

  ## (f) implied TFP is larger than average of true underlying productivity because larger (more productive) firms produce more.

  ## (g)
  #find coordinates of maximal mass
  val, cord = findmax(μ_ergodic)
  p11 = plot(zgrid./kgrid[cord[2]],kgrid[pols[:,cord[2]]], ylabel = "k'(z=1,k)", xlabel = "z/k", title = "k(z=1,k) - $(strpara[ii]) ")
  eqpolicy = plot(p11, legend = false)
  savefig(eqpolicy,"./src/eqpolicy_$(strpara[ii]).png")
  display(eqpolicy)

end

max_stars_repo_name: shanemcmiken/PEinvestment

max_issues_repo_name: shanemcmiken/PEinvestment

max_forks_repo_name: shanemcmiken/PEinvestment


text:
using CSV
using StatsBase
using Distributions
using NPZ
using ArgParse

using StatisticalVerificationOfIsingEnergies
import StatisticalVerificationOfIsingEnergies: bootstrap_hists_of_mins, squared_error

function estimate_min(file::String)

    D = npzread(file)
    α = 0.19

    ens = D["energies01"]

    S = 1_000

    l = size(ens, 2)

    min_data = [minimum(ens[:,i]) for i in 1:l]
    eta = [skewness(ens[:,i]) for i in 1:l]
    min_estimated = [estimate_ground_state_energy(ens[:,i], α) for i in 1:l]
    p_values = [bootstrap_get_pvalue(ens[:,i], α, S) for i in 1:l]
    p_values_10 = [bootstrap_get_pvalue(ens[:,i], 0.10, S) for i in 1:l]
    p_values_39 = [bootstrap_get_pvalue(ens[:,i], 0.39, S) for i in 1:l]

    ys = [bootstrap_hists_of_mins(ens[:,i], α, S) for i in 1:l]
    q95 = [quantile(y, 0.95) for y in ys]
    bootstrap_std = [std(y) for y in ys]

    squared_error_from_cums = [squared_error(ens[:,i], α) for i in 1:l]

    e_min = D["ground"]

    betas = [estiamte_temperature(ens[:,i], e_min) for i in 1:l]

    push!(D, "eta" => eta)
    push!(D, "alpha" => α)
    push!(D, "minimum_from_data" => min_data)
    push!(D, "minimum_estimated" => min_estimated)
    push!(D, "p_values" => p_values)
    push!(D, "p_values_10" => p_values_10)
    push!(D, "p_values_39" => p_values_39)
    push!(D, "estimated_betas" => betas)
    push!(D, "bootstrap_std" => bootstrap_std)
    push!(D, "bootstrap_q95" => q95)
    push!(D, "squared_error" => squared_error_from_cums)

    outfile = replace(file, "input_data" => "output")
    npzwrite(outfile, D)
end

s = ArgParseSettings("description")
  @add_arg_table! s begin
  "file"
  arg_type = String
  help = "file with data to be plotted"
end

file = parse_args(s)["file"]
estimate_min(file)

ext: jl


text:
function get_volumes(setup)
    if isa(setup, CylinderSetup)
        dim = 2
        get_vol = r -> π * r^2 * setup.height
    else
        dim = 3
        get_vol = r -> 4 * π / 3 * r^3
    end

    # Get spherical volumes
    volumes = get_vol.(r * 1e6)

    # Subtract inner volumes (each volume is contained within the next)
    volumes[2:end] = volumes[2:end] - volumes[1:end-1]

    volumes
end


text:
# Deform domain by bending and twisting.
function deform_domain!(points, bend, twist)

    height = maximum(points[3, :]) - minimum(points[3, :])
    width = maximum(points[1, :]) - minimum(points[1, :])

    # Ben points in x-direction as a percentage of height.
    zcenter = (maximum(points[3, :]) + minimum(points[3, :])) / 2
    @. points[1, :] += bend * 30.0 * width * ((points[3, :] - zcenter) / height)^2.0

    # Twist angles by height
    thvec = points[[3], :] / height * twist

    # Twist around center
    center = (maximum(points[1:2, :], dims = 2) .+ minimum(points[1:2, :], dims = 2)) ./ 2
    points[1:2, :] .-= center

    points[1:2, :] .= [
        cos.(thvec) .* points[[1], :] .- sin.(thvec) .* points[[2], :]
        sin.(thvec) .* points[[1], :] .+ cos.(thvec) .* points[[2], :]
    ]

    # points[1, :] = points[1, :] - rmax
    points[1:2, :] .+= center

    nothing
end


text:
function CTOH(ts) #This is Convert to one hot
    try
        ts = broadcast(Int, ts)
    catch
        throw("Perhaps the array of objective variables you passed has a non-zero decimal point in the name of the class (as in 1.1). Please change it to an integer type.")
    end
    if 0 in ts
        ts .+= 1
    end
    c = maximum(ts)
    oh = zeros(length(ts), c)
    i = 1
    for t in ts
        oh[i, t] = 1.0
        i += 1
    end
    return oh
end


text:
#=============================================
Generate 𝐹 and ∇ₓ𝐹 from user input
=============================================#

"""
    F, ∇ₓF = state_function_and_Jacobian(Ts, Gs, nb)

Returns the state function `F` and its jacobian, `∇ₓF`.
"""
function state_function_and_Jacobian(Ts::Tuple, Gs::Tuple, nb)
    nt = length(Ts)
    tracers(x) = state_to_tracers(x, nb, nt)
    tracer(x,i) = state_to_tracer(x, nb, nt, i)
    T(p) = blockdiag([Tⱼ(p) for Tⱼ in Ts]...) # Big T (linear part)
    function G(x,p)
        xs = tracers(x)
        return reduce(vcat, [Gⱼ(xs..., p) for Gⱼ in Gs]) # nonlinear part
    end
    F(x,p) = G(x,p) - reduce(vcat, [Tⱼ(p) * tracer(x,j) for (j,Tⱼ) in enumerate(Ts)])
    ∇ₓG(x,p) = local_jacobian(Gs, x, p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓG(x,p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, ∇ₓF
end
function state_function_and_Jacobian(T, G)
    F(x,p) = G(x,p) - T(p) * x                     # full 𝐹(𝑥) = -T 𝑥 + 𝐺(𝑥)
    ∇ₓG(x,p) = sparse(Diagonal(localderivative(G, x, p))) # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓG(x,p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, ∇ₓF
end

"""
    localderivative(G, x, p)
    localderivative(Gᵢ, xs, i, p)
    localderivative(Gᵢ, dx, xs, i, p)

Returns the "local" derivative of `G` (or `Gᵢ`), i.e., equivalent to the vector

```
∇ₓG(x,p) * ones(size(x))
```

but using ForwardDiff's Jacobian instead.
"""
function localderivative(G, x, p) # for single tracer
    return ForwardDiff.derivative(λ -> G(x .+ λ, p), 0.0)
end
function localderivative(Gᵢ, xs, j, p) # for multiple tracers
    return ForwardDiff.derivative(λ -> Gᵢ(perturb_tracer(xs,j,λ)..., p), 0.0)
end
function localderivative(Gᵢ!, dx, xs, j, p) # if Gᵢ are in-place
    return ForwardDiff.derivative((dx,λ) -> Gᵢ!(dx, perturb_tracer(xs,j,λ)..., p), dx, 0.0)
end
perturb_tracer(xs, j, λ) = (xs[1:j-1]..., xs[j] .+ λ, xs[j+1:end]...)

function inplace_state_function_and_Jacobian(Ts::Tuple, Gs::Tuple, nb)
    nt = length(Ts)
    tracers(x) = state_to_tracers(x, nb, nt)
    tracer(x,i) = state_to_tracer(x, nb, nt, i)
    T(p) = blockdiag([Tⱼ(p) for Tⱼ in Ts]...) # Big T (linear part)
    #F(x,p) = G(x,p) - T(p) * x                     # full 𝐹(𝑥) = -T 𝑥 + 𝐺(𝑥)
    function F!(dx,x,p)
        xs = tracers(x)
        for (j, (Tⱼ, Gⱼ!)) in enumerate(zip(Ts, Gs))
            ij = tracer_indices(nb,nt,j)
            @views dx[ij] .= Gⱼ!(dx[ij], xs..., p)
            @views dx[ij] .-= Tⱼ(p) * x[ij]
        end
        return dx
    end
    ∇ₓG(x,p) = inplace_local_jacobian(Gs, x, p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓG(x,p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F!, ∇ₓF
end

export state_function_and_Jacobian, inplace_state_function_and_Jacobian

"""
    F, ∇ₓF = state_function_and_Jacobian(Ts, Gs, nb)

Returns the state function `F` and its jacobian, `∇ₓF`.
"""
function split_state_function_and_Jacobian(Ts::Tuple, Ls::Tuple, NLs::Tuple, nb)
    nt = length(Ts)
    tracers(x) = state_to_tracers(x, nb, nt)
    T(p) = blockdiag([Tⱼ(p) for Tⱼ in Ts]...) # Big T (linear part)
    NL(x,p) = reduce(vcat, [NLⱼ(tracers(x)..., p) for NLⱼ in NLs]) # nonlinear part
    L(x,p) = reduce(vcat, [Lⱼ(tracers(x)..., p) for Lⱼ in Ls]) # nonlinear part
    F(x,p) = NL(x,p) + L(x,p) - T(p) * x                     # full 𝐹(𝑥) = -T 𝑥 + 𝐺(𝑥)
    ∇ₓNL(x,p) = local_jacobian(NLs, x, p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓL(p) = local_jacobian(Ls, zeros(nt*nb), p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓNL(x,p) + ∇ₓL(p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, L, NL, ∇ₓF, ∇ₓL, ∇ₓNL, T
end
function split_state_function_and_Jacobian(T, L, NL, nb)
    F(x,p) = NL(x,p) + L(x,p) - T(p) * x                     # full 𝐹(𝑥)
    ∇ₓNL(x,p) = sparse(Diagonal(localderivative(NL, x, p))) # Jacobian of nonlinear part
    ∇ₓL(p) = sparse(Diagonal(localderivative(L, zeros(nb), p)))     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓNL(x,p) + ∇ₓL(p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, L, NL, ∇ₓF, ∇ₓL, ∇ₓNL, T
end
export split_state_function_and_Jacobian

function local_jacobian(Gs, x, p, nt, nb)
    return reduce(vcat, [local_jacobian_row(Gⱼ, x, p, nt, nb) for Gⱼ in Gs])
end
function inplace_local_jacobian(Gs, x, p, nt, nb)
    return reduce(vcat, [inplace_local_jacobian_row(Gⱼ!, x, p, nt, nb) for Gⱼ! in Gs])
end

function local_jacobian_row(Gᵢ, x, p, nt, nb)
    tracers(x) = state_to_tracers(x, nb, nt)
    return reduce(hcat, [sparse(Diagonal(localderivative(Gᵢ, tracers(x), j, p))) for j in 1:nt])
end
function inplace_local_jacobian_row(Gᵢ!, x, p, nt, nb)
    tracers(x) = state_to_tracers(x, nb, nt)
    dx = Vector{Float64}(undef,nb)
    return reduce(hcat, [sparse(Diagonal(localderivative(Gᵢ!, dx, tracers(x), j, p))) for j in 1:nt])
end

#=============================================
Generate 𝑓 and derivatives from user input
=============================================#

function generate_objective(ωs, μx, σ²x, v, ωp)
    nt, nb = length(ωs), length(v)
    tracers(x) = state_to_tracers(x, nb, nt)
    f(x, p) = ωp * mismatch(p) +
        sum([ωⱼ * mismatch(xⱼ, μⱼ, σⱼ², v) for (ωⱼ, xⱼ, μⱼ, σⱼ²) in zip(ωs, tracers(x), μx, σ²x)])
    return f
end
function generate_objective(ωs, ωp, grd, obs; kwargs...)
    nt, nb = length(ωs), count(iswet(grd))
    tracers(x) = state_to_tracers(x, nb, nt)
    Ms = [interpolationmatrix(grd, obsⱼ) for obsⱼ in obs]
    cs = get(kwargs, :cs, (collect(identity for i in 1:nt)...,))
    f(x, p) = ωp * mismatch(p) +
        sum([ωⱼ * mismatch(xⱼ, grd, obsⱼ, M=Mⱼ, c=cⱼ) for (ωⱼ, xⱼ, obsⱼ, Mⱼ, cⱼ) in zip(ωs, tracers(x), obs, Ms, cs)])
    return f
end
function generate_objective(ωs, ωp, grd, modify::Function, obs)
    nt, nb = length(ωs), count(iswet(grd))
    Ms = [interpolationmatrix(grd, obsⱼ) for obsⱼ in obs]
    iwets = [iswet(grd, obsⱼ) for obsⱼ in obs]
    function f(x, p)
        xs = unpack_tracers(x, grd)
        return ωp * mismatch(p) + sum([ωᵢ * indirectmismatch(xs, grd, modify, obs, i, Mᵢ, iwetᵢ) for (i, (ωᵢ, Mᵢ, iwetᵢ)) in enumerate(zip(ωs, Ms, iwets))])
    end
    return f
end

function generate_∇ₓobjective(ωs, μx, σ²x, v)
    nt, nb = length(ωs), length(v)
    tracers(x) = state_to_tracers(x, nb, nt)
    ∇ₓf(x, p) = reduce(hcat, [ωⱼ * ∇mismatch(xⱼ, μⱼ, σⱼ², v) for (ωⱼ, xⱼ, μⱼ, σⱼ²) in zip(ωs, tracers(x), μx, σ²x)])
    return ∇ₓf
end
function generate_∇ₓobjective(ωs, grd, obs; kwargs...)
    nt, nb = length(ωs), count(iswet(grd))
    tracers(x) = state_to_tracers(x, nb, nt)
    Ms = [interpolationmatrix(grd, obsⱼ) for obsⱼ in obs]
    cs = get(kwargs, :cs, (collect(identity for i in 1:nt)...,))
    ∇ₓf(x, p) = reduce(hcat, [ωⱼ * ∇mismatch(xⱼ, grd, obsⱼ, M=Mⱼ, c=cⱼ) for (ωⱼ, xⱼ, obsⱼ, Mⱼ, cⱼ) in zip(ωs, tracers(x), obs, Ms, cs)])
    return ∇ₓf
end
function generate_∇ₓobjective(ωs, grd, modify::Function, obs)
    nt, nb = length(ωs), count(iswet(grd))
    Ms = [interpolationmatrix(grd, obsⱼ) for obsⱼ in obs]
    iwets = [iswet(grd, obsⱼ) for obsⱼ in obs]
    function ∇ₓf(x, p)
        xs = unpack_tracers(x, grd)
        return sum([ωᵢ * ∇indirectmismatch(unpack_tracers(x, grd), grd, modify, obs, i, Mᵢ, iwetᵢ) for (i, (ωᵢ, Mᵢ, iwetᵢ)) in enumerate(zip(ωs, Ms, iwets))])
    end
    return ∇ₓf
end

function generate_∇ₚobjective(ωp)
    ∇ₚf(x, p) = ωp * ∇mismatch(p)
    return ∇ₚf
end

generate_objective_and_derivatives(ωs, μx, σ²x, v, ωp) =
    generate_objective(ωs, μx, σ²x, v, ωp),
    generate_∇ₓobjective(ωs, μx, σ²x, v),
    generate_∇ₚobjective(ωp)
generate_objective_and_derivatives(ωs, ωp, grd, obs; kwargs...) =
    generate_objective(ωs, ωp, grd, obs; kwargs...),
    generate_∇ₓobjective(ωs, grd, obs; kwargs...),
    generate_∇ₚobjective(ωp)
function generate_objective_and_derivatives(ωs, ωp, grd, modify::Function, obs)
    return generate_objective(ωs, ωp, grd, modify, obs),
             generate_∇ₓobjective(ωs, grd, modify, obs),
             generate_∇ₚobjective(ωp)
end

export generate_objective, generate_∇ₓobjective, generate_∇ₚobjective
export generate_objective_and_derivatives

"""
    mismatch(x, xobs, σ²xobs, v)

Volume-weighted mismatch of modelled tracer `x` against observed mean, `xobs`, given observed variance, `σ²xobs`, and volumes `v`.
"""
function mismatch(x, xobs, σ²xobs, v)
    δx = x - xobs
    W = Diagonal(v ./ σ²xobs)
    return 0.5 * transpose(δx) * W * δx / (transpose(xobs) * W * xobs)
end
mismatch(x, ::Missing, args...) = 0

"""
    ∇mismatch(x, xobs, σ²xobs, v)

Adjoint of the gradient of `mismatch(x, xobs, σ²xobs, v)`.
"""
function ∇mismatch(x, xobs, σ²xobs, v)
    δx = x - xobs
    W = Diagonal(v ./ σ²xobs)
    return transpose(W * δx) / (transpose(xobs) * W * xobs)
end
∇mismatch(x, ::Missing, args...) = transpose(zeros(length(x)))

## new functions for more generic obs packages
# TODO Add an optional function argument to transform the data before computingn the mismatch
# Example if for isotope tracers X where one ususally wants to minimize the mismatch in δ or ε.
function mismatch(x, grd::OceanGrid, obs; c=identity, W=I, M=interpolationmatrix(grd, obs.metadata), iwet=iswet(grd, obs))
    o = view(obs, iwet)
    δx = M * c(x) - o
    return 0.5 * transpose(δx) * W * δx / (transpose(o) * W * o)
end
mismatch(x, grd::OceanGrid, ::Missing; kwargs...) = 0
function ∇mismatch(x, grd::OceanGrid, obs; c=identity, W=I, M=interpolationmatrix(grd, obs.metadata), iwet=iswet(grd, obs))
    ∇c = Diagonal(ForwardDiff.derivative(λ -> c(x .+ λ), 0.0))
    o = view(obs, iwet)
    δx = M * c(x) - o
    return transpose(W * δx) * M * ∇c / (transpose(o) * W * o)
end
∇mismatch(x, grd::OceanGrid, ::Missing; kwargs...) = transpose(zeros(length(x)))

# In case the mismatch is not based on the tracer but on some function of it
function indirectmismatch(xs::Tuple, grd::OceanGrid, modify::Function, obs, i, M = interpolationmatrix(grd, obs[i].metadata), iwet = iswet(grd, obs[i]))
    x2 = modify(xs...)
    out = 0.0
    M = interpolationmatrix(grd, obs[i].metadata)
    iwet = iswet(grd, obs[i])
    o = view(obs[i], iwet)
    δx = M * x2[i] - o
    return 0.5 * transpose(δx) * δx / (transpose(o) * o)
end
function ∇indirectmismatch(xs::Tuple, grd::OceanGrid, modify::Function, obs, i, M = interpolationmatrix(grd, obs[i].metadata), iwet = iswet(grd, obs[i]))
    nt, nb = length(xs), length(iswet(grd))
    x2 = modify(xs...)
    o = view(obs[i], iwet)
    δx = M * x2[i] - o
    ∇modᵢ = ∇modify(modify, xs, i)
    return transpose(δx) * M * ∇modᵢ / (transpose(o) * o)
end
# TODO think of more efficient way to avoid recomputing ∇modify whole for each i
function ∇modify(modify, xs, i, j)
    return sparse(Diagonal(ForwardDiff.derivative(λ -> modify(perturb_tracer(xs,j,λ)...)[i], 0.0)))
end
∇modify(modify, xs, i) = reduce(hcat, [∇modify(modify, xs, i, j) for j in 1:length(xs)])

#=============================================
multi-tracer norm
=============================================#

function volumeweighted_norm(nt, v)
    w = repeat(v, nt)
    return nrm(x) = transpose(x) * Diagonal(w) * x
end

#=============================================
unpacking of multi-tracers
=============================================#

state_to_tracers(x, nb, nt) = ntuple(i -> state_to_tracer(x, nb, nt, i), nt)
state_to_tracer(x, nb, nt, i) = x[tracer_indices(nb, nt, i)]
function state_to_tracers(x, grd)
    nb = number_of_wet_boxes(grd)
    nt = Int(round(length(x) / nb))
    return state_to_tracers(x, nb, nt)
end
tracer_indices(nb, nt, i) = (i-1)*nb+1 : i*nb
tracers_to_state(xs) = reduce(vcat, xs)
export state_to_tracers, state_to_tracer, tracers_to_state, tracer_indices
# Alias for better name
unpack_tracers = state_to_tracers
export unpack_tracers


text:
abstract type ABasic end

struct Node{T<:ABasic}
    data::String
    class::Type{T}
    args::NTuple{N,Node} where N  # array does not work since Node([]]) == Node([]]) returns false
end
Node(data::String, class::Type{<:ABasic}) = Node(data, class, ())
data(self::Node) = self.data
class(self::Node) = self.class
args(self::Node) = self.args

abstract type ANode <: ABasic end
abstract type AAtom <: ANode end
abstract type AOperator <: ANode end

abstract type AUnknown <: AAtom end
abstract type AVariable <: AUnknown end
struct Variable{T} <: AVariable
end

abstract type AWild <: AUnknown end
struct Wild{T} <: AWild
end

abstract type ASingleton <: AAtom end

abstract type ASingletonStatement <: ASingleton end
abstract type ASingletonSet <: ASingleton end
abstract type ASingletonExpression <: ASingleton end

wild_symbol(name::String, set::Type{<:ASingletonSet}) = Node(name, Wild{set})
wild_symbol(name::String, set::Node) = Node(name, Wild{class(set)})

symbol(name::String, set::Type{<:ASingletonSet}) = Node(name, Variable{set})
symbol(name::String, set::Node) = Node(name, Variable{class(set)})
export symbol

# External types
abstract type ANonBasic <: ABasic end
convert(::Type{Node}, x::Number) = Node(String(x), ANonBasic)


text:
"""
    arbitraryparam(params)

returns the first field in the struct that is a subtype of `ClapeyronParam`. errors if it finds none.
"""
function arbitraryparam(params)
    paramstype = typeof(params)
    idx = findfirst(z->z <: ClapeyronParam,fieldtypes(paramstype))
    if isnothing(idx)
        error("The paramater struct ", paramstype, " must contain at least one ClapeyronParam")
    end
     return fieldnames(paramstype)[idx] |> z->getfield(params,z)
end

"""
    @comps

This macro is an alias to

    model.icomponents

The caveat is that `model` has to exist in the local namespace.
`model` is expected to be an EoSModel type that contains the `icomponents` field.
`icomponents` is an iterator that goes through all component indices.
"""
macro comps()
    return quote
        1:length(model)
    end |> esc
end

"""
    @groups

This macro is an alias to

    model.groups.i_flattenedgroups

`iflattenedgroups` is an iterator that goes through all groups in flattenedgroups.
"""
macro groups()
    return :($(esc(:(model.groups.i_flattenedgroups::UnitRange{Int64}))))
end

"""
    @groups(component)

This macro is an alias to

    model.groups.i_groups[component]

`i_groups[component]` is an iterator that goes through all groups in relevent to a given component.
"""
macro groups(component)
    return :($(esc(:(model.groups.i_groups[$(component)]::Vector{Int}))))
end

"""
    @sites(component)

This macro is an alias to

    model.sites.i_sites[component]

`i_sites[component]` is an iterator that goes through all sites relevant to
each group in a GC model, and to each main component in a non-GC model.
"""
macro sites(component)
    return :($(esc(:(model.sites.i_sites[$(component)]::Vector{Int}))))
end

"""
    @f(func,a,b,c,...)

This macro is an alias to

    func(model, V, T, z, a, b, c, ...)

where `func` is the name of the function, `model` is the model struct,
`V` is the volume, `T` is the absolute temperature, `z` is an array of number
of moles of each component, and `a`, `b`, `c`, ... are arbitrary parameters
that get passed to `func`.

It is very common for functions that are involved in the models to contain the
`model`, `V`, `T` and `z` parameters, so this macro helps reduce code repetition
as long as the first four parameters in the function are written exactly as above.

"""
macro f(func, args...)
    f = func
    model = :model
    V = :V
    T = :T
    z = :z
    quote
        $f($model,$V,$T,$z,$(args...))
    end |> esc
end

"""
    @nan(function_call,default=NaN)

Wraps the function in a `try-catch` block, and if a `DomainError` or `DivideError` is raised, then returns `default`.
for better results, its best to generate the default result beforehand
"""
macro nan(Base.@nospecialize(fcall),default = nothing)
    e = gensym(:error)
    quote
      try $fcall
      catch $e
        if $e isa Union{DomainError,DivideError}
          $default
        else
          rethrow($e)
        end
      end
    end |> esc
  end

"""
    @newmodelgc modelname parent paramstype

This is a data type that contains all the information needed to use an EoS model.
It also functions as an identifier to ensure that the right functions are called.

The user is expected to create an outter constructor that takes this signature

    function modelname(components::Array{String,1})

It should then return name(params::paramtype, groups::GroupParam, sites::SiteParam, idealmodel::IdealModel)

= Fields =
The Struct consists of the following fields:

* components: a string lists of components
* icomponents: an iterator that goes through the indices corresponding to each component
* groups: a [`GroupParam`](@ref)
* sites: a [`SiteParam`](@ref)
* params: the Struct paramstype that contains all parameters in the model
* idealmodel: the IdealModel struct that determines which ideal model to use
* assoc_options: struct containing options for the association solver. see [AssocOptions](@ref)
* references: reference for this EoS

See the tutorial or browse the implementations to see how this is used.
"""
macro newmodelgc(name, parent, paramstype)
    quote
    struct $name{T <: IdealModel} <: $parent
        components::Array{String,1}
        icomponents::UnitRange{Int}
        groups::GroupParam
        sites::SiteParam
        params::$paramstype
        idealmodel::T
        assoc_options::AssocOptions
        references::Array{String,1}
    end

    has_sites(::Type{<:$name}) = true
    has_groups(::Type{<:$name}) = true

    function Base.show(io::IO, mime::MIME"text/plain", model::$name)
        return gc_eosshow(io, mime, model)
    end

    function Base.show(io::IO, model::$name)
        return eosshow(io, model)
    end

    Base.length(model::$name) = Base.length(model.icomponents)

    molecular_weight(model::$name,z=SA[1.0]) = group_molecular_weight(model.groups,mw(model),z)

end |> esc
end

"""

    @newmodel name parent paramstype

This is exactly the same as the above but for non-GC models.
All group parameters are absent in this struct.
The sites are associated to the main component rather than the groups,
and the respective fieldnames are named correspondingly.
"""
macro newmodel(name, parent, paramstype)
    quote
    struct $name{T <: IdealModel} <: $parent
        components::Array{String,1}
        icomponents::UnitRange{Int}
        sites::SiteParam
        params::$paramstype
        idealmodel::T
        assoc_options::AssocOptions
        references::Array{String,1}
    end
    has_sites(::Type{<:$name}) = true
    has_groups(::Type{<:$name}) = false

    function Base.show(io::IO, mime::MIME"text/plain", model::$name)
        return eosshow(io, mime, model)
    end

    function Base.show(io::IO, model::$name)
        return eosshow(io, model)
    end
    molecular_weight(model::$name,z=SA[1.0]) = comp_molecular_weight(mw(model),z)
    Base.length(model::$name) = Base.length(model.icomponents)
    end |> esc
end

"""
    @newmodelsimple

Even simpler model, primarily for the ideal models.
Contains neither sites nor ideal models.
"""
macro newmodelsimple(name, parent, paramstype)
    quote
    struct $name <: $parent
        components::Array{String,1}
        icomponents::UnitRange{Int}
        params::$paramstype
        references::Array{String,1}
    end
    has_sites(::Type{<:$name}) = false
    has_groups(::Type{<:$name}) = false

    function Base.show(io::IO, mime::MIME"text/plain", model::$name)
        return eosshow(io, mime, model)
    end

    function Base.show(io::IO, model::$name)
        return eosshow(io, model)
    end

    Base.length(model::$name) = Base.length(model.icomponents)

    end |> esc
end

const IDEALTYPE = Union{T,Type{T}} where T<:IdealModel

function (::Type{model})(params::EoSParam,
        groups::GroupParam,
        sites::SiteParam,
        idealmodel::IDEALTYPE = BasicIdeal;
        ideal_userlocations::Vector{String}=String[],
        references::Vector{String}=String[],
        assoc_options::AssocOptions = AssocOptions(),
        verbose::Bool = false) where model <:EoSModel

        components = groups.components
        icomponents = 1:length(components)
        init_idealmodel = init_model(idealmodel,components,ideal_userlocations,verbose)
        return model(components, icomponents,
        groups,
        sites,
        params, init_idealmodel, assoc_options, references)
end

function (::Type{model})(params::EoSParam,
        groups::GroupParam,
        idealmodel::IDEALTYPE = BasicIdeal;
        ideal_userlocations::Vector{String}=String[],
        references::Vector{String}=String[],
        assoc_options::AssocOptions = AssocOptions(),
        verbose::Bool = false) where model <:EoSModel

    sites = SiteParam(groups.components)
    return model(params,groups,sites,idealmodel;ideal_userlocations,references,assoc_options,verbose)
end

#non GC
function (::Type{model})(params::EoSParam,
        sites::SiteParam,
        idealmodel::IDEALTYPE = BasicIdeal;
        ideal_userlocations::Vector{String}=String[],
        references::Vector{String}=String[],
        assoc_options::AssocOptions = AssocOptions(),
        verbose::Bool = false) where model <:EoSModel

    components = sites.components
    icomponents = 1:length(components)

    init_idealmodel = init_model(idealmodel,components,ideal_userlocations,verbose)
    return model(components, icomponents,
    sites, params, init_idealmodel, assoc_options, references)

end

#non GC may be shared with model simple
function (::Type{model})(params::EoSParam,
        idealmodel::IDEALTYPE = BasicIdeal;
        ideal_userlocations::Vector{String}=String[],
        references::Vector{String}=String[],
        assoc_options::AssocOptions = AssocOptions(),
        verbose::Bool = false) where model <:EoSModel

    if has_sites(model)
        arbparam = arbitraryparam(params)
        components = arbparam.components
        sites = SiteParam(components)
        return model(params,sites,idealmodel;ideal_userlocations,references,assoc_options,verbose)
    end
    #With sites out of the way, this is a simplemodel, no need to initialize the ideal model
    #if there isnt any params, just put empty values.
    if length(fieldnames(typeof(params))) > 0
    arbparam = arbitraryparam(params)
    components = arbparam.components
    icomponents = 1:length(components)
    else
        components = String[]
        icomponents =1:0
    end
    return model(components,icomponents,params,references)
end

function init_model(idealmodel::IdealModel,components,userlocations,verbose)
    return idealmodel
end

function init_model(::Nothing,components,userlocations,verbose)
    return nothing
end

function init_model(idealmodel::Type{<:IdealModel},components,userlocations,verbose)
    verbose && @info("""Now creating ideal model:
    $idealmodel""")
    return idealmodel(components;userlocations,verbose)
end
"""
    @registermodel(model)

given an existing model, composed of Clapeyron EoS models, ClapeyronParams or EoSParams, it will generate
the necessary traits to make the model compatible with Clapeyron routines.

"""
macro registermodel(model)
    _model = @eval $model
    _has_components = hasfield(_model,:components)
    splittable = _has_components

    _has_sites = hasfield(_model,:sites)
    _has_groups = hasfield(_model,:groups)
    _eos_show = if _has_groups
        :(gc_eosshow(io, mime, model))
    else
        :(eosshow(io, mime, model))
    end

    _len = if hasfield(_model,:icomponents)
        :(Base.length(model.icomponents))
    else
        :(Base.length(model.components))
    end

    _length = if _has_components
            :(Base.length(model::$model) = $_len)
        else
            :()
        end

    _mw = if _has_groups
        :(group_molecular_weight(model.groups,mw(model),z))
    else
        :(comp_molecular_weight(mw(model),z))
    end

    _molecular_weight = if _has_components
        :(molecular_weight(model::$model,z=SA[1.0]) =$_mw)
    else
        :()
    end

return quote
    has_sites(::Type{<:$model}) = $_has_sites
    has_groups(::Type{<:$model}) = $_has_groups

    function Base.show(io::IO, mime::MIME"text/plain", model::$model)
        return $_eos_show
    end

    function Base.show(io::IO, model::$model)
        return eosshow(io, model)
    end

    $_length

    $_molecular_weight

    end |> esc
end

export @newmodel, @f, @newmodelgc, @newmodelsimple


text:
#!/usr/bin/env julia
# coding: utf-8

# In[1]:

#boyer moore is a string search algorithm
#the main objective is to skip as many loops as possible
#the skipping depends on two heuristics
#bad character and good suffix
#a more detailed tutorial can be found in jhu material
# https://www.cs.jhu.edu/~langmea/resources/lecture_notes/strings_matching_boyer_moore.pdf

# In[2]:

#naïve search iterates the letter one by one
function naive_search(pattern,rawtext)

    len_pattern=length(pattern)
    len_rawtext=length(rawtext)
    output=[]

    #this part is stupid
    #as julia allows us to do
    #rawtext[i:i+len_pattern-1]==pattern
    for i in 1:(len_rawtext-len_pattern+1)
        ignore=false
        j=1
        while !ignore
            if rawtext[i+j-1]!=pattern[j]
                ignore=true
            end
            j+=1
            if j>len_pattern
                if !ignore
                    push!(output,i)
                end
                ignore=true
            end
        end
    end
    return output
end

# In[3]:

#preprocess for bad character
#for each index
#look back and search for every distinct letter in front of the index
#find the nearest index (as opposed to the current index)
#where every distinct letter occurs
function get_hbc(pattern)
    heuristics_bc=[]
    for i in 1:length(pattern)
        bc=Dict()
        for j in i:-1:1
            if !(pattern[j] in keys(bc)) && pattern[j]!=pattern[i]
                bc[pattern[j]]=i-j
            end
        end
        push!(heuristics_bc,bc)
    end
    return heuristics_bc
end

# In[4]:

#the concept of good suffix is very similar to lps in kmp algo
#the only difference is bm algo starts from the right side of the pattern
#u can check kmp for reference
# https//github.com/je-suis-tm/search-and-sort/blob/master/knuth%20morris%20pratt%20search.jl
function get_hgs(pattern)

    #initialize
    heuristics_gs=zeros(Int8,length(pattern))

    #iterate
    for i in length(pattern)-1:-1:1
        stop=false
        left=i
        right=i+1
        while !stop

            #compute the longest proper prefix
            if pattern[left:left+length(pattern[right:end])-1]==pattern[right:end]
                stop=true
                break

            #if no match
            #increment right side and keep trying
            else
                right+=1
            end

            #if right side has reached its limit
            #try left side
            if right>length(pattern)
                left-=1
                right=i+1
            end

            #avoid index error
            if left==0
                stop=true
                break
            end
        end

        #remember to revert the pointer back to the end of the pattern
        heuristics_gs[i]=right-left+length(pattern)-i
    end

    #the last one is always one
    heuristics_gs[end]=1

    return heuristics_gs
end

# In[5]:

#bm algo with bad character and good suffix heuristics
function boyer_moore(pattern,rawtext)

    #get heuristics
    heuristics_bc=get_hbc(pattern)
    heuristics_gs=get_hgs(pattern)

    #initialize
    i=length(pattern)
    j=length(pattern)
    pos=[]

    #iterate
    while i<=length(rawtext)

        #if no match,take the maximum skip from two different heuristics
        if pattern[j]!=rawtext[i]
            if rawtext[i] in keys(heuristics_bc[j])
                hbc=heuristics_bc[j][rawtext[i]]
                hbc+=(length(pattern)-j)
            else
                hbc=length(pattern)
            end
            hgs=heuristics_gs[j]

            #move the pointer to the end of the pattern before skipping
            i+=max(hbc,hgs)
            j=length(pattern)

        #keep matching
        else
            i-=1
            j-=1
        end

        #if matched,create output
        if j==0
            i+=1
            push!(pos,i)
            i+=length(pattern)
            j=length(pattern)
        end
    end

    return pos
end

# In[6]:

#for some strange reasons
#cyrillic letters dont work the same way in julia as in python
#have to convert to latin letters
cyrillic2latin=Dict([("а", "a"),
 ("б", "b"),
 ("в", "v"),
 ("г", "h"),
 ("ґ", "g"),
 ("д", "d"),
 ("е", "e"),
 ("є", "ye"),
 ("ж", "zh"),
 ("з", "z"),
 ("и", "y"),
 ("і", "i"),
 ("ї", "yi"),
 ("й", "y"),
 ("к", "k"),
 ("л", "l"),
 ("м", "m"),
 ("н", "n"),
 ("о", "o"),
 ("п", "p"),
 ("р", "r"),
 ("с", "s"),
 ("т", "t"),
 ("у", "u"),
 ("ф", "f"),
 ("х", "kh"),
 ("ц", "ts"),
 ("ч", "ch"),
 ("ш", "sh"),
 ("щ", "shch"),
 ("ь", "()"),
 ("ю", "yu"),
 ("я", "ya"),
 ("А", "A"),
 ("Б", "B"),
 ("В", "V"),
 ("Г", "H"),
 ("Ґ", "G"),
 ("Д", "D"),
 ("Е", "E"),
 ("Є", "YE"),
 ("Ж", "ZH"),
 ("З", "Z"),
 ("И", "Y"),
 ("І", "I"),
 ("Ї", "YI"),
 ("Й", "Y"),
 ("К", "K"),
 ("Л", "L"),
 ("М", "M"),
 ("Н", "N"),
 ("О", "O"),
 ("П", "P"),
 ("Р", "R"),
 ("С", "S"),
 ("Т", "T"),
 ("У", "U"),
 ("Ф", "F"),
 ("Х", "KH"),
 ("Ц", "TS"),
 ("Ч", "CH"),
 ("Ш", "SH"),
 ("Щ", "SHCH"),
 ("Ь", "()"),
 ("Ю", "YU"),
 ("Я", "YA")])
rawtext="Знаменитості продовжують підтримувати Україну у війні, яку веде Російська Федерація. Серед них - актори, ведучі, співаки, письменники та найбагатші люди планети.\nТретій тиждень триває повномасштабне вторгнення російських загарбників на територію України. За цей час висловили підтримку та надали фінансову допомогу українцям, зокрема, канадський бізнесмен Ілон Маск, американська акторка українського походження Міла Куніс з чоловіком-колегою Ештоном Кутчером, американська артистка Мадонна, голлівудська кінозірка Леонардо ді Капріо та інші. А британський актор Бенедикт Камбербетч запропонував власне житло для біженців з України."
pattern="Украї";

# In[7]:

#alphabet conversion
cleantext=[]
for i in rawtext
    if string(i) in keys(cyrillic2latin)
        push!(cleantext,cyrillic2latin[string(i)])
    else
        push!(cleantext,i)
    end
end

#list2string
clean_pattern=join([cyrillic2latin[string(i)] for i in pattern])
clean_text=join(cleantext);

# In[8]:

println(naive_search(clean_pattern,clean_text)==boyer_moore(clean_pattern,clean_text))

# In[9]:

#0.000008 seconds (3 allocations: 144 bytes)
@time naive_search(clean_pattern,clean_text);

# In[10]:

#0.000170 seconds (483 allocations: 12.703 KiB)
#as usual,naïve search by me is always faster than improvements...
#in python bm is actually faster
@time boyer_moore(clean_pattern,clean_text);

# In[ ]:

# In[ ]:

max_stars_repo_path: boyer moore search.jl

max_stars_repo_name: tattooday/search-and-sort

max_issues_repo_path: boyer moore search.jl

max_issues_repo_name: tattooday/search-and-sort

max_forks_repo_path: boyer moore search.jl

max_forks_repo_name: tattooday/search-and-sort


text:
using StaticArrays, LinearAlgebra

"""
    GaussHermiteQuadrature

As described in

* [Gauss-Hermite quadrature on Wikipedia](http://en.wikipedia.org/wiki/Gauss-Hermite_quadrature)

*Gauss-Hermite* quadrature uses a weighted sum of values of `f(x)` at specific `x` values to approximate

```math
\\int_{-\\infty}^\\infty f(x) e^{-x^2} dx
```

An `n`-point rule, as returned by `hermite(n)` from the
[`GaussQuadrature``](https://github.com/billmclean/GaussQuadrature.jl) package provides `n` abscicca
values (i.e. values of `x`) and `n` weights.

As noted in the Wikipedia article, a modified version can be used to evaluate the expectation `E[h(x)]`
with respect to a `Normal(μ, σ)` density as
```julia
using MixedModels

gn5 = GHnorm(5)
μ = 3.
σ = 2.
sum(@. abs2(σ*gn5.z + μ)*gn5.w) # E[X^2] where X ∼ N(μ, σ)
```

For evaluation of the log-likelihood of a GLMM the integral to evaluate for each level of the grouping
factor is approximately Gaussian shaped.
"""
GaussHermiteQuadrature
"""
    GaussHermiteNormalized{K}

A struct with 2 SVector{K,Float64} members
- `z`: abscissae for the K-point Gauss-Hermite quadrature rule on the Z scale
- `wt`: Gauss-Hermite weights normalized to sum to unity
"""
struct GaussHermiteNormalized{K}
    z::SVector{K, Float64}
    w::SVector{K,Float64}
end
function GaussHermiteNormalized(k::Integer)
    ev = eigen(SymTridiagonal(zeros(k), sqrt.(1:k-1)))
    w = abs2.(ev.vectors[1,:])
    GaussHermiteNormalized(
        SVector{k}((ev.values .- reverse(ev.values)) ./ 2),
        SVector{k}(LinearAlgebra.normalize((w .+ reverse(w)) ./ 2, 1)))
end

Base.iterate(g::GaussHermiteNormalized{K}, i=1) where {K} = (K < i ? nothing : ((z = g.z[i], w = g.w[i]), i + 1))

Base.length(g::GaussHermiteNormalized{K}) where {K} = K

"""
    GHnormd

Memoized values of `GHnorm`{@ref} stored as a `Dict{Int,GaussHermiteNormalized}`
"""
const GHnormd = Dict{Int,GaussHermiteNormalized}(
    1 => GaussHermiteNormalized(SVector{1}(0.),SVector{1}(1.)),
    2 => GaussHermiteNormalized(SVector{2}(-1.0,1.0),SVector{2}(0.5,0.5)),
    3 => GaussHermiteNormalized(SVector{3}(-sqrt(3),0.,sqrt(3)),SVector{3}(1/6,2/3,1/6))
    )

"""
    GHnorm(k::Int)

Return the (unique) GaussHermiteNormalized{k} object.

The function values are stored (memoized) when first evaluated.  Subsequent evaluations
for the same `k` have very low overhead.
"""
GHnorm(k::Int) = get!(GHnormd, k) do
    GaussHermiteNormalized(k)
end
GHnorm(k) = GHnorm(Int(k))


text:
using BlackBoxOptim

require("../covar_matrix_self_adaptation_es.jl")
require("../experiment_framework.jl")

ps = BlackBoxOptim.as_fixed_dim_problem_set(BlackBoxOptim.example_problems, 2^5)

sf(p) = begin
  p = BlackBoxOptim.shifted(p)
  n = numdims(p)
  lambda = 4 + 3 * int(ceil(log10(n)))
  mu = 2 # int(ceil(lambda/2))
  cmsa_es(p, mu = mu, lambda = lambda)
end

@time repeated_runs(sf, ps, 25; experiment = "cmsa_es_2")


text:
# RanDAGen/src/task.jl
# Licensed under the MIT License. See LICENSE.md file in the project root for
# full license information.

"""
Task

Represents Stuff

# Fields
- `ID::Int64`: The unique ID of this Task.
- `work::Array{Pair(Symbol,Int64)}`: The (resource, work) pairs that makes up this tasks needs.
- `alpha::Float64`: The alpha parameter for Amdahl law calculation
- `complexity::Symbol`: The complexity of this Task (e.g :LOG_N, :N, :N_2)
- `comm::Symbol`: Direction of communication :send, :recieve, or :none
- `comm_target::Int64`: Identity of the process that needs to h
- `communication_cost::Float64`: The cost of communicating
- `data::Int64`: The amount of data this task needs to transfer

# Methods
- `val::Type{Any}`: words
"""
mutable struct Task
    ID::Int64
    alpha::Float64
    complexity::Symbol
    work::Dict{Symbol,Int64}

    data::Int64

    comm::Symbol
    target::Int64
    cost::Float64

    function Task(
        ID::Int64,
        alpha::Float64,
        complexity::Symbol,
        work::Dict{Symbol,Int64},
        data::Int64=0,
        comm::Symbol=:none,
        target::Int64=-1,
        cost::Float64=0.0,
    ) begin
        return new(ID, alpha, complexity, work, data, comm, target, cost)
    end
end

work(task::Task, n::Int64) = (task.work - n) < 0 ? 0 : task.work -= n
workRemaining(task::Task) = sum(x->x.second, task)
isDone(task::Task) = task.work.first === 0


text:
module MultiUse
    using Test
    using OOPMacro

    @class MUParent begin
        pfield::Int
        pfun(self) = self.pfield
    end

    module SomeOtherModule
        using OOPMacro
        @class MUParent begin
            otherField::Int
        end
    end

    @class MUChild(MUParent) begin
        cfield::Int
        pfun(self) = self.cfield
    end

    p = MUParent(0)
    c = MUChild(0,1)
    @test c.pfield == 0
    @test c.cfield == 1
    pvalue = p.pfield
    cvalue = c.cfield
    @test pfun(p) == pvalue
    @test pfun(c) == cvalue

    mup = SomeOtherModule.MUParent(2)
    @test mup.otherField == 2
end


text:
using OpenCL, CLBLAS, CLBlast, BenchmarkTools
@static if VERSION < v"0.7-"
    LA = Base.LinAlg
else
    using LinearAlgebra, Random, Printf
    LA = LinearAlgebra
end

CLBLAS.setup()

function run_gbmv(m::Integer, n::Integer, kl::Integer, ku::Integer, T=Float32)
    srand(12345)
    A = rand(T, kl+ku+1, n)
    y = rand(T, m)
    x = rand(T, n)
    α = rand(T)
    β = rand(T)

    @printf("\nm = %d, n = %d, kl = %d, ku = %d, eltype = %s\n", m, n, kl, ku, T)

    println("BLAS:")
    y_true = copy(y)
    LinAlg.BLAS.gbmv!('N', m, kl, ku, α, A, x, β, y_true)
    yy = copy(y)
    display(@benchmark LinAlg.BLAS.gbmv!($('N'), $m, $kl, $ku, $α, $A, $x, $β, $yy))
    println()

    for device in cl.devices()
        println("-"^70)
        @printf("Platform name   : %s\n", device[:platform][:name])
        @printf("Platform version: %s\n", device[:platform][:version])
        @printf("Device name     : %s\n", device[:name])
        @printf("Device type     : %s\n", device[:device_type])
        println()

        ctx = cl.Context(device)
        queue = cl.CmdQueue(ctx)

        println("CLBLAS:")
        A_cl = cl.CLArray(queue, A)
        y_cl = cl.CLArray(queue, y)
        x_cl = cl.CLArray(queue, x)
        CLBLAS.gbmv!('N', m, kl, ku, α, A_cl, x_cl, β, y_cl)
        @assert cl.to_host(y_cl) ≈ y_true
        display(@benchmark CLBLAS.gbmv!($('N'), $m, $kl, $ku, $α, $A_cl, $x_cl, $β, $y_cl))
        println()

        println("CLBlast:")
        A_cl = cl.CLArray(queue, A)
        y_cl = cl.CLArray(queue, y)
        x_cl = cl.CLArray(queue, x)
        CLBlast.gbmv!('N', m, kl, ku, α, A_cl, x_cl, β, y_cl)
        @assert cl.to_host(y_cl) ≈ y_true
        display(@benchmark CLBlast.gbmv!($('N'), $m, $kl, $ku, $α, $A_cl, $x_cl, $β, $y_cl))
        println()
    end
end

run_gbmv(2^10, 2^10, 5, 5, Float32)


text:
module UtilitiesWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AuxiliaryElements = @jimport org.orekit.propagation.semianalytical.dsst.utilities.AuxiliaryElements
const CjSjCoefficient = @jimport org.orekit.propagation.semianalytical.dsst.utilities.CjSjCoefficient
const Class = @jimport java.lang.Class
const CoefficientsFactory = @jimport org.orekit.propagation.semianalytical.dsst.utilities.CoefficientsFactory
const Field = @jimport org.hipparchus.Field
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FieldAuxiliaryElements = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldAuxiliaryElements
const FieldCjSjCoefficient = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldCjSjCoefficient
const FieldFixedNumberInterpolationGrid = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldFixedNumberInterpolationGrid
const FieldGHIJjsPolynomials = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldGHIJjsPolynomials
const FieldGHmsjPolynomials = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldGHmsjPolynomials
const FieldGammaMnsFunction = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldGammaMnsFunction
const FieldGradient = @jimport org.hipparchus.analysis.differentiation.FieldGradient
const FieldInterpolationGrid = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldInterpolationGrid
const FieldLnsCoefficients = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldLnsCoefficients
const FieldMaxGapInterpolationGrid = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldMaxGapInterpolationGrid
const FieldOrbit = @jimport org.orekit.orbits.FieldOrbit
const FieldShortPeriodicsInterpolatedCoefficient = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FieldShortPeriodicsInterpolatedCoefficient
const FieldVector3D = @jimport org.hipparchus.geometry.euclidean.threed.FieldVector3D
const FixedNumberInterpolationGrid = @jimport org.orekit.propagation.semianalytical.dsst.utilities.FixedNumberInterpolationGrid
const Frame = @jimport org.orekit.frames.Frame
const GHIJjsPolynomials = @jimport org.orekit.propagation.semianalytical.dsst.utilities.GHIJjsPolynomials
const GHmsjPolynomials = @jimport org.orekit.propagation.semianalytical.dsst.utilities.GHmsjPolynomials
const GammaMnsFunction = @jimport org.orekit.propagation.semianalytical.dsst.utilities.GammaMnsFunction
const Gradient = @jimport org.hipparchus.analysis.differentiation.Gradient
const InterpolationGrid = @jimport org.orekit.propagation.semianalytical.dsst.utilities.InterpolationGrid
const JString = @jimport java.lang.String
const JacobiPolynomials = @jimport org.orekit.propagation.semianalytical.dsst.utilities.JacobiPolynomials
const LnsCoefficients = @jimport org.orekit.propagation.semianalytical.dsst.utilities.LnsCoefficients
const MaxGapInterpolationGrid = @jimport org.orekit.propagation.semianalytical.dsst.utilities.MaxGapInterpolationGrid
const NewcombOperators = @jimport org.orekit.propagation.semianalytical.dsst.utilities.NewcombOperators
const Object = @jimport java.lang.Object
const Orbit = @jimport org.orekit.orbits.Orbit
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const ShortPeriodicsInterpolatedCoefficient = @jimport org.orekit.propagation.semianalytical.dsst.utilities.ShortPeriodicsInterpolatedCoefficient
const TreeMap = @jimport java.util.TreeMap
const UpperBounds = @jimport org.orekit.propagation.semianalytical.dsst.utilities.UpperBounds
const Vector3D = @jimport org.hipparchus.geometry.euclidean.threed.Vector3D

export AuxiliaryElements
export CjSjCoefficient
export CoefficientsFactory
export FieldAuxiliaryElements
export FieldCjSjCoefficient
export FieldFixedNumberInterpolationGrid
export FieldGHIJjsPolynomials
export FieldGHmsjPolynomials
export FieldGammaMnsFunction
export FieldInterpolationGrid
export FieldLnsCoefficients
export FieldMaxGapInterpolationGrid
export FieldShortPeriodicsInterpolatedCoefficient
export FixedNumberInterpolationGrid
export GHIJjsPolynomials
export GHmsjPolynomials
export GammaMnsFunction
export InterpolationGrid
export JacobiPolynomials
export LnsCoefficients
export MaxGapInterpolationGrid
export NewcombOperators
export ShortPeriodicsInterpolatedCoefficient
export UpperBounds
export add_grid_point
export clear_history
export compute_gs_hs
export compute_qns
export compute_vns
export equals
export get_alpha
export get_b
export get_beta
export get_c
export get_cj
export get_class
export get_date
export get_dcj_dh
export get_dcj_dk
export get_derivative
export get_dnl
export get_dsj_dh
export get_dsj_dk
export get_ecc
export get_frame
export get_gamma
export get_gjs
export get_gmsj
export get_grid_points
export get_h
export get_hjs
export get_hmsj
export get_ijs
export get_jjs
export get_k
export get_keplerian_period
export get_le
export get_lf
export get_lm
export get_lns
export get_lv
export get_mean_motion
export get_p
export get_q
export get_retrograde_factor
export get_rnml
export get_sj
export get_sma
export get_value
export get_vector_f
export get_vector_g
export get_vector_w
export get_vmns
export getd_gjsd_alpha
export getd_gjsd_beta
export getd_gjsdh
export getd_gjsdk
export getd_gmsd_alpha
export getd_gmsd_beta
export getd_gmsdh
export getd_gmsdk
export getd_hjsd_alpha
export getd_hjsd_beta
export getd_hjsdh
export getd_hjsdk
export getd_hmsd_alpha
export getd_hmsd_beta
export getd_hmsdh
export getd_hmsdk
export getd_ijsd_alpha
export getd_ijsd_beta
export getd_ijsdh
export getd_ijsdk
export getd_jjsd_alpha
export getd_jjsd_beta
export getd_jjsdh
export getd_jjsdk
export getd_lnsd_gamma
export hash_code
export notify
export notify_all
export to_string
export value
export wait

include("HansenWrapper/HansenWrapper.jl")
include("auxiliary_elements.jl")
include("cj_sj_coefficient.jl")
include("coefficients_factory.jl")
include("field_auxiliary_elements.jl")
include("field_cj_sj_coefficient.jl")
include("field_fixed_number_interpolation_grid.jl")
include("field_g_hmsj_polynomials.jl")
include("field_gamma_mns_function.jl")
include("field_ghi_jjs_polynomials.jl")
include("field_interpolation_grid.jl")
include("field_lns_coefficients.jl")
include("field_max_gap_interpolation_grid.jl")
include("field_short_periodics_interpolated_coefficient.jl")
include("fixed_number_interpolation_grid.jl")
include("g_hmsj_polynomials.jl")
include("gamma_mns_function.jl")
include("ghi_jjs_polynomials.jl")
include("interpolation_grid.jl")
include("jacobi_polynomials.jl")
include("lns_coefficients.jl")
include("max_gap_interpolation_grid.jl")
include("newcomb_operators.jl")
include("short_periodics_interpolated_coefficient.jl")
include("upper_bounds.jl")

end


text:
latexinline(args...;kwargs...) = process_latexify(args...;kwargs...,env=:inline)

function _latexinline(x; kwargs...)
    latexstr = latexstring( process_latexify(x; kwargs...,env=:raw) )
    COPY_TO_CLIPBOARD && clipboard(latexstr)
    return latexstr
end


text:
__precompile__()

module VersatileHDPMixtureModels

using Distributed
using StatsBase
using Distributions
using SpecialFunctions
using LinearAlgebra
using JLD2
using Clustering
using Random
using NPZ
using Base
using PDMats

#DS:
include("ds.jl")

#Distributions:
include("distributions/mv_gaussian.jl")
include("distributions/mv_group_gaussian.jl")
include("distributions/multinomial_dist.jl")
include("distributions/topic_modeling_dist.jl")
include("distributions/compact_mnm.jl")

#Priors:
include("priors/multinomial_prior.jl")
include("priors/topic_modeling_prior.jl")
include("priors/niw.jl")
include("priors/niw_stable_var.jl")
include("priors/bayes_network_model.jl")
include("priors/compact_mnm_prior.jl")

#Rest:
include("params_base.jl")
include("utils.jl")
include("shared_actions.jl")
include("local_clusters_actions.jl")
include("global_clusters_actions.jl")
include("crf_hdp.jl")
include("gaussian_generator.jl")
include("hdp_shared_features.jl")

#Data Generators:
export
    generate_sph_gaussian_data,
    generate_gaussian_data,
    generate_mnmm_data,
    generate_grouped_mnm_data,
    generate_grouped_gaussian_data,
    create_mnmm_data,
    create_gaussian_data,
    create_grouped_gaussian_data,
    create_grouped_mnmm_data,
    hdp_prior_crf_draws,
    generate_grouped_gaussian_from_hdp_group_counts,
#Model DS:
    hdp_shared_features,
    multinomial_dist,
    mv_gaussian,
    mv_group_gaussian,
    topic_modeling_dist,
    bayes_network_model,
    niw_hyperparams,
    niw_stable_hyperparams,
    topic_modeling_hyper,
    compact_mnm_hyper,
    compact_mnm_dist,
#Functions
    hdp_fit,
    vhdp_fit,
    create_default_priors,
    get_model_global_pred,
    create_global_labels

end # module


text:
mutable struct MyRecur{T,S,I,P}
  cell::T
  state::S
  input::I
  p::P
end
function MyRecur(cell)
  p = initial_params(cell)
  MyRecur(cell, cell.state0, zero(cell.state0), p)
end
function (m::MyRecur)(x, p=m.p)
  batchsize = size(x,2)
  m.input = x
  h′, y = m.cell(m.state, m.input, x, p)
  Inf ∈ h′ && return fill(Inf32, size(y,1), batchsize)
  m.state = h′
  return y
end
# function (m::MyRecur)(x::AbstractArray{T, 3}, p=m.p) where T
#   h = [m(view(x, :, :, i), p) for i in 1:size(x, 3)]
#   sze = size(h[1])
#   reshape(reduce(hcat, h), sze[1], sze[2], length(h))
# end
Base.show(io::IO, m::MyRecur) = print(io, "MyRecur(", m.cell, ")")
initial_params(m::MyRecur) = initial_params(m.cell)
paramlength(m::MyRecur) = length(m.p)
Flux.functor(::Type{<:MyRecur}, m) = (m.p), re -> MyRecur(m.cell, m.state, m.input, re...)
Flux.trainable(m::MyRecur) = (m.p,)
get_bounds(m::MyRecur, T::DataType=eltype(m.state)) = get_bounds(m.cell, T)
Flux.reset!(m::MyRecur) = reset_state!(m, m.p)
@views function reset_state!(m::MyRecur, p=m.p)
  #m.p = p
  pl = length(p)
  state = p[reshape(pl-size(m.cell.state0,1)+1:pl, :, 1)]
  m.state = state
  m.input = reshape(zeros(eltype(m.input), size(state,1)), :, 1)
  nothing
end
# TODO: reset_state! for cell with train_u0=false

function Base.getproperty(m::MyRecur{T,S,I,P}, s::Symbol) where {T,S,I,P}
  if s === :cell
    return getfield(m, :cell)
  elseif s === :state
    return getfield(m, :state)
  elseif s === :input
    return getfield(m, :input)
  elseif s === :p
    return getfield(m, :p)
  else
    return getfield(m, s)
  end
end

ode_solve_kwargs(; abstol=1e-4, reltol=1e-3, save_everystep=false, save_start=false, save_end=true, kwargs...) =
  (abstol=abstol, reltol=reltol, save_everystep=save_everystep, save_start=save_start, save_end=save_end, kwargs...)

struct BCTRNNCell{W,SOLVER,SENSE,PROB,LB,UB,S,P,KW}
  wiring::W

  solver::SOLVER
  sensealg::SENSE
  prob::PROB
  lb::LB
  ub::UB
  state0::S
  p::P
  kwargs::KW

  function BCTRNNCell(wiring, solver, sensealg, prob, lb, ub, state0, p; kwargs...)
    new{typeof(wiring),typeof(solver),typeof(sensealg),typeof(prob),typeof(lb),typeof(ub),typeof(state0),typeof(p), typeof(kwargs)}(
                      wiring, solver, sensealg, prob, lb, ub, state0, p, kwargs)
  end
end

function Base.getproperty(m::BCTRNNCell{W,SOLVER,SENSE,PROB,LB,UB,S,P,KW}, s::Symbol) where {W,SOLVER,SENSE,PROB,LB,UB,S,P,KW}
  if s === :s_in
    return getproperty(getfield(m, :wiring), s)
  elseif s === :n_total
    return getproperty(getfield(m, :wiring), s)
  # elseif s === :n_in
  #   return getproperty(getfield(m, :wiring), s)
  elseif s === :n_out
    return getproperty(getfield(m, :wiring), s)
  elseif s === :wiring
    return getfield(m, :wiring)
  elseif s === :solver
    return getfield(m, :solver)
  elseif s === :sensealg
    return getfield(m, :sensealg)
  elseif s === :prob
    return getfield(m, :prob)
  elseif s === :lb
    return getfield(m, :lb)
  elseif s === :ub
    return getfield(m, :ub)
  elseif s === :state0
    return getfield(m, :state0)
  elseif s === :p
    return getfield(m, :p)
  elseif s === :kwargs
    return getfield(m, :kwargs)
  else
    return getfield(m, s)
  end
end

function BCTRNNCell(wiring, solver, sensealg, prob, lb, ub; kwargs...)
  u0 = prob.u0
  p = prob.p
  p_ode = p[wiring.s_in+1:end]
  state0 = reshape(u0, :, 1)
  θ = vcat(p_ode, u0)
  BCTRNNCell(wiring, solver, sensealg, prob, lb, ub, state0, θ; kwargs...)
end

function BCTRNNCell(wiring, solver, sensealg, odef, u0, tspan, p, lb, ub; mtkize=false, gen_jac=false, kwargs...)
  _prob = ODEProblem{true}(odef, u0, tspan, p)
  p_ode = @view p[wiring.s_in+1:end]
  prob = mtkize_prob(_prob, odef, u0, tspan, p, mtkize, gen_jac)

  # eqs = ModelingToolkit.equations(sys)
  # sts = ModelingToolkit.states(sys)
  # noiseeqs = 0.1f0 .* sts
  # ps = ModelingToolkit.parameters(sys)
  # @named sde = SDESystem(eqs,noiseeqs,ModelingToolkit.get_iv(sys),sts,ps)
  # prob = SDEProblem(sde,u0,tspan,p)

  state0 = reshape(u0, :, 1)
  θ = vcat(p_ode, u0)

  BCTRNNCell(wiring, solver, sensealg, prob, lb, ub, state0, θ; kwargs...)
end

function (m::BCTRNNCell)(h, _last_input, x::AbstractArray, p=m.p)
  # size(h) == (N,1) at the first MTKNODECell invocation. Need to duplicate batchsize times
  nstates = size(h,1)
  batchsize = size(x,2)
  num_reps = batchsize-size(h,2)+1
  u0s = repeat(h, 1, num_reps)
  last_input = repeat(_last_input, 1, num_reps)

  tspan_h = m.prob.tspan
  tspan_o = (tspan_h[2], tspan_h[2]+1)

  p_ode = @view p[1:end-nstates]

  out_idxs = _output_idxs(m.wiring.output_mapping, m.wiring.n_total)

  #sol = solve_ode(m, u0s, tspan_h, p_ode, last_input)
  sol = solve_ode(m, u0s, tspan_h, p_ode, x)
  h′ = sol[:, end, :]

  Inf ∈ h′ && return (fill(Inf32, nstates, batchsize), fill(Inf32, length(out_idxs), batchsize))
  NaN ∈ h′ && return (fill(Inf32, nstates, batchsize), fill(Inf32, length(out_idxs), batchsize))

  return (h′, (@view h′[out_idxs, :]))

  # sol = solve_ode(m, h′, tspan_o, p_ode, x)
  # h′ = sol[:, end, :]

  # return (h′, h′[out_idxs, :])
end

function solve_ode(m, u0s, tspan, p_ode, x)
  batchsize = size(u0s, 2)
  prob = m.prob
  function prob_func(prob,i,repeat)
    u0 = @view u0s[:,i]
    stim = @view x[:,i]
    p = vcat(stim, p_ode)

    # condition(u,t,integrator) = any((u .> 10) .| (u .< -10))
    # function affect!(integrator)
    #   for ii in 1:length(integrator.u)
    #     integrator.u[ii] = clamp(integrator.u[ii], -one(integrator.u[ii]), one(integrator.u[ii]))
    #   end
    # end
    # cb = ContinuousCallback(condition,affect!)

    remake(prob; u0, tspan, p)
  end
  function output_func(sol,i)

    # Zygote.@ignore begin
    #   if i == 1 && length(sol.t) > 1
    #     @show length(sol.t)
    #     fig = plot(xlabel="t", ylabel="step size",title="Adaptive step size")
    #     steps1 = (@view sol.t[2:end]) - (@view sol.t[1:end-1])
    #     plot!(fig, (@view sol.t[2:end]),steps1)
    #     display(fig)
    #   end
    # end
    (sol.retcode != :Success || NaN ∈ sol) && return fill(Inf32, size(u0s,1)), false
    (@view sol[:, end]), false
  end

  kwargs = ode_solve_kwargs(; m.kwargs...)
  ensemble_prob = EnsembleProblem(prob; prob_func, output_func, safetycopy=false) # TODO: safetycopy ???
  sa = Array(solve(ensemble_prob, m.solver, EnsembleThreads(), trajectories=batchsize;
              sensealg=m.sensealg,
              kwargs...
  ))

  sa
end

Base.show(io::IO, m::BCTRNNCell) = print(io, "BCTRNNCell(", m.s_in, ",", m.n_out, ")")
initial_params(m::BCTRNNCell) = m.p
paramlength(m::BCTRNNCell) = length(m.p)
Flux.functor(::Type{<:BCTRNNCell}, m) = (m.p,), re -> BCTRNNCell(m.wiring, m.solver, m.sensealg, m.prob, m.lb, m.ub, m.state0, re...)
Flux.trainable(m::BCTRNNCell) = (m.p,)

function get_bounds(m::BCTRNNCell, T::DataType=eltype(m.p))
  m.lb, m.ub
end

function mtkize_prob(prob, f, u0, tspan, p, mtkize, gen_jac)
  mtkize == false && return prob

  sys = ModelingToolkit.modelingtoolkitize(prob)
  #sys = ModelingToolkit.structural_simplify(sys)

  gen_jac == false && return ODEProblem{true}(sys,u0,tspan,p, tgrad=true)

  jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
  odef = ODEFunction{true}(f, jac=jac)
  ODEProblem{true}(odef,u0,tspan,p, tgrad=true)
end


text:
export PacketHeader

type PacketHeader
    typ :: UInt8 # hsa_packet_type_t
    barrier :: Bool
    acquire_fence_scope :: UInt8 # hsa_fence_scope_t
    release_fence_scope :: UInt8 # hsa_fence_scope_t
end

function PacketHeader(typ;
    barrier::Bool = false,
    acquire = FenceScopeNone,
    release = FenceScopeNone
    )
    typ = convert(UInt8, typ)
    acquire = convert(UInt8, acquire)
    release = convert(UInt8, release)

    return PacketHeader(typ, barrier, acquire, release)
end

import Base.==

function ==(h1 :: PacketHeader, h2 :: PacketHeader)
    return h1.typ == h2.typ &&
        h1.barrier == h2.barrier &&
        h1.acquire_fence_scope == h2.acquire_fence_scope &&
        h1.release_fence_scope == h2.release_fence_scope
end

import Base.unsafe_convert

function unsafe_convert(::Type{PacketHeader}, ptr :: Ptr{Void})
    if ptr == C_NULL
        error("invalid packet pointer")
    end

    h_ptr = convert(Ptr{UInt8}, ptr)

    h_type = unsafe_load(h_ptr)
    byte2 = unsafe_load(h_ptr + 1)
    h_barrier = (byte2 & 0x01) != 0x00
    h_acq_fen = (byte2 >> 1) & 0x03
    h_rel_fen = (byte2 >> 3) & 0x03

    return PacketHeader(h_type, h_barrier, h_acq_fen, h_rel_fen)
end

import Base.unsafe_store!

function unsafe_store!(ptr :: Ptr{Void}, hdr :: PacketHeader)
    h_ptr = convert(Ptr{UInt8}, ptr)

    h_byte2 =
        ((hdr.barrier) ? 0x01 : 0x00) |
        hdr.acquire_fence_scope << 1 |
        hdr.release_fence_scope << 3
    unsafe_store!(h_ptr + 1, h_byte2)

    # set the type last, once the packet is complete
    unsafe_store!(h_ptr, hdr.typ)
end

import Base.copy

function copy(hdr :: PacketHeader)
    return PacketHeader(hdr.typ, hdr.barrier, hdr.acquire_fence_scope, hdr.release_fence_scope)
end

abstract AQLPacket

import Base.==

function ==(p1 :: AQLPacket, p2 :: AQLPacket)
    bytes1 = Array(UInt8, 64)
    bytes2 = Array(UInt8, 64)
    bytes1_ptr = convert(Ptr{Void}, pointer(bytes1))
    bytes2_ptr = convert(Ptr{Void}, pointer(bytes2))

    unsafe_store!(bytes1_ptr, p1)
    unsafe_store!(bytes2_ptr, p2)

    return bytes1 == bytes2
end

export KernelDispatchPacket

type KernelDispatchPacket <: AQLPacket
    header :: PacketHeader
    dimensions :: UInt16
    workgroup_size_x :: UInt16
    workgroup_size_y :: UInt16
    workgroup_size_z :: UInt16
    grid_size_x :: UInt32
    grid_size_y :: UInt32
    grid_size_z :: UInt32
    private_segment_size :: UInt32
    group_segment_size :: UInt32
    kernel_object :: UInt64
    kernarg_address :: Ptr{Void}
    completion_signal :: hsa_signal_t

    function KernelDispatchPacket(
        kernel_object,
        range;
        dimensions = 0,
        wg_size = nothing,
        kernarg_address = 0,
        header :: PacketHeader = PacketHeader(
                PacketTypeKernelDispatch,
                barrier = false,
                # acquire
                # release
            ),
        private_segment_size = 0,
        group_segment_size = 0,
        completion_signal = hsa_signal_t(0)
        )

        grid_size = [1,1,1]
        wg_size_3d = [1,1,1]

        if dimensions == 0
            if isa(range, Tuple)
                dimensions = length(range)
            else
                dimensions = 1
            end
        end

        if dimensions < 1 || dimensions > 3
            error("range dimensions out of range, expected 1-3, got $dimensions")
        end

        if isa(range, Tuple)
            if length(range) < dimensions
                error("too few workgroup dimensions specified")
            end

            for i = 1:dimensions
                grid_size[i] = range[i]
            end
        else
            grid_size[1] = range
        end

        if wg_size != nothing
            if isa(wg_size, Tuple)
                if length(wg_size) < dimensions
                    error("too few workgroup dimensions specified")
                end
                for i = 1:dimensions
                    wg_size_3d[i] = wg_size[i]
                end
            else
                if 1 < dimensions
                    error("too few workgroup dimensions specified")
                end
                wg_size_3d[1] = wg_size
            end
        end

        return new(
            header,
            dimensions, # dimensions
            wg_size_3d[1],
            wg_size_3d[2],
            wg_size_3d[3],
            grid_size[1],
            grid_size[2],
            grid_size[3],
            private_segment_size, # private_segment_size
            group_segment_size, # group_segment_size
            kernel_object, # kernel_object
            kernarg_address, # kernarg_address
            completion_signal, # completion_signal
        )
    end
end

function load(::Type{AQLPacket}, ::Type{Val{PacketTypeKernelDispatch}}, ptr :: Ptr{Void}, p_hdr :: PacketHeader)
    if ptr == C_NULL
        error("invalid packet pointer")
    end

    p_dims = unsafe_load(convert(Ptr{UInt16}, ptr + 2)) & 0x03
    p_wg_x = unsafe_load(convert(Ptr{UInt16}, ptr + 4))
    p_wg_y = unsafe_load(convert(Ptr{UInt16}, ptr + 6))
    p_wg_z = unsafe_load(convert(Ptr{UInt16}, ptr + 8))
    # UInt16 reserved
    p_gr_x = unsafe_load(convert(Ptr{UInt32}, ptr + 12))
    p_gr_y = unsafe_load(convert(Ptr{UInt32}, ptr + 16))
    p_gr_z = unsafe_load(convert(Ptr{UInt32}, ptr + 20))

    p_pseg_size = unsafe_load(convert(Ptr{UInt32}, ptr + 24))
    p_gseg_size = unsafe_load(convert(Ptr{UInt32}, ptr + 28))
    p_kobj_addr = unsafe_load(convert(Ptr{UInt64}, ptr + 32))
    p_karg_addr = unsafe_load(convert(Ptr{UInt64}, ptr + 40))
    # UInt64 reserved
    p_comp_sign = unsafe_load(convert(Ptr{hsa_signal_t}, ptr + 56))

    res = KernelDispatchPacket(
        p_kobj_addr,
        (p_gr_x, p_gr_y, p_gr_z),
        wg_size = (p_wg_x, p_wg_y, p_wg_z),
        dimensions = p_dims,
        header = p_hdr,
        private_segment_size = p_pseg_size,
        group_segment_size = p_gseg_size,
        completion_signal = p_comp_sign
        )
    res.kernarg_address = p_karg_addr

    return res
end

function unsafe_store!(ptr :: Ptr{Void}, dp :: KernelDispatchPacket)
    if dp.header.typ != PacketTypeKernelDispatch
        error("not a dispatch packet")
    end

    unsafe_store!(convert(Ptr{UInt16}, ptr + 2), convert(UInt16, dp.dimensions)) # setup
    unsafe_store!(convert(Ptr{UInt16}, ptr + 4), dp.workgroup_size_x)
    unsafe_store!(convert(Ptr{UInt16}, ptr + 6), dp.workgroup_size_y)
    unsafe_store!(convert(Ptr{UInt16}, ptr + 8), dp.workgroup_size_z)
    unsafe_store!(convert(Ptr{UInt16}, ptr + 10), 0x0000)    # UInt16 reserved
    unsafe_store!(convert(Ptr{UInt32}, ptr + 12), dp.grid_size_x)
    unsafe_store!(convert(Ptr{UInt32}, ptr + 16), dp.grid_size_y)
    unsafe_store!(convert(Ptr{UInt32}, ptr + 20), dp.grid_size_z)
    unsafe_store!(convert(Ptr{UInt32}, ptr + 24), dp.private_segment_size)
    unsafe_store!(convert(Ptr{UInt32}, ptr + 28), dp.group_segment_size)
    unsafe_store!(convert(Ptr{UInt64}, ptr + 32), dp.kernel_object)
    unsafe_store!(convert(Ptr{UInt64}, ptr + 40), dp.kernarg_address)
    unsafe_store!(convert(Ptr{UInt64}, ptr + 48), 0x00)    # UInt64 reserved
    unsafe_store!(convert(Ptr{UInt64}, ptr + 56), dp.completion_signal.handle)

    unsafe_store!(ptr, dp.header)
end

export AgentDispatchPacket

type AgentDispatchPacket <: AQLPacket
    header :: PacketHeader
    typ :: UInt16
    return_address :: UInt64
    arg :: Array{UInt64, 1}
    completion_signal :: hsa_signal_t

    function AgentDispatchPacket(typ;
        return_address = 0,
        header :: PacketHeader = PacketHeader(PacketTypeAgentDispatch),
        args :: Array{UInt64,1} = Array(UInt64, 4),
        completion_signal = hsa_signal_t(0)
        )

        assert(header.typ == PacketTypeAgentDispatch)
        assert(length(args) == 4)

        new(
            header,
            typ,
            return_address,
            args,
            completion_signal
        )
    end
end

function load(::Type{AQLPacket}, ::Type{Val{PacketTypeAgentDispatch}}, ptr :: Ptr{Void}, p_hdr :: PacketHeader)
    p_type = unsafe_load(convert(Ptr{UInt16}, ptr + 2))
    # UInt32 reserved
    p_retu = unsafe_load(convert(Ptr{UInt64}, ptr + 8))
    p_arg1 = unsafe_load(convert(Ptr{UInt64}, ptr + 16))
    p_arg2 = unsafe_load(convert(Ptr{UInt64}, ptr + 24))
    p_arg3 = unsafe_load(convert(Ptr{UInt64}, ptr + 32))
    p_arg4 = unsafe_load(convert(Ptr{UInt64}, ptr + 40))
    # UInt64 reserved
    p_comp = unsafe_load(convert(Ptr{UInt64}, ptr + 56))

    return AgentDispatchPacket(p_type;
        header = p_hdr,
        return_address = p_retu,
        args = UInt64[p_arg1, p_arg2, p_arg3, p_arg4],
        completion_signal = hsa_signal_t(p_comp))
end

function unsafe_store!(ptr :: Ptr{Void}, ad :: AgentDispatchPacket)
    unsafe_store!(convert(Ptr{UInt16}, ptr + 2), ad.typ)
    unsafe_store!(convert(Ptr{UInt32}, ptr + 4), 0x00000000) # UInt32 reserved
    unsafe_store!(convert(Ptr{UInt64}, ptr + 8), ad.return_address)
    unsafe_copy!(
        convert(Ptr{UInt64}, ptr + 16),
        convert(Ptr{UInt64}, pointer(ad.arg)),
        8 # Bytes
        )
    unsafe_store!(convert(Ptr{UInt64}, ptr + 48), 0x0000000000000000)    # UInt64 reserved
    unsafe_store!(convert(Ptr{UInt64}, ptr + 56), ad.completion_signal.handle)

    unsafe_store!(ptr, ad.header)
end

export BarrierPacket

type BarrierPacket <: AQLPacket
    header :: PacketHeader
    is_or :: Bool
    dep_signal :: Array{UInt64, 1}
    completion_signal :: hsa_signal_t
end

function load(::Type{AQLPacket}, ::Type{Val{PacketTypeBarrierAnd}}, ptr :: Ptr{Void}, p_hdr :: PacketHeader)
    p_dep = Array(UInt64, 5)
    p_dep_ptr = convert(Ptr{UInt64}, pointer(p_dep))

    unsafe_copy!(p_dep_ptr, convert(Ptr{UInt64}, ptr + 8), 5)
    p_comp = unsafe_load(convert(Ptr{UInt64}, ptr + 56))

    return BarrierPacket(p_hdr, false, p_dep, p_comp)
end

function unsafe_store!(ptr :: Ptr{Void}, bp :: BarrierPacket)
    p_dep_ptr = convert(Ptr{UInt64}, pointer(bp.dep_signal))

    unsafe_store!(convert(Ptr{UInt16}, ptr + 2), 0)
    unsafe_store!(convert(Ptr{UInt32}, ptr + 4), 0)
    unsafe_copy!(
        convert(Ptr{UInt64}, ptr + 8),
        p_dep_ptr,
        5
    )
    unsafe_store!(convert(Ptr{UInt64}, ptr + 48), 0)
    unsafe_store!(convert(Ptr{UInt64}, ptr + 56), bp.completion_signal)

    unsafe_store!(ptr, bp.header)
end

export InvalidPacket

type InvalidPacket
    header :: PacketHeader
    bytes :: Array{UInt8, 1}

    function InvalidPacket(hdr, bytes)
        @assert hdr.typ == HSA.PacketTypeInvalid
        @assert length(bytes) == 64

        return new(hdr, bytes)
    end
end

function load(::Type{AQLPacket}, ::Type{Val{PacketTypeInvalid}}, ptr :: Ptr{Void}, p_hdr :: PacketHeader)
    bytes = Array(UInt8, 64)

    unsafe_copy!(pointer(bytes), convert(Ptr{UInt8}, ptr), 64)

    return InvalidPacket(p_hdr, bytes)
end

function convert(::Type{KernelDispatchPacket}, p :: InvalidPacket)
    ptr = convert(Ptr{Void}, pointer(p.bytes))
    return load(AQLPacket, Val{PacketTypeKernelDispatch}, ptr, p.header)
end

function unsafe_convert(::Type{AQLPacket}, ptr :: Ptr{Void})
    hdr = unsafe_convert(PacketHeader, ptr)

    packet = load(AQLPacket, Val{hdr.typ}, ptr, hdr)

    return packet
end


text:
using Gnuplot
Θ = LinRange(0, 2π, 100) # 50
Φ = LinRange(0, π, 20)
r = 0.8
x = [r * cos(θ) * sin(ϕ)      for θ in Θ, ϕ in Φ]
y = [r * sin(θ) * sin(ϕ)      for θ in Θ, ϕ in Φ]
z = [r * cos(ϕ) for θ in Θ, ϕ in Φ]
@gsp "set pm3d depthorder" :-
@gsp :- "set style fill transparent solid 0.5" "set view equal xyz"
@gsp :-  x y z "w pm3d notit" "set xyplane 0" palette(:viridis)
save(term="pngcairo size 600,600", output="sphere2.png") # hide

max_stars_repo_name: lazarusA/gnuplot-examples

max_issues_repo_name: lazarusA/gnuplot-examples

max_forks_repo_name: lazarusA/gnuplot-examples


text:
# struct TableauVPRKpLegendre{T} <: AbstractPartitionedTableau{T}
#     @HeaderTableau
#
#     q::Tableau{T}
#     p::Tableau{T}
#
#     R∞::Int
#
#     ω::Matrix{T}
#     d::Vector{T}
#
#     function TableauVPRKpLegendre{T}(name, o, q, p, R∞, ω, d) where {T}
#         @assert q.s == p.s == length(d)
#         new(name, o, q.s, q, p, R∞, ω, d)
#     end
#
#     function TableauVPRKpLegendre{T}(name, o, q, p, R∞, ω) where {T}
#         @assert q.s == p.s
#         new(name, o, q.s, q, p, R∞, ω)
#     end
# end
#
# function TableauVPRKpLegendre(name::Symbol, order::Int, q::Tableau{T}, p::Tableau{T}, R∞::Int, ω::Matrix{T}, d::Vector{T}) where {T}
#     TableauVPRKpLegendre{T}(name, order, q, p, R∞, ω, d)
# end
#
# function TableauVPRKpLegendre(name::Symbol, order::Int, q::Tableau{T}, p::Tableau{T}, R∞::Int, ω::Matrix{T}) where {T}
#     TableauVPRKpLegendre{T}(name, order, q, p, R∞, ω)
# end
#
# function TableauVPRKpLegendre(name::Symbol, order::Int, q::Tableau{T}, R∞::Int, ω::Matrix{T}, d::Vector{T}) where {T}
#     TableauVPRKpLegendre{T}(name, order, q, get_symplectic_conjugate_coefficients(q), R∞, ω, d)
# end
#
# function TableauVPRKpLegendre(name::Symbol, order::Int, q::Tableau{T}, R∞::Int, ω::Matrix{T}) where {T}
#     TableauVPRKpLegendre{T}(name, order, q, get_symplectic_conjugate_coefficients(q), R∞, ω)
# end

"Parameters for right-hand side function of Variational Partitioned Runge-Kutta methods."
const ParametersVPRKpLegendre = AbstractParametersVPRK{:vprk_plegendre}

"Variational special partitioned additive Runge-Kutta integrator."
struct IntegratorVPRKpLegendre{DT, TT, D, S,
                PT <: ParametersVPRKpLegendre{DT,TT},
                ST <: NonlinearSolver{DT},
                IT <: InitialGuessIODE{TT}} <: AbstractIntegratorVPRK{DT,TT,D,S}
    params::PT
    solver::ST
    iguess::IT
    caches::CacheDict{PT}

    function IntegratorVPRKpLegendre(params::ParametersVPRKpLegendre{DT,TT,D,S}, solver::ST, iguess::IT, caches) where {DT,TT,D,S,ST,IT}
        new{DT, TT, D, S, typeof(params), ST, IT}(params, solver, iguess, caches)
    end

    function IntegratorVPRKpLegendre{DT,D}(equations::NamedTuple, tableau::TableauVPRK{TT}, Δt::TT) where {DT,TT,D}
        # get number of stages
        S = tableau.s

        N = (3*S+2)*D

        if isdefined(tableau, :d) && length(tableau.d) > 0
            N += D
        end

        # create params
        params = ParametersVPRKpLegendre{DT,D}(equations, tableau, Δt)

        # create cache dict
        caches = CacheDict(params)

        # create solver
        solver = create_nonlinear_solver(DT, N, params, caches)

        # create initial guess
        iguess = InitialGuessIODE(get_config(:ig_extrapolation), equations[:v̄], equations[:f̄], Δt)

        # create integrator
        IntegratorVPRKpLegendre(params, solver, iguess, caches)
    end

    function IntegratorVPRKpLegendre(equation::Union{IODE{DT},LODE{DT}}, tableau, Δt; kwargs...) where {DT}
        IntegratorVPRKpLegendre{DT, equation.d}(get_functions(equation), tableau, Δt; kwargs...)
    end
end

function Integrators.initialize!(int::IntegratorVPRKpLegendre, sol::AtomicSolutionPODE)
    sol.t̄ = sol.t - timestep(int)

    equation(int, :v̄)(sol.t, sol.q, sol.v)
    equation(int, :f̄)(sol.t, sol.q, sol.v, sol.f)

    initialize!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f,
                            sol.t̄, sol.q̄, sol.p̄, sol.v̄, sol.f̄)
end

function initial_guess!(int::IntegratorVPRKpLegendre{DT,TT}, sol::AtomicSolutionPODE{DT,TT},
                        cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT,TT}
    for i in eachstage(int)
        evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                              sol.q, sol.p, sol.v, sol.f,
                              cache.q̃, cache.p̃, cache.v, cache.f,
                              tableau(int).q.c[i], tableau(int).p.c[i])

        for k in eachdim(int)
            offset = 3*(ndims(int)*(i-1)+k-1)
            int.solver.x[offset+1] = (cache.q̃[k] - sol.q[k]) / timestep(int)
            int.solver.x[offset+2] = (cache.p̃[k] - sol.p[k]) / timestep(int)
            int.solver.x[offset+3] = cache.v[k]
        end
    end

    evaluate!(int.iguess, sol.q, sol.p, sol.v, sol.f,
                          sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                          cache.q̃, cache.p̃,
                          one(TT), one(TT))

    offset = (3*nstages(int)+0)*ndims(int)
    for k in eachdim(int)
        int.solver.x[offset+k] = cache.q̃[k]
    end

    offset = (3*nstages(int)+1)*ndims(int)
    for k in eachdim(int)
        int.solver.x[offset+k] = cache.p̃[k]
    end

    if isdefined(tableau(int), :d)  && length(tableau(int).d) > 0
        offset = (3*nstages(int)+2)*ndims(int)
        for k in eachdim(int)
            int.solver.x[offset+k] = 0
        end
    end
end

function compute_stages!(y::Vector{ST}, Q, V, P, F, Y, Z, Φ, q, v, p, ϕ, μ,
                params::ParametersVPRKpLegendre{DT,TT,D,S}) where {ST,DT,TT,D,S}

    local offset::Int
    local tqᵢ::TT

    for i in 1:S
        # copy y to Q and V
        for k in 1:D
            offset = 3*(D*(i-1)+k-1)
            Y[i][k] = y[offset+1]
            Z[i][k] = y[offset+2]
            V[i][k] = y[offset+3]

            Q[i][k] = params.q̄[k] + params.Δt * Y[i][k]
            P[i][k] = params.p̄[k] + params.Δt * Z[i][k]
        end

        # compute P=p(t,Q) and F=f(t,Q,V)
        tqᵢ = params.t̄ + params.Δt * params.tab.q.c[i]
        params.equ[:ϑ](tqᵢ, Q[i], V[i], Φ[i])
        params.equ[:f](tqᵢ, Q[i], V[i], F[i])
    end

    # copy y to q and p
    for k in 1:D
        q[k] = y[(3*S+0)*D+k]
        p[k] = y[(3*S+1)*D+k]
    end

    # compute p=p(t,q)
    params.equ[:ϑ](params.t̄ + params.Δt, q, v, ϕ)

    # for Lobatto-type methods, copy y to μ
    if isdefined(params.tab, :d) && length(params.tab.d) > 0
        offset = (3*S+2)*D
        for k in 1:D
            μ[k] = y[offset+k]
        end
    end
end

function compute_rhs!(b::Vector{ST},
                      Q::Vector{Vector{ST}}, V::Vector{Vector{ST}},
                      P::Vector{Vector{ST}}, F::Vector{Vector{ST}},
                      Y::Vector{Vector{ST}}, Z::Vector{Vector{ST}},
                      Φ::Vector{Vector{ST}},
                      q::Vector{ST}, p::Vector{ST}, ϕ::Vector{ST}, μ::Vector{ST},
                      params::ParametersVPRKpLegendre{DT,TT,D,S}) where {ST,DT,TT,D,S}
    local ω::ST
    local y::ST
    local z::ST

    # compute b = - (Q-q-AV)
    for i in 1:S
        offset = 0*S*D + D*(i-1)
        for k in 1:D
            y = 0
            for j in 1:S
                y += params.tab.q.a[i,j] * V[j][k]
            end
            b[offset+k] = Y[i][k] - y
        end
    end

    # compute b = - (P-p-AF)
    for i in 1:S
        offset = 1*S*D + D*(i-1)
        for k in 1:D
            z = 0
            for j in 1:S
                z += params.tab.p.a[i,j] * F[j][k]
            end
            b[offset+k] = Z[i][k] - z
        end
    end

    # compute b = - (P-Φ)
    for i in 1:S-1
        offset = 2*S*D + D*(i-1)
        for k in 1:D
            z = 0
            for j in 1:S
                ω = params.tab.p.b[j] * params.tab.p.c[j]^(i-1)
                # ω = params.tab.q.a[i+1,j]
                z += ω * (Φ[j][k] - P[j][k])
            end
            b[offset+k] = z
        end
    end

    # compute b = - (p-ϕ)
    offset = (3*S-1)*D
    for k in 1:D
        b[offset+k] = ϕ[k] - p[k]
    end

    # compute b = - (q-q-AV)
    offset = (3*S+0)*D
    for k in 1:D
        y = 0
        for j in 1:S
            y += params.tab.q.b[j] * V[j][k]
        end
        b[offset+k] = (params.q̄[k] - q[k]) + params.Δt * y
    end

    # compute b = - (p̄-p-AF)
    offset = (3*S+1)*D
    for k in 1:D
        z = 0
        for j in 1:S
            z += params.tab.p.b[j] * F[j][k]
        end
        b[offset+k] = (params.p̄[k] - p[k]) + params.Δt * z
    end

    # if isdefined(params.tab, :d) && length(params.tab.d) > 0
    #     offset = (3*S+2)*D
    #     for k in 1:D
    #         b[offset+k] = 0
    #         for i in 1:S
    #             b[offset+k] -= V[k,i] * params.tab.d[i]
    #         end
    #     end
    # end

    if isdefined(params.tab, :d) && length(params.tab.d) > 0
        sl     = div(S+1, 2)
        offset = D*(S+sl-1)

        # compute μ
        z = params.tab.p.b[sl] / params.tab.d[sl]
        for k in 1:D
            μ[k] = z * b[offset+k]
        end

        # replace equation for Pₗ with constraint on V
        for k in 1:D
            z = 0
            for i in 1:S
                z += V[k,i] * params.tab.d[i]
            end
            b[offset+k] = z
        end

        # modify P₁, ..., Pₛ except for Pₗ
        for i in 1:S
            if i ≠ sl
                offset = D*(S+i-1)
                z = params.tab.d[i] / params.tab.p.b[i]
                for k in 1:D
                    b[offset+k] -= z * μ[k]
                end
            end
        end
    end
end

"Compute stages of variational special partitioned additive Runge-Kutta methods."
function Integrators.function_stages!(y::Vector{ST}, b::Vector{ST},
                params::ParametersVPRKpLegendre{DT,TT,D,S},
                caches::CacheDict) where {ST,DT,TT,D,S}

    @assert length(y) == length(b)

    # get cache for internal stages
    cache = caches[ST]

    compute_stages!(y, cache.Q, cache.V, cache.P, cache.F, cache.Y, cache.Z, cache.Φ, cache.q̃, cache.ṽ, cache.p̃, cache.ϕ, cache.μ, params)

    compute_rhs!(b, cache.Q, cache.V, cache.P, cache.F, cache.Y, cache.Z, cache.Φ, cache.q̃, cache.p̃, cache.ϕ, cache.μ, params)

    # debug output
    # println()
    # for k in 1:D
    #     println(params.q[k], ",  ", params.p[k], ",  ", cache.Q[k,1], ",  ", cache.P[k,1], ",  ", cache.q[k], ",  ", cache.p[k])
    # end
    # println()
end

function Integrators.integrate_step!(int::IntegratorVPRKpLegendre{DT,TT}, sol::AtomicSolutionPODE{DT,TT},
                                     cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT,TT}
    local offset::Int

    # update nonlinear solver parameters from cache
    update_params!(int.params, sol)

    # compute initial guess
    initial_guess!(int, sol, cache)

    # debug output
    # println()
    # println(int.solver.x)

    # reset solution
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute vector fields at internal stages
    compute_stages!(int.solver.x, cache.Q, cache.V, cache.P,
                    cache.F, cache.Y, cache.Z, cache.Φ,
                    cache.q̃, cache.ṽ, cache.p̃, cache.ϕ,
                    cache.μ, int.params)

    # compute final update
    update_solution!(sol.q, cache.V, tableau(int).q.b, tableau(int).q.b̂, timestep(int))
    update_solution!(sol.p, cache.F, tableau(int).p.b, tableau(int).p.b̂, timestep(int))

    # copy solution to initial guess for next time step
    update_vector_fields!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f)
end


text:
type NoRedeployModel <: RedeployModel
    assignment::Vector{Int} # which location the ambulance is assigned to
    ambulances::Vector{Vector{Int}} # list of ambulances assigned to each location
    status::Vector{Symbol} # the current status of the ambulance
        # possible statuses: :available, :responding, :atscene, :conveying, :returning
    fromtime::Vector{Int} # the time it started the new status
    hospital::Vector{Int} # the hospital the ambulance is at (0 otherwise)
end

function NoRedeployModel(
        p::DeploymentProblem,
        available::Vector{Int},
        # utilization::Vector{Float64},
        hospitals::DataFrame,
        stations::DataFrame;
        lambda::Float64 = 100.0
    )

    nambulances = sum(available)
    nlocations = length(available)
    assignment = zeros(Int, nambulances)

    k = 1
    ambulances = [Int[] for i in 1:nlocations]
    for i in eachindex(available), j in 1:available[i]
        assignment[k] = i
        push!(ambulances[i], k)
        k += 1
    end
    @assert k == nambulances + 1

    @assert sum(length(a) for a in ambulances) == nambulances
    status = fill(:available, nambulances)
    fromtime = zeros(Int, nambulances)
    hospital = zeros(Int, nambulances)

    NoRedeployModel(assignment, ambulances, status, fromtime, hospital)
end

reassign_ambulances!(ems, problem::DispatchProblem, redeploy::NoRedeployModel, t::Int) = nothing