struct ODEInterpolatingAdjointSensitivityFunction{C<:AdjointDiffCache,Alg<:InterpolatingAdjoint,
                                                  uType,SType,CPS,pType,fType<:DiffEqBase.AbstractDiffEqFunction} <: SensitivityFunction
  diffcache::C
  sensealg::Alg
  discrete::Bool
  y::uType
  sol::SType
  checkpoint_sol::CPS
  prob::pType
  f::fType
  noiseterm::Bool
end

mutable struct CheckpointSolution{S,I,T}
  cpsol::S # solution in a checkpoint interval
  intervals::I # checkpoint intervals
  cursor::Int # sol.prob.tspan = intervals[cursor]
  tols::T
end

function ODEInterpolatingAdjointSensitivityFunction(g,sensealg,discrete,sol,dg,f,checkpoints,tols;noiseterm=false)
  tspan = reverse(sol.prob.tspan)
  checkpointing = ischeckpointing(sensealg, sol)
  (checkpointing && checkpoints === nothing) && error("checkpoints must be passed when checkpointing is enabled.")

  checkpoint_sol = if checkpointing
    intervals = map(tuple, @view(checkpoints[1:end-1]), @view(checkpoints[2:end]))
    interval_end = intervals[end][end]
    tspan[1] > interval_end && push!(intervals, (interval_end, tspan[1]))
    cursor = lastindex(intervals)
    interval = intervals[cursor]

    if typeof(sol.prob) <: SDEProblem
      # replicated noise
      _sol = deepcopy(sol)
      sol.W.save_everystep = false
      _sol.W.save_everystep = false
      idx1 = searchsortedfirst(_sol.t, interval[1]-1000eps(interval[1]))
      #idx2 = searchsortedfirst(_sol.t, interval[2])
      #forwardnoise = DiffEqNoiseProcess.NoiseGrid(_sol.t[idx1:idx2], _sol.W.W[idx1:idx2])
      forwardnoise = DiffEqNoiseProcess.NoiseWrapper(_sol.W, indx=idx1)
      dt = abs(_sol.W.dt)
      if dt < 1000eps(_sol.t[end])
        dt = interval[2] - interval[1]
      end
      cpsol = solve(remake(sol.prob, tspan=interval, u0=sol(interval[1]), noise=forwardnoise), sol.alg, save_noise=false; dt=dt, tstops=_sol.t[idx1:end] ,tols...)
    else
      cpsol = solve(remake(sol.prob, tspan=interval, u0=sol(interval[1])), sol.alg; tols...)
    end
    CheckpointSolution(cpsol, intervals, cursor, tols)
  else
    nothing
  end

  diffcache, y = adjointdiffcache(g,sensealg,discrete,sol,dg,f;quad=false,noiseterm=noiseterm)

  return ODEInterpolatingAdjointSensitivityFunction(diffcache,sensealg,
                                                    discrete,y,sol,
                                                    checkpoint_sol,sol.prob,f,noiseterm)
end

function findcursor(intervals, t)
  # equivalent with `findfirst(x->x[1] <= t <= x[2], intervals)`
  lt(x, t) = <(x[2], t)
  return searchsortedfirst(intervals, t, lt=lt)
end

# u = λ'
# add tstop on all the checkpoints
function (S::ODEInterpolatingAdjointSensitivityFunction)(du,u,p,t)
  @unpack sol, y, checkpoint_sol, discrete, prob, f = S
  idx = length(y)

  if checkpoint_sol === nothing
    if typeof(t) <: ForwardDiff.Dual && eltype(S.y) <: AbstractFloat
      y = sol(t)
    else
      sol(y,t)
    end
  else
    intervals = checkpoint_sol.intervals
    interval = intervals[checkpoint_sol.cursor]
    if !(interval[1] <= t <= interval[2])
      cursor′ = findcursor(intervals, t)
      interval = intervals[cursor′]
      cpsol_t = checkpoint_sol.cpsol.t
      if typeof(t) <: ForwardDiff.Dual && eltype(S.y) <: AbstractFloat
        y = sol(interval[1])
      else
        sol(y, interval[1])
      end
      if typeof(sol.prob) <: SDEProblem
        #idx1 = searchsortedfirst(sol.t, interval[1])
        _sol = deepcopy(sol)
        _sol.W.save_everystep = false
        idx1 = searchsortedfirst(_sol.t, interval[1]-100eps(interval[1]))
        idx2 = searchsortedfirst(_sol.t, interval[2])
        #forwardnoise = DiffEqNoiseProcess.NoiseGrid(_sol.t[idx1:idx2], _sol.W.W[idx1:idx2])
        forwardnoise = DiffEqNoiseProcess.NoiseWrapper(_sol.W, indx=idx1)
        prob′ = remake(prob, tspan=intervals[cursor′], u0=y, noise=forwardnoise)
        dt = abs(cpsol_t[end]-cpsol_t[end-1])
        if dt < 10000eps(cpsol_t[end])
          dt = interval[2] - interval[1]
        end
        cpsol′ = solve(prob′, sol.alg, noise=forwardnoise, save_noise=false; dt=dt, tstops=_sol.t[idx1:idx2], checkpoint_sol.tols...)
      else
        prob′ = remake(prob, tspan=intervals[cursor′], u0=y)
        cpsol′ = solve(prob′, sol.alg; dt=abs(cpsol_t[end] - cpsol_t[end-1]), checkpoint_sol.tols...)
      end
      checkpoint_sol.cpsol = cpsol′
      checkpoint_sol.cursor = cursor′
    end
    checkpoint_sol.cpsol(y, t)
  end

  λ     = @view u[1:idx]
  grad  = @view u[idx+1:end]

  if length(u) == length(du)
    dλ    = @view du[1:idx]
    dgrad = @view du[idx+1:end]

  elseif length(u) != length(du) &&  StochasticDiffEq.is_diagonal_noise(prob) && !isnoisemixing(S.sensealg)
    idx1 = [length(u)*(i-1)+i for i in 1:idx] # for diagonal indices of [1:idx,1:idx]

    dλ    = @view du[idx1]
    dgrad = @view du[idx+1:end,1:idx]

  else
    # non-diagonal noise and noise mixing case
    dλ    = @view du[1:idx,1:idx]
    dgrad = @view du[idx+1:end,1:idx]
  end

  if S.noiseterm
    if length(u) == length(du)
      vecjacobian!(dλ, y, λ, p, t, S, dgrad=dgrad)
    elseif length(u) != length(du) &&  StochasticDiffEq.is_diagonal_noise(prob) && !isnoisemixing(S.sensealg)
      vecjacobian!(dλ, y, λ, p, t, S)
      jacNoise!(λ, y, p, t, S, dgrad=dgrad)
    else
      jacNoise!(λ, y, p, t, S, dgrad=dgrad, dλ=dλ)
    end
  else
    vecjacobian!(dλ, y, λ, p, t, S, dgrad=dgrad)
  end

  dλ .*= -one(eltype(λ))

  discrete || accumulate_cost!(dλ, y, p, t, S, dgrad)
  return nothing
end

# g is either g(t,u,p) or discrete g(t,u,i)
@noinline function ODEAdjointProblem(sol,sensealg::InterpolatingAdjoint,
                                     g,t=nothing,dg=nothing;
                                     checkpoints=sol.t,
                                     callback=CallbackSet(),
                                     reltol=nothing, abstol=nothing,
                                     kwargs...)
  @unpack f, p, u0, tspan = sol.prob
  tspan = reverse(tspan)
  discrete = t != nothing

  numstates = length(u0)
  numparams = p === nothing || p === DiffEqBase.NullParameters() ? 0 : length(p)

  len = numstates+numparams

  λ = p === nothing || p === DiffEqBase.NullParameters() ? similar(u0) : one(eltype(u0)) .* similar(p, len)
  λ .= false

  sense = ODEInterpolatingAdjointSensitivityFunction(g,sensealg,discrete,sol,dg,f,
                                                     checkpoints,
                                                     (reltol=reltol,abstol=abstol))

  init_cb = t !== nothing && tspan[1] == t[end]
  cb = generate_callbacks(sense, g, λ, t, callback, init_cb)
  z0 = vec(zero(λ))
  original_mm = sol.prob.f.mass_matrix
  if original_mm === I || original_mm === (I,I)
    mm = I
  else
    adjmm = copy(sol.prob.f.mass_matrix')
    zzz = similar(adjmm, numstates, numparams)
    fill!(zzz, zero(eltype(zzz)))
    # using concrate I is slightly more efficient
    II = Diagonal(I, numparams)
    mm = [adjmm       zzz
          copy(zzz')   II]
  end

  jac_prototype = sol.prob.f.jac_prototype
  if !sense.discrete || jac_prototype === nothing
    adjoint_jac_prototype = nothing
  else
    _adjoint_jac_prototype = copy(jac_prototype')
    zzz = similar(_adjoint_jac_prototype, numstates, numparams)
    fill!(zzz, zero(eltype(zzz)))
    II = Diagonal(I, numparams)
    adjoint_jac_prototype = [_adjoint_jac_prototype zzz
                             copy(zzz')             II]
  end

  odefun = ODEFunction(sense, mass_matrix=mm, jac_prototype=adjoint_jac_prototype)
  return ODEProblem(odefun,z0,tspan,p,callback=cb)
end

@noinline function SDEAdjointProblem(sol,sensealg::InterpolatingAdjoint,
                                     g,t=nothing,dg=nothing;
                                     checkpoints=sol.t,
                                     callback=CallbackSet(),
                                     reltol=nothing, abstol=nothing,
                                     diffusion_jac=nothing, diffusion_paramjac=nothing,
                                     kwargs...)
  @unpack f, p, u0, tspan = sol.prob
  tspan = reverse(tspan)
  discrete = t != nothing

  if length(unique(round.(checkpoints, digits=13))) != length(checkpoints)
    @warn "The given checkpoints are not unique. To avoid issues in the interpolation the checkpoints were redefined. You may want to check sol.t if default checkpoints were used."
    checkpoints = unique(round.(checkpoints, digits=13))
  end

  p === DiffEqBase.NullParameters() && error("Your model does not have parameters, and thus it is impossible to calculate the derivative of the solution with respect to the parameters. Your model must have parameters to use parameter sensitivity calculations!")
  numstates = length(u0)
  numparams = length(p)

  len = numstates+numparams

  λ = one(eltype(u0)) .* similar(p, len)
  λ .= false

  sense_drift = ODEInterpolatingAdjointSensitivityFunction(g,sensealg,discrete,sol,dg,sol.prob.f,
                                                     checkpoints,(reltol=reltol,abstol=abstol))

  diffusion_function = ODEFunction(sol.prob.g, jac=diffusion_jac, paramjac=diffusion_paramjac)
  sense_diffusion = ODEInterpolatingAdjointSensitivityFunction(g,sensealg,discrete,sol,dg,diffusion_function,
                                                     checkpoints,(reltol=reltol,abstol=abstol);noiseterm=true)

  init_cb = t !== nothing && tspan[1] == t[end]
  cb = generate_callbacks(sense_drift, g, λ, t, callback, init_cb)
  z0 = vec(zero(λ))
  original_mm = sol.prob.f.mass_matrix
  if original_mm === I || original_mm === (I,I)
    mm = I
  else
    adjmm = copy(sol.prob.f.mass_matrix')
    zzz = similar(adjmm, numstates, numparams)
    fill!(zzz, zero(eltype(zzz)))
    # using concrate I is slightly more efficient
    II = Diagonal(I, numparams)
    mm = [adjmm       zzz
          copy(zzz')   II]
  end

  jac_prototype = sol.prob.f.jac_prototype
  if !sense_drift.discrete || jac_prototype === nothing
    adjoint_jac_prototype = nothing
  else
    _adjoint_jac_prototype = copy(jac_prototype')
    zzz = similar(_adjoint_jac_prototype, numstates, numparams)
    fill!(zzz, zero(eltype(zzz)))
    II = Diagonal(I, numparams)
    adjoint_jac_prototype = [_adjoint_jac_prototype zzz
                             copy(zzz')             II]
  end

  sdefun = SDEFunction(sense_drift,sense_diffusion,mass_matrix=mm,jac_prototype=adjoint_jac_prototype)

  # replicated noise
  _sol = deepcopy(sol)
  backwardnoise = DiffEqNoiseProcess.NoiseWrapper(_sol.W, reverse=true)

  if StochasticDiffEq.is_diagonal_noise(sol.prob) && typeof(sol.W[end])<:Number
    # scalar noise case
    noise_matrix = nothing
  else
    noise_matrix = similar(z0,length(z0),numstates)
    noise_matrix .= false
  end

  return SDEProblem(sdefun,sense_diffusion,z0,tspan,p,
    callback=cb,
    noise=backwardnoise,
    noise_rate_prototype = noise_matrix
    )
end


module RTE
using CUDA
using Adapt
using ..AtmosphericStates
using DocStringExtensions
import GaussQuadrature
using ..Sources
using ..Fluxes
using ..Optics
using ..BCs

import ..Parameters as RP

export Solver

"""
    struct Solver{
        FT,
        I,
        FTA1D,
        FTA2D,
        AS,
        OP,
        SL,
        SS,
        BCL,
        BCS,
        FXBL,
        FXBS,
        FXL,
        FXS,
    }

The high-level RRTMGP data structure storing the atmospheric state,
optical properties, sources, boundary conditions and fluxes configurations
for a given simulation.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct Solver{
    FT <: AbstractFloat,
    I <: Int,
    FTA1D <: AbstractArray{FT, 1},
    FTA2D <: AbstractArray{FT, 2},
    AS <: AbstractAtmosphericState{FT, I, FTA1D},
    OP <: AbstractOpticalProps{FT, FTA2D},
    SL <: Union{AbstractSourceLW{FT, FTA1D, FTA2D}, Nothing},
    SS <: Union{SourceSW2Str{FT, FTA1D, FTA2D}, Nothing},
    BCL <: Union{LwBCs{FT}, Nothing},
    BCS <: Union{SwBCs{FT}, Nothing},
    FXBL <: Union{FluxLW{FT, FTA2D}, Nothing},
    FXBS <: Union{FluxSW{FT, FTA2D}, Nothing},
    FXL <: Union{FluxLW{FT, FTA2D}, Nothing},
    FXS <: Union{FluxSW{FT, FTA2D}, Nothing},
}
    as::AS         # atmospheric state
    op::OP         # optical properties
    src_lw::SL     # source functions
    src_sw::SS     # source functions
    bcs_lw::BCL    # boundary conditions
    bcs_sw::BCS    # boundary conditions
    fluxb_lw::FXBL # temporary storage for bandwise calculations
    fluxb_sw::FXBS # temporary storage for bandwise calculations
    flux_lw::FXL   # fluxes for longwave problem
    flux_sw::FXS   # fluxes for shortwave problem
end

Solver(as, op, src_lw, src_sw, bcs_lw, bcs_sw, fluxb_lw, fluxb_sw, flux_lw, flux_sw) = Solver{
    eltype(as.p_lev),
    typeof(as.ncol),
    typeof(as.t_sfc),
    typeof(as.p_lev),
    typeof(as),
    typeof(op),
    typeof(src_lw),
    typeof(src_sw),
    typeof(bcs_lw),
    typeof(bcs_sw),
    typeof(fluxb_lw),
    typeof(fluxb_sw),
    typeof(flux_lw),
    typeof(flux_sw),
}(
    as,
    op,
    src_lw,
    src_sw,
    bcs_lw,
    bcs_sw,
    fluxb_lw,
    fluxb_sw,
    flux_lw,
    flux_sw,
)

Adapt.@adapt_structure Solver
end


## Copyright (c) 2015, Júlio Hoffimann Mendes <juliohm@stanford.edu>
##
## Permission to use, copy, modify, and/or distribute this software for any
## purpose with or without fee is hereby granted, provided that the above
## copyright notice and this permission notice appear in all copies.
##
## THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
## WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
## MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
## ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
## WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
## ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
## OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

function iqsim(training_image::AbstractArray,
               tplsizex::Integer, tplsizey::Integer, tplsizez::Integer,
               gridsizex::Integer, gridsizey::Integer, gridsizez::Integer;
               overlapx=1/6, overlapy=1/6, overlapz=1/6,
               soft=nothing, hard=nothing, tol=.1,
               cut=:dijkstra, path=:rasterup, categorical=false,
               nreal=1, gpu=false, debug=false)

  # GPU setup
  global GPU = gpu ? gpu_setup() : nothing

  # use all CPU cores in FFT
  FFTW.set_num_threads(CPU_CORES)

  # sanity checks
  @assert ndims(training_image) == 3 "training image is not 3D (add ghost dimension for 2D)"
  @assert all(0 .< [tplsizex, tplsizey, tplsizez] .≤ [size(training_image)...]) "invalid template size"
  @assert all([gridsizex, gridsizey, gridsizez] .≥ [tplsizex, tplsizey, tplsizez]) "invalid grid size"
  @assert all(0 .< [overlapx, overlapy, overlapz] .< 1) "overlaps must be in range (0,1)"
  @assert 0 < tol ≤ 1 "tolerance must be in range (0,1]"
  @assert cut ∈ [:dijkstra,:boykov] "invalid cut algorithm"
  @assert path ∈ [:rasterup,:rasterdown,:dilation,:random] "invalid simulation path"
  @assert nreal > 0 "invalid number of realizations"

  # GPU checks
  if gpu
    @assert all([size(training_image)...] .> 1) "GPU support for 3D training images only"
  end

  # soft data checks
  if soft ≠ nothing
    @assert isa(soft, SoftData) || isa(soft, AbstractArray{SoftData})

    # encapsulate single auxiliary variable in an array
    isa(soft, SoftData) && (soft = [soft])

    for aux in soft
      @assert ndims(aux.data) == 3 "soft data is not 3D (add ghost dimension for 2D)"
      @assert all([size(aux.data)...] .≥ [gridsizex, gridsizey, gridsizez]) "soft data size < grid size"
    end
  end

  # hard data checks
  if hard ≠ nothing
    @assert isa(hard, HardData)
    locations = Int[loc[i] for loc in keys(hard), i=1:3]
    @assert all(maximum(locations, 1) .≤ [gridsizex gridsizey gridsizez]) "hard data locations outside of grid"
    @assert all(minimum(locations, 1) .> 0) "hard data locations must be positive indexes"
  end

  # calculate the overlap from given percentage
  overlapx = ceil(Int, overlapx * tplsizex)
  overlapy = ceil(Int, overlapy * tplsizey)
  overlapz = ceil(Int, overlapz * tplsizez)

  # spacing in raster path
  spacingx = tplsizex - overlapx
  spacingy = tplsizey - overlapy
  spacingz = tplsizez - overlapz

  # calculate the number of tiles from grid size
  ntilex = ceil(Int, gridsizex / max(spacingx, 1))
  ntiley = ceil(Int, gridsizey / max(spacingy, 1))
  ntilez = ceil(Int, gridsizez / max(spacingz, 1))

  # simulation grid dimensions
  nx = ntilex * (tplsizex - overlapx) + overlapx
  ny = ntiley * (tplsizey - overlapy) + overlapy
  nz = ntilez * (tplsizez - overlapz) + overlapz

  # total overlap volume in simulation grid
  overlap_volume = nx*ny*nz - (nx - (ntilex-1)overlapx)*
                              (ny - (ntiley-1)overlapy)*
                              (nz - (ntilez-1)overlapz)

  # warn in case of 1-voxel overlaps
  if any(([tplsizex, tplsizey, tplsizez] .>  1) &
         ([overlapx, overlapy, overlapz] .== 1))
    warn("1-voxel overlaps. Check template/overlap configuration.")
  end

  # hard data in grid format
  hardgrid = []; preset = []; activated = []
  if hard ≠ nothing
    hardgrid = zeros(nx, ny, nz)
    preset = falses(nx, ny, nz)
    activated = trues(nx, ny, nz)
    for loc in keys(hard)
      if isnan(hard[loc])
        activated[loc...] = false
      else
        hardgrid[loc...] = hard[loc]
        preset[loc...] = true
      end
    end
    activated[gridsizex+1:nx,:,:] = false
    activated[:,gridsizey+1:ny,:] = false
    activated[:,:,gridsizez+1:nz] = false
  end

  # keep track of hard data and inactive voxels
  skipped = Set(); datum = Set(); rastered = []
  if hard ≠ nothing
    rastered = falses(nx, ny, nz)
    for k=1:ntilez, j=1:ntiley, i=1:ntilex
      # tile corners are given by (iₛ,jₛ,kₛ) and (iₑ,jₑ,kₑ)
      iₛ = (i-1)spacingx + 1
      jₛ = (j-1)spacingy + 1
      kₛ = (k-1)spacingz + 1
      iₑ = iₛ + tplsizex - 1
      jₑ = jₛ + tplsizey - 1
      kₑ = kₛ + tplsizez - 1

      if all(!activated[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ])
        push!(skipped, (i,j,k))
      else
        rastered[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ] = true
        if any(preset[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ])
          push!(datum, (i,j,k))
        end
      end
    end
    rastered[!activated] = false

    # grid must contain active voxels
    any_activated = any(activated[1:gridsizex,1:gridsizey,1:gridsizez])
    @assert any_activated "simulation grid has no active voxel"
  end

  # always work with floating point
  TI = map(Float64, training_image)

  # inactive voxels in the training image
  NaNTI = isnan(TI); TI[NaNTI] = 0

  # perform simplex transform
  simplexTI = Any[TI]; nvertices = 1
  if categorical
    categories = Set(TI[!NaNTI])
    ncategories = nvertices = length(categories) - 1

    @assert categories == Set(0:ncategories) "categories should be labeled 1, 2, 3,..."

    simplexTI = simplex_transform(TI, nvertices)
  end

  # disable tiles in the training image if they contain inactive voxels
  mₜ, nₜ, pₜ = size(TI)
  disabled = falses(mₜ-tplsizex+1, nₜ-tplsizey+1, pₜ-tplsizez+1)
  for nanidx in find(NaNTI)
    iₙ, jₙ, kₙ = ind2sub(size(TI), nanidx)

    # tile corners are given by (iₛ,jₛ,kₛ) and (iₑ,jₑ,kₑ)
    iₛ = max(iₙ-tplsizex+1, 1)
    jₛ = max(jₙ-tplsizey+1, 1)
    kₛ = max(kₙ-tplsizez+1, 1)
    iₑ = min(iₙ, mₜ-tplsizex+1)
    jₑ = min(jₙ, nₜ-tplsizey+1)
    kₑ = min(kₙ, pₜ-tplsizez+1)

    disabled[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ] = true
  end

  # pad soft data and soft transform training image
  softgrid = []; softTI = []
  if soft ≠ nothing
    for aux in soft
      mx, my, mz = size(aux.data)
      lx = min(mx,nx); ly = min(my,ny); lz = min(mz,nz)

      auxpad = padarray(aux.data, [0,0,0], [nx-lx,ny-ly,nz-lz], "symmetric")
      auxpad[isnan(auxpad)] = 0

      push!(softgrid, auxpad)

      auxTI = copy(aux.transform(training_image))

      @assert size(auxTI) == size(TI) "auxiliary TI must have the same size as TI"

      # inactive voxels in the auxiliary training image
      auxTI[NaNTI] = 0

      push!(softTI, auxTI)
    end
  end

  # overwrite path option if data is available
  !isempty(datum) && (path = :datum)

  # select cut algorithm
  boundary_cut = cut == :dijkstra ? dijkstra_cut : boykov_kolmogorov_cut

  # main output is a vector of 3D grids
  realizations = []

  # for each realization we have:
  boundarycuts = [] # boundary cut
  voxelreuse = [] # voxel reuse

  for real=1:nreal
    # allocate memory for current simulation
    simgrid = zeros(nx, ny, nz)
    cutgrid = debug ? zeros(nx, ny, nz) : []

    # keep track of pasted tiles
    pasted = Set()

    # construct simulation path
    simpath = genpath((ntilex,ntiley,ntilez), path, datum)

    # loop simulation grid tile by tile
    for pathidx in simpath
      i, j, k = ind2sub((ntilex,ntiley,ntilez), pathidx)

      # skip tile if all voxels are inactive
      (i,j,k) ∈ skipped && continue

      # if not skipped, proceed and paste tile
      push!(pasted, (i,j,k))

      # tile corners are given by (iₛ,jₛ,kₛ) and (iₑ,jₑ,kₑ)
      iₛ = (i-1)spacingx + 1
      jₛ = (j-1)spacingy + 1
      kₛ = (k-1)spacingz + 1
      iₑ = iₛ + tplsizex - 1
      jₑ = jₛ + tplsizey - 1
      kₑ = kₛ + tplsizez - 1

      # current simulation dataevent
      simdev = simgrid[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ]

      # compute overlap distance
      distance = zeros(mₜ-tplsizex+1, nₜ-tplsizey+1, pₜ-tplsizez+1)
      if overlapx > 1 && (i-1,j,k) ∈ pasted
        ovx = simdev[1:overlapx,:,:]
        xsimplex = categorical ? simplex_transform(ovx, nvertices) : Any[ovx]

        D = convdist(simplexTI, xsimplex)
        distance += D[1:mₜ-tplsizex+1,:,:]
      end
      if overlapx > 1 && (i+1,j,k) ∈ pasted
        ovx = simdev[spacingx+1:end,:,:]
        xsimplex = categorical ? simplex_transform(ovx, nvertices) : Any[ovx]

        D = convdist(simplexTI, xsimplex)
        distance += D[spacingx+1:end,:,:]
      end
      if overlapy > 1 && (i,j-1,k) ∈ pasted
        ovy = simdev[:,1:overlapy,:]
        ysimplex = categorical ? simplex_transform(ovy, nvertices) : Any[ovy]

        D = convdist(simplexTI, ysimplex)
        distance += D[:,1:nₜ-tplsizey+1,:]
      end
      if overlapy > 1 && (i,j+1,k) ∈ pasted
        ovy = simdev[:,spacingy+1:end,:]
        ysimplex = categorical ? simplex_transform(ovy, nvertices) : Any[ovy]

        D = convdist(simplexTI, ysimplex)
        distance += D[:,spacingy+1:end,:]
      end
      if overlapz > 1 && (i,j,k-1) ∈ pasted
        ovz = simdev[:,:,1:overlapz]
        zsimplex = categorical ? simplex_transform(ovz, nvertices) : Any[ovz]

        D = convdist(simplexTI, zsimplex)
        distance += D[:,:,1:pₜ-tplsizez+1]
      end
      if overlapz > 1 && (i,j,k+1) ∈ pasted
        ovz = simdev[:,:,spacingz+1:end]
        zsimplex = categorical ? simplex_transform(ovz, nvertices) : Any[ovz]

        D = convdist(simplexTI, zsimplex)
        distance += D[:,:,spacingz+1:end]
      end

      # disable dataevents that contain inactive voxels
      distance[disabled] = Inf

      # compute hard and soft distances
      auxdistances = []
      if hard ≠ nothing && any(preset[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ])
        harddev = hardgrid[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ]
        hsimplex = categorical ? simplex_transform(harddev, nvertices) : Any[harddev]
        D = convdist(simplexTI, hsimplex, weights=preset[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ])

        # disable dataevents that contain inactive voxels
        D[disabled] = Inf

        push!(auxdistances, D)
      end
      for n=1:length(softTI)
        softdev = softgrid[n][iₛ:iₑ,jₛ:jₑ,kₛ:kₑ]
        D = convdist(Any[softTI[n]], Any[softdev])

        # disable dataevents that contain inactive voxels
        D[disabled] = Inf

        push!(auxdistances, D)
      end

      # current pattern database
      patterndb = isempty(auxdistances) ? find(distance .≤ (1+tol)minimum(distance)) :
                                          relaxation(distance, auxdistances, tol)
      patternprobs = tau_model(patterndb, distance, auxdistances)

      # pick a pattern at random from the database
      idx = sample(patterndb, weights(patternprobs))
      iᵦ, jᵦ, kᵦ = ind2sub(size(distance), idx)

      # selected training image dataevent
      TIdev = TI[iᵦ:iᵦ+tplsizex-1,jᵦ:jᵦ+tplsizey-1,kᵦ:kᵦ+tplsizez-1]

      # boundary cut mask
      M = falses(simdev)
      if overlapx > 1 && (i-1,j,k) ∈ pasted
        A = simdev[1:overlapx,:,:]; B = TIdev[1:overlapx,:,:]
        M[1:overlapx,:,:] |= boundary_cut(A, B, :x)
      end
      if overlapx > 1 && (i+1,j,k) ∈ pasted
        A = simdev[spacingx+1:end,:,:]; B = TIdev[spacingx+1:end,:,:]
        M[spacingx+1:end,:,:] |= flipdim(boundary_cut(flipdim(A, 1), flipdim(B, 1), :x), 1)
      end
      if overlapy > 1 && (i,j-1,k) ∈ pasted
        A = simdev[:,1:overlapy,:]; B = TIdev[:,1:overlapy,:]
        M[:,1:overlapy,:] |= boundary_cut(A, B, :y)
      end
      if overlapy > 1 && (i,j+1,k) ∈ pasted
        A = simdev[:,spacingy+1:end,:]; B = TIdev[:,spacingy+1:end,:]
        M[:,spacingy+1:end,:] |= flipdim(boundary_cut(flipdim(A, 2), flipdim(B, 2), :y), 2)
      end
      if overlapz > 1 && (i,j,k-1) ∈ pasted
        A = simdev[:,:,1:overlapz]; B = TIdev[:,:,1:overlapz]
        M[:,:,1:overlapz] |= boundary_cut(A, B, :z)
      end
      if overlapz > 1 && (i,j,k+1) ∈ pasted
        A = simdev[:,:,spacingz+1:end]; B = TIdev[:,:,spacingz+1:end]
        M[:,:,spacingz+1:end] |= flipdim(boundary_cut(flipdim(A, 3), flipdim(B, 3), :z), 3)
      end

      # paste contributions from simulation grid and training image
      simgrid[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ] = M.*simdev + !M.*TIdev

      # save boundary cut
      debug && (cutgrid[iₛ:iₑ,jₛ:jₑ,kₛ:kₑ] = M)
    end

    # hard data and shape correction
    if hard ≠ nothing
      simgrid[preset] = hardgrid[preset]
      simgrid[!activated] = NaN
      debug && (cutgrid[!activated] = NaN)
    end

    # throw away voxels that are outside of the grid
    simgrid = simgrid[1:gridsizex,1:gridsizey,1:gridsizez]

    # do the same for boundary cut after saving voxel reuse
    debug && push!(voxelreuse, sum(cutgrid)/overlap_volume)
    debug && (cutgrid = cutgrid[1:gridsizex,1:gridsizey,1:gridsizez])

    # save and continue
    push!(realizations, simgrid)
    debug && push!(boundarycuts, cutgrid)
  end

  debug ? (realizations, boundarycuts, voxelreuse) : realizations
end


using Knet
using ArgParse
using JLD
using Combinatorics

include("env.jl")
include("model.jl")
include("data.jl")
include("vocab.jl")

const CAPACITY = 20000
const EPS_INIT = 0.9
const EPS_FINAL = 0.005
const EPS_DECAY = 200

function main(args)
    s = ArgParseSettings()
    s.description = ""

    @add_arg_table s begin
        # load/save files
        ("--loadfile"; default=nothing)
        ("--savefile"; default=nothing)
        ("--task"; default="copy")
        ("--seed"; default=-1; arg_type=Int64; help="random seed")
        ("--optim"; default="Rmsprop()")
        ("--units"; default=200)
        ("--controller"; default="feedforward"; help="feedforward or lstm")
        ("--discount"; default=-1.; arg_type=Float64)
        ("--start"; default=6; arg_type=Int64)
        ("--end"; default=50; arg_type=Int64)
        ("--step"; default=4; arg_type=Int64)
        ("--batchsize"; default=20; arg_type=Int64)
        ("--nvalid"; default=60; arg_type=Int64)
        ("--atype"; default=(gpu()>=0 ? "KnetArray{Float32}":"Array{Float32}"))
        ("--period"; default=100; arg_type=Int64)
        ("--dist"; arg_type=String; default="randn"; help="[randn|xavier]")
        ("--supervised"; action=:store_true; help="if not, q-learning")
        ("--capacity"; default=CAPACITY; arg_type=Int64)
        ("--nsteps"; default=20; arg_type=Int64)
        ("--update"; default=5000; arg_type=Int64)
        ("--nepisodes"; default=5000; arg_type=Int64)
        ("--threshold"; default=0.98; arg_type=Float64)
        ("--noreward"; default=10; arg_type=Int64)
    end

    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true); println(o); flush(STDOUT)
    sr = o[:seed] > 0 ? srand(o[:seed]) : srand()
    o[:atype] = eval(parse(o[:atype]))
    data_generator = get_data_generator(o[:task])

    # init model, params etc.
    w = opts = s2i = i2s = nothing
    a2i, i2a = initvocab(get_actions(o[:task]))
    s2i, i2s = initvocab(get_symbols(o[:task]))
    if o[:loadfile] == nothing
        w = initweights(
            o[:atype],o[:units],length(s2i),length(a2i),o[:controller],o[:dist])
    else
        o[:task] = load(o[:loadfile], "task")
        w = load(o[:loadfile], "w")
        w = Dict(k=>convert(o[:atype],v) for (k,v) in w)
        # opts - not yet!
        # opts = load(o[:loadfile])
        # o[:start] = load(o[:loadfile], "complexity")
    end
    opts = initopts(w,o[:optim])

    # C => complexity
    # c => controller cell
    steps_done = 0
    for C = o[:start]:o[:step]:o[:end]
        # prepare validation data
        seqlen = div(C, complexities[o[:task]])
        data = map(xi->data_generator(seqlen), [1:o[:nvalid]...])
        valid = []; actions = []
        for (input,output,action) in data
            push!(valid, Game(input, output, o[:task]))
            push!(actions, action)
        end
        empty!(data)

        iter = 1
        lossval = 0
        while true
            # get examples for training
            trn = map(xi->data_generator(seqlen), 1:o[:nepisodes])

            # build environments
            games = []; actions = []
            for (input,output,action) in trn
                push!(games, Game(input, output, o[:task]))
                push!(actions, [action..., STOP_ACTION])
            end

            # train network while running episodes
            this_loss = train!(
                w,games,s2i,i2s,a2i,i2a,actions,opts,o[:supervised]; o=o)
            lossval = update_lossval(lossval, this_loss, iter)

            lossval = update_lossval(lossval,this_loss,iter)
            if iter % o[:period] == 0
                accuracy, average_reward = validate(w,valid,s2i,i2s,a2i,i2a;o=o)
                println("(iter:$iter,loss:$lossval,accuracy:$accuracy,reward:$average_reward)")
                if accuracy >= o[:threshold] && o[:savefile] != nothing
                    save(o[:savefile],
                         "w", Dict(k=>Array(v) for (k,v) in w),
                         # need something like above for opts
                         # "opts", opts,
                         "task", o[:task],
                         "complexity", C)
                end

                if accuracy >= o[:threshold]
                    println("$C converged in $iter iterations")
                    break
                end
            end

            iter += 1
        end # while true
    end # one complexity step
end

# train while running episodes
function train!(w,games,s2i,i2s,a2i,i2a,actions,opts,supervised; o=Dict())
    train = true
    run_episodes!(w,games,s2i,i2s,a2i,i2a,train,supervised;
                  o=o,opts=opts,actions=actions)
end

# validate games
function validate(w,games,s2i,i2s,a2i,i2a; o=Dict())
    train = supervised = false
    acc, reward = run_episodes!(w,games,s2i,i2s,a2i,i2a,train,supervised; o=o)
    return acc, reward
end

# this is skeleton for all methods
function run_episodes!(w,games,s2i,i2s,a2i,i2a,train,supervised;
                       o=Dict(), opts=Dict(), actions=[])

    # init state parameters
    atype = get(o, :atype, typeof(w[:wcont]))
    hidden = size(w[:wcont], 1)
    controller = get(o, :controller, "lstm")

    # unrolled
    nsteps = get(o, :nsteps, 20)

    # supervised learning
    if train && supervised
        isempty(actions) && error("actions must not be empty for supervision")

        # (1) prepare input data
        inputs, outputs, masks = make_data(games,s2i,a2i,actions)

        # (2) init controller states
        h, c = initstates(atype, hidden, length(games), controller)

        # (3) train network
        batchloss = sltrain!(w,h,c,inputs,outputs,masks,opts; o=o)

        # (4) return batchloss
        return batchloss
    end

    # needed by qwatkins
    histories = []
    for game in games; push!(histories, []); end

    # run episodes - for both validation and q-learning
    ncorrect = 0
    iter = 1

    cumulative_reward = 0
    for (i,game) in enumerate(games)
        episode_reward = 0
        h, c = initstates(atype, hidden, 1, controller)
        while !game.is_done
            # (1) prepare input data
            input_symbol = read_symbol(game)
            input_action = game.prev_actions[end]
            input = zeros(Cuchar, length(s2i)+length(a2i), 1)
            input[s2i[input_symbol],1] = 1
            input[length(s2i)+a2i[input_action],1] = 1

            # (2) propage controller
            prev_h = prev_c = nothing
            if train && h != nothing && c != nothing
                prev_h = Array(h)
                prev_c = Array(c)
            end
            cout, h, c = propagate(w[:wcont],w[:bcont],convert(atype,input),h,c)

            # (3) take action
            eps = train ? 0.0 : 0.20
            next_action = first(take_action(w[:wact],w[:bact],cout; eps=eps))
            next_action = i2a[next_action]
            move_action, write_action = next_action

            # (4) predict symbol
            y = NO_SYMBOL
            if write_action == WRITE
                y = predict(w[:wsymb],w[:bsymb],cout)
                # symbol = mapslices(indmax, Array(symbol), 1)[1]
                y = indmax(Array(y))
                y = i2s[y]
            end

            move_timestep!(game, y, move_action)
            reward = get_reward(game)
            episode_reward += reward

            # (5) add transition to history if phase is RL train
            if train
                remaining_steps = get_remaining_steps(game)

                # new transition
                transition = Transition(
                    reward,
                    input_symbol,
                    input_action,
                    remaining_steps,
                    prev_h, prev_c,
                    next_action,
                    y,
                    game.is_done
                )

                add_history = true
                if game.is_done
                    add_history = game.gold_tape == game.output_tape
                end

                # push it do history
                add_history && push!(histories[i], transition)
            end
        end

        cumulative_reward += episode_reward
        if game.output_tape == game.gold_tape
            ncorrect += 1
        end
    end

    # reset games
    for (i,game) in enumerate(games)
        reset!(game)
    end

    # q-watkins training if phase is RL train
    if train && !supervised
        discount = get(o, :discount, -1)
        batchsize = get(o, :batchsize, length(games))

        # (1) prepare input data
        batches = make_batches(
            w,histories,s2i,a2i,discount,nsteps,batchsize)

        # (2) train network
        batchloss = rltrain!(w,batches,opts; o=o)

        # (3) return batchloss
        return batchloss
    end

    return ncorrect/length(games), cumulative_reward/length(games)
end

function sltrain!(w,h,c,inputs,outputs,masks,opt; o=Dict())
    dw = similar(w)
    for k in keys(w); dw[k] = similar(w[k]); fill!(dw[k], 0); end

    total = num_samples = 0
    maxlen = get(o, :maxlen, 50)
    for i = 1:maxlen:length(inputs)
        lower = i; upper = min(i+maxlen-1,length(inputs))
        x = inputs[lower:upper]
        y = outputs[lower:upper]
        m = masks[lower:upper]

        values = []
        gloss = slgrad(w,x,y,m,h,c; values=values)

        # track loss and training sample values
        total += values[1]; num_samples += values[2]

        # accumulate gradients
        for k in keys(dw); dw[k] += gloss[k]; end

        # unbox states
        h = AutoGrad.getval(h)
        c = AutoGrad.getval(c)
    end

    # clip with num_samples
    for k in keys(dw); dw[k] = dw[k]/num_samples; end
    # FIXME: try a different thing in here

    # update weights
    update!(w,dw,opt)

    return total/num_samples
end

function rltrain!(w,batches,opt; o=Dict())
    dw = similar(w)
    for k in keys(w); dw[k] = similar(w[k]); fill!(dw[k], 0); end
    atype = get(o, :atype, AutoGrad.getval(typeof(w[:wcont])))

    total = num_samples = 0
    mapconvert(args...) = map(i->convert(atype,i), args)
    asarray(x) = typeof(x) <: Array ? x : typeof(x)[x]
    asarrays(xs...) = map(x->asarray(x), xs)
    for (ts,xs,ys,as,ms,hs,cs) in batches
        ts,ys,as,ms = asarrays(ts,ys,as,ms)
        # @show ts,ys,as,ms
        ts,xs = mapconvert(ts,xs)
        if hs != nothing && cs != nothing
            hs,cs = mapconvert(hs,cs)
        end

        values = []
        gloss = rlgrad(w,ts,xs,ys,as,ms,hs,cs; values=values)
        total += values[1]; num_samples += values[2]
        for k in keys(gloss); dw[k] += gloss[k]; end
    end

    for k in keys(w); dw[k] = dw[k]/num_samples; end
    update!(w,dw,opt)
    return total/num_samples
end

function update_lossval(lossval,batchloss,iter,alpha=0.01)
    if iter < 100
        lossval = (iter-1)*lossval + batchloss
        lossval = lossval / iter
    else
        lossval = alpha * batchloss + (1-alpha) * lossval
    end
    return lossval
end

!isinteractive() && !isdefined(Core.Main, :load_only) && main(ARGS)


module Chuffed

import ConstraintProgrammingExtensions
using Chuffed_jll

const CP = ConstraintProgrammingExtensions

function run_chuffed(args)
    return String(read(`$(Chuffed_jll.fznchuffed()) $(args)`))
end

function Optimizer()
    return CP.FlatZinc.Optimizer(Chuffed_jll.fznchuffed())
end

end # module


using SpikingNeuralNetworks
SNN = SpikingNeuralNetworks
using Evolutionary, Test, Random

function get_trains(p)
    fire = p.records[:fire]
    x, y = Float32[], Float32[]
    for time in eachindex(fire)
        for neuron_id in findall(fire[time])
            push!(x, time)
            push!(y, neuron_id)
        end
    end
    cellsa = Array{Union{Missing,Any}}(undef, 1, Int(findmax(y)[1]))
    nac = Int(findmax(y)[1])
    for (inx, cell_id) in enumerate(1:nac)
        cellsa[inx] = []
    end
    @inbounds for cell_id in unique(y)
        @inbounds for (time, cell) in collect(zip(x, y))
            if Int(cell_id) == cell
                append!(cellsa[Int(cell_id)], time)
            end
        end
    end
    cellsa

end

function Evolutionary.trace!(
    record::Dict{String,Any},
    objfun,
    state,
    population,
    method::GA,
    options,
)
    idx = sortperm(state.fitpop)
    record["fitpop"] = state.fitpop[:]
    record["pop"] = population[:]
end

function get_ranges(ranges)
    ###
    # Code appropriated from:
    # https://github.com/JuliaML/MLPlots.jl/blob/master/src/optional/onlineai.jl
    ###

    lower = Float32[]
    upper = Float32[]
    for (k, v) in ranges
        append!(lower, v[1])
        append!(upper, v[2])
    end
    lower, upper
end

function init_b(lower, upper)
    gene = [] # TODO: should be a typed list one day
    for (i, (l, u)) in enumerate(zip(lower, upper))
        p1 = rand(l:u, 1)
        append!(gene, p1)
    end
    gene
end

function initf(n)
    genesb = []  # TODO: should be a typed list one day
    for i = 1:n
        genes = init_b(lower, upper)
        append!(genesb, [genes])
    end
    genesb
end

function custom_raster(P1::Array, P2::Array)
    y0 = Int32[0]
    X = Float32[]
    Y = Float32[]
    N = 1
    y = 0
    for p in [P1, P2]
        y += 1
        append!(X, P1)
        append!(Y, 1 .+ sum(y0))
        push!(y0, 1)
        append!(X, P2)
        append!(Y, 1 .+ sum(y0))
        push!(y0, 1)
        N += 1
    end
    plt = scatter(
        X,
        Y,
        w = 50,
        m = (1, :black),
        leg = :none,
        marker = :vline,
        xaxis = ("t", (0, Inf)),
        yaxis = ("neuron",),
    )
    y0 = y0[2:end-1]
    !isempty(y0) && hline!(plt, cumsum(y0), linecolor = :red)
    return plt
end

function get_vm(p, simulation_duration)
    vm = p.records[:v]
    vm = [i[1] for i in vm]
    vm = signal(vm, length(vm) / simulation_duration)
    vm
end

function checkmodel(param, cell_type, ngt_spikes)
    if cell_type == "IZHI"
        pp = SNN.IZParameter(; a = param[1], b = param[2], c = param[3], d = param[4])
        E = SNN.IZ(; N = 1, param = pp)
    end
    if cell_type == "ADEXP"
        adparam = SNN.ADEXParameter(;
            a = param[1],
            b = param[2],
            cm = param[3],
            v0 = param[4],
            τ_m = param[5],
            τ_w = param[6],
            θ = param[7],
            delta_T = param[8],
            v_reset = param[9],
            spike_delta = param[10],
        )
        E = SNN.AD(; N = 1, param = adparam)
    end

    ALLEN_DURATION = 2000 * ms
    ALLEN_DELAY = 1000 * ms
    current = current_search(cell_type, param, ngt_spikes)
    E.I = [current * nA]

    SNN.monitor(E, [:v])
    SNN.monitor(E, [:fire])

    ###
    # TODO buid t into neuron models
    # SNN.monitor(E, [:t])
    ###
    simulation_duration = ALLEN_DURATION + ALLEN_DELAY + 343ms
    SNN.sim!(
        [E];
        dt = 1 * ms,
        delay = ALLEN_DELAY,
        stimulus_duration = ALLEN_DURATION,
        simulation_duration = simulation_duration,
    )
    spikes = get_spikes(E)
    spikes = [s / 1000.0 for s in spikes]
    vm = get_vm(E, simulation_duration)
    (vm, spikes)
end

#=
function Evolutionary.value!(
    ::Val{:multiproc},
    fitness,
    objfun,
    population::AbstractVector{IT},
) where {IT}
    fitness = SharedArrays.SharedArray{Float32}(fitness)
    @time @sync @distributed for i = 1:length(population)
        fitness[i] = value(objfun, population[i])
    end
    fitness
end
=#
function Evolutionary.value!(
    ::Val{:serial},
    fitness,
    objfun,
    population::AbstractVector{IT},
) where {IT}
    Threads.@threads for i = 1:length(population)
        fitness[i] = value(objfun, population[i])
    end
end
function get_data()
    file = "../JLD/ground_truth.jld"
    if file
        vmgtv = load(file, "vmgtv")
        ngt_spikes = load(file, "ngt_spikes")
        gt_spikes = load(file, "gt_spikes")

        ground_spikes = gt_spikes
        ngt_spikes = size(gt_spikes)[1]
        vmgtt = load(file, "vmgtt")

        plot(plot(vmgtv, vmgtt, w = 1))

    else
        # if the file doesn't exist complete the expensive operation of making the file.
        # using Python
        py"""
        from neo import AnalogSignal
        from neuronunit.allenapi import make_allen_tests_from_id

        """

        py"""
        specimen_id = (
            325479788,
            324257146,
            476053392,
            623893177,
            623960880,
            482493761,
            471819401
        )
        specimen_id = specimen_id[1]
        target_num_spikes=7
        sweep_numbers, data_set, sweeps = make_allen_tests_from_id.allen_id_to_sweeps(specimen_id)
        (vmm,stimulus,sn,spike_times) = make_allen_tests_from_id.get_model_parts_sweep_from_spk_cnt(
            target_num_spikes, data_set, sweep_numbers, specimen_id
        )
        """
        gt_spikes = py"spike_times"
        ground_spikes = gt_spikes
        ngt_spikes = size(gt_spikes)[1]
        vmgtv = py"vmm.magnitude"
        vmgtt = py"vmm.times"
        julia_version_vm = signal(vmgtv, length(vmgtt) / last(vmgtt))
        plot(plot(vmgtv, vmgtt, w = 1))

        save(
            "ground_truth.jld",
            "vmgtv",
            vmgtv,
            "vmgtt",
            vmgtt,
            "ngt_spikes",
            ngt_spikes,
            "gt_spikes",
            gt_spikes,
            "julia_version_vm",
            julia_version_vm,
        )
    end

    vmgtv = load("ground_truth.jld", "vmgtv")
    ngt_spikes = load("ground_truth.jld", "ngt_spikes")
    ngt_spikes = size(gt_spikes)[1]
    ground_spikes = load("ground_truth.jld", "gt_spikes")

    vmgtt = load("ground_truth.jld", "vmgtt")
    julia_version_vm = load("ground_truth.jld", "julia_version_vm")
    return (vmgtv, vmgtt, ngt_spikes, ground_spikes, julia_version_vm)
end
function get_izhi_ranges()
    ranges_izhi = DataStructures.OrderedDict{Char,Float32}()
    ranges_izhi =
        ("a" => (0.002, 0.3), "b" => (0.02, 0.36), "c" => (-75, -35), "d" => (0.005, 16))#,"I"=>[100,9000])
    lower, upper = get_ranges(ranges_izhi)
    return lower, upper
end

function get_adexp_ranges()
    ranges_adexp = DataStructures.OrderedDict{String,Tuple{Float32,Float32}}()
    ranges_adexp[:"a"] = (2.0, 10)
    ranges_adexp[:"b"] = (5.0, 10)
    ranges_adexp[:"cm"] = (700.0, 983.5)
    ranges_adexp[:"v0"] = (-70, -55)
    ranges_adexp[:"τ_m"] = (10.0, 42.78345)
    ranges_adexp[:"τ_w"] = (300.0, 454.0)  # Tau_w 0, means very low adaption
    ranges_adexp[:"θ"] = (-45.0, -10)
    ranges_adexp[:"delta_T"] = (1.0, 5.0)
    ranges_adexp[:"v_reset"] = (-70.0, -15.0)
    ranges_adexp[:"spike_delta"] = (1.25, 20.0)
    lower, upper = get_ranges(ranges_adexp)
    return lower, upper
end

function vecplot(p, sym)
    v = SNN.getrecord(p, sym)
    y = hcat(v...)'
    x = 1:length(v)
    plot(x, y, leg = :none, xaxis = ("t", extrema(x)), yaxis = (string(sym), extrema(y)))
end

function vecplot(P::Array, sym)
    plts = [vecplot(p, sym) for p in P]
    N = length(plts)
    plot(plts..., size = (600, 400N), layout = (N, 1))
end

function spike_train_difference(spkd_ground, spkd_found)
    if length(spkd_found) == 0
        return sum(ones(Int(length(spkd_ground))))
    end
    if length(spkd_found) > 0
        maxi0 = findmax(spkd_ground[:, :])[1]
        maxi1 = findmax(spkd_found[:, :])[1]
        maxi1 = maxi1[1]
        maxi0 = maxi0[1]
        mini = findmin([maxi0, maxi1])[1]
        l0 = length(spkd_ground)
        l1 = length(spkd_found)
        maxi = findmin([l0, l1])[1]
        spkd = zeros(Int(maxi))
        @inbounds for (_, i) in zip(spkd, eachindex(spkd))
            if !isempty(spkd_ground[i]) && !isempty(spkd_found[i])

                maxt1 = findmax(spkd_ground[i])[1]
                maxt2 = findmax(spkd_found[i])[1]
                maxt = findmax([maxt1, maxt2])[1]

                if maxt1 > 0.0 && maxt2 > 0.0
                    t, S = SpikeSynchrony.SPIKE_distance_profile(
                        unique(sort(spkd_ground[i])),
                        unique(sort(spkd_found[i]));
                        t0 = 0.0,
                        tf = maxt,
                    )
                    spkd[i] = SpikeSynchrony.trapezoid_integral(t, S) / (t[end] - t[1]) # == SPIKE_distance(y1, y2)

                end
            else
                spkd[i] = 0.0
            end
        end
        scatter([i for i = 1:mini], spkd) |> display
        error = sum(spkd)
    end
end

function initd()
    population = initd(10)
    garray = zeros((length(population)[1], length(population[1])))
    for (i, p) in enumerate(population)
        garray[i, :] = p
    end
    garray[1, :]
end

function initd(n)
    genesb = []
    lower = Float32[0.0 0.0 0.0 0.0]
    upper = Float32[1.0 1.0 1.0 1.0]
    lower = vec(lower)
    upper = vec(upper)

    for i = 1:n
        genes = initd(lower, upper)
        append!(genesb, [genes])
    end
    genesb
end

function initd(lower, upper)
    gene = []

    for (i, (l, u)) in enumerate(zip(lower, upper))
        p1 = rand(l:u, 1)
        append!(gene, p1)
    end
    gene
end


"""
	map function for parallelized processing of PGV and PGA

2021.1.4 Kurama Okubo
"""

using CSV, DataFrames

"""
	map_seisPGVPGA(startid, InputDict::Dict)

Download 3 components seismogram and compute PGV, PGA and noise strength d_rms.
"""
function map_seisPGVPGA(startid, InputDict::Dict)

	download_time_unit = InputDict["download_time_unit"]
	fodir              = InputDict["fodir"]
	starttimelist = InputDict["starttimelist"]

	#SeisIO getdata option
	!haskey(InputDict, "get_data_opt") && (InputDict["get_data_opt"] = [true, true, true, true, true])#: [unscale, demean, detrend, taper, ungap]
	!haskey(InputDict, "sampling_frequency") && (InputDict["sampling_frequency"] = false) # default margin: 5 minutes
	!haskey(InputDict, "download_margin") && (InputDict["download_margin"] = 5 * 60) # default margin: 5 minutes

	#show progress
	if starttimelist[startid][end-8:end] == "T00:00:00"
	    println("start downloading $(starttimelist[startid])")
	end

	request_src_chanks = InputDict["requeststations"]

	for src in keys(request_src_chanks)
	    # replace source name
	    if src=="IRISDMC"
			reqsrc="IRIS"
		else
			reqsrc=src
		end

	    #---download data---#
	    requeststrs = get_requeststr(request_src_chanks[src], 1) #  request single station by station

	    for requeststr in requeststrs

    		#---Skip if this request station has been already performed due to restart processing etc.---#
    		s_str = string(DateTime(starttimelist[startid]))[1:19]
	        # compute end time:
	        et = DateTime(starttimelist[startid]) + Second(download_time_unit)
	        e_str=string(et)[1:19]
	        # split network, station, location and channel
	        net, sta, loc, cha = split(requeststr[1], ".")
	        varname = join([net, sta], ".")*"__"*s_str*"__"*e_str

	        # make hierarchical directory
	        dir1 = joinpath(InputDict["pgvpga_dir"], join([net, sta], ".")) #ex. BP.LCCB..BP1
	        ts_year = split(split(varname, "__")[2], "-")[1]# start year
	        dir2 = joinpath(dir1, ts_year) #ex. 2014
	        mkpath(dir2)
	        # dump file
	        if ispath(joinpath(dir2, varname*".dat"))
	        	println("skip $(varname*".dat").")
	        	continue;
	        end

	        #NOTE: method is currently fixed "FDSN", which downloads waveforms.
	        method = "FDSN"

	        # including download margin
	        starttime = string(DateTime(starttimelist[startid]) - Second(InputDict["download_margin"]))
	        dltime = download_time_unit + 2 * InputDict["download_margin"]

	        ex = :(get_data($(method), $(requeststr), s=$(starttime), t=$(dltime),
	         v=$(0), src=$(reqsrc), unscale=$(InputDict["get_data_opt"][1]),
	          demean=$(InputDict["get_data_opt"][2]), detrend=$(InputDict["get_data_opt"][3]),taper=$(InputDict["get_data_opt"][4]),
	          ungap=$(InputDict["get_data_opt"][5]), rr=true))

	        t_dl = @elapsed Stemp = check_and_get_data(ex)
	        # If download is failed, skip
	        Stemp == 1 && continue
	        # If Stemp does not have 3 components, skip
	        Stemp.n != 3 && continue

	        # manipulate download_margin
	        # NOTE: using SeisIO.sync()
	        stsync = DateTime(starttimelist[startid])
	        etsync = u2d(d2u(stsync) + download_time_unit)
	        SeisIO.sync!(Stemp, s=stsync, t=etsync, v=0)

	        #---downsampling---#
	        if InputDict["sampling_frequency"] isa Number
	            for j = 1:Stemp.n
	                if Stemp.misc[j]["dlerror"] == 0 && !isempty(Stemp[j].t)
	                    # downsample
	                    if Stemp.fs[j] > InputDict["sampling_frequency"]
	                        SeisIO.resample!(Stemp, chans=j, fs=float(InputDict["sampling_frequency"]))
	                    end
	                end
	            end
	        end

	        #---compute pgv, pga and d_rms---#
	        pgv, pga, drms_mean = compute_pgvpga_drms(Stemp, InputDict["bp_fmin"], InputDict["bp_fmax"],
	                                                    InputDict["drms_fmin"], InputDict["drms_fmax"])

	        #---output result into dat file---#
	        # dump file
	        open(joinpath(dir2, varname*".dat"), "w") do io
	            write(io, "id,starttime,endtime,PGV,PGA,drms_mean\n")
	            write(io, "$(join([net, sta], ".")), $(s_str), $(e_str), $(pgv), $(pga), $(drms_mean)")
	        end;

	        #---save seisdata---#
	        if InputDict["save_seisdata"]
	            # make hierarchical directory
	            dir1 = joinpath(InputDict["seisdata_dir"], join([net, sta], ".")) #ex. BP.LCCB..BP1
	            ts_year = split(split(varname, "__")[2], "-")[1]# start year
	            dir2 = joinpath(dir1, ts_year) #ex. 2014
	            mkpath(dir2)
	            # dump file
	            wseis(joinpath(dir2, varname*".seisio"), Stemp)
	        end
	    end
	end

end

"""
    map_readPGVPGA(fi_csv::String)

read pgv, pga and d_rms from dat file.
"""
function map_readPGVPGA(fi_csv::String)
    df = CSV.File(fi_csv,  header=1, delim=',') |> DataFrame
    df["midtime"] = u2d((d2u(df["starttime"][1]) + d2u(df["endtime"][1]))/2)
    return df
end


# Analyzing PowerModels' raw parser, it was concluded that b_to & b_fr data was present in DSHL section
function _handle_b(pwf_data::Dict, f_bus::Int, t_bus::Int, susceptance::Float64, circuit::Int, dict_dshl)
    b_fr, b_to = susceptance / 200, susceptance / 200
    if haskey(pwf_data, "DSHL") && haskey(dict_dshl, (f_bus, t_bus)) && haskey(dict_dshl[(f_bus, t_bus)], circuit)
        group = dict_dshl[(f_bus, t_bus)][circuit]
        b_fr += group["SHUNT FROM"] !== nothing && group["STATUS FROM"] == " L" ? group["SHUNT FROM"] / 100 : 0
        b_to += group["SHUNT TO"] !== nothing && group["STATUS TO"] == " L" ? group["SHUNT TO"] / 100 : 0
    end
    return b_fr, b_to
end

function _create_dict_dshl(data::Dict)
    dshl_dict = Dict{Tuple, Any}()

    for (k,v) in data
        sub_data = Dict{String, Any}()
        sub_data["SHUNT FROM"] = v["SHUNT FROM"]
        sub_data["SHUNT TO"] = v["SHUNT TO"]
        sub_data["STATUS FROM"] = v["STATUS FROM"]
        sub_data["STATUS TO"] = v["STATUS TO"]
        current_value = get(dshl_dict, (v["FROM BUS"], v["TO BUS"]), Dict{Int, Any}())
        current_value[v["CIRCUIT"]] = sub_data
        dshl_dict[(v["FROM BUS"], v["TO BUS"])] = current_value
    end
    return dshl_dict
end

function _pwf2pm_branch!(pm_data::Dict, pwf_data::Dict, branch::Dict; add_control_data::Bool=false)
    sub_data = Dict{String,Any}()

    sub_data["f_bus"] = pop!(branch, "FROM BUS")
    sub_data["t_bus"] = pop!(branch, "TO BUS")
    sub_data["br_r"] = pop!(branch, "RESISTANCE") / 100
    sub_data["br_x"] = pop!(branch, "REACTANCE") / 100

    dshl_dict = haskey(pwf_data, "DSHL") ? _create_dict_dshl(pwf_data["DSHL"]) : nothing
    b = _handle_b(pwf_data, sub_data["f_bus"], sub_data["t_bus"], branch["SHUNT SUSCEPTANCE"], branch["CIRCUIT"], dshl_dict)
    sub_data["g_fr"] = 0.0
    sub_data["b_fr"] = b[1]
    sub_data["g_to"] = 0.0
    sub_data["b_to"] = b[2]

    sub_data["tap"] = pop!(branch, "TAP")
    sub_data["shift"] = -pop!(branch, "LAG")
    sub_data["angmin"] = -60.0 # PowerModels.jl standard
    sub_data["angmax"] = 60.0 # PowerModels.jl standard
    sub_data["transformer"] = false

    if branch["STATUS"] == branch["OPENING FROM BUS"] == branch["OPENING TO BUS"] == 'L'
        sub_data["br_status"] = 1
    else
        sub_data["br_status"] = 0
    end

    if add_control_data
        sub_data["control_data"] = Dict{String,Any}()
        sub_data["control_data"]["tapmin"] = sub_data["tap"]
        sub_data["control_data"]["tapmax"] = sub_data["tap"]
        sub_data["control_data"]["circuit"] = branch["CIRCUIT"]

        # Transformer control_data fields are created and set to nothing
        sub_data["control_data"]["control_type"] = nothing
        sub_data["control_data"]["constraint_type"] = nothing
        sub_data["control_data"]["controlled_bus"] = nothing
        sub_data["control_data"]["vmsp"] = nothing
        sub_data["control_data"]["vmmin"] = nothing
        sub_data["control_data"]["vmmax"] = nothing
        sub_data["control_data"]["shift_control_variable"] = nothing
        sub_data["control_data"]["shiftmin"] = nothing
        sub_data["control_data"]["shiftmax"] = nothing
        sub_data["control_data"]["valsp"] = nothing
        sub_data["control_data"]["control"] = nothing
    end

    sub_data["source_id"] = ["branch", sub_data["f_bus"], sub_data["t_bus"], "01"]
    sub_data["index"] = length(pm_data["branch"]) + 1

    sub_data["rate_a"] = pop!(branch, "NORMAL CAPACITY")
    sub_data["rate_b"] = pop!(branch, "EMERGENCY CAPACITY")
    sub_data["rate_c"] = pop!(branch, "EQUIPAMENT CAPACITY")

    if sub_data["rate_a"] >= 9999
        delete!(sub_data, "rate_a")
    end
    if sub_data["rate_b"] >= 9999
        delete!(sub_data, "rate_b")
    end
    if sub_data["rate_c"] >= 9999
        delete!(sub_data, "rate_c")
    end

    idx = string(sub_data["index"])
    pm_data["branch"][idx] = sub_data

end

function _pwf2pm_DCSC_branch!(pm_data::Dict, pwf_data::Dict, branch::Dict; add_control_data::Bool=false)
    sub_data = Dict{String,Any}()

    sub_data["f_bus"] = pop!(branch, "FROM BUS")
    sub_data["t_bus"] = pop!(branch, "TO BUS")
    sub_data["br_r"] = 0
    sub_data["br_x"] = pop!(branch, "INITIAL VALUE") / 100

    sub_data["g_fr"] = 0.0
    sub_data["b_fr"] = 0.0
    sub_data["g_to"] = 0.0
    sub_data["b_to"] = 0.0

    sub_data["tap"] = 1.0
    sub_data["shift"] = 0
    sub_data["angmin"] = -60.0 # PowerModels.jl standard
    sub_data["angmax"] = 60.0 # PowerModels.jl standard
    sub_data["transformer"] = false

    if branch["STATUS"] == 'L'
        sub_data["br_status"] = 1
    else
        sub_data["br_status"] = 0
    end

    if add_control_data
        sub_data["control_data"] = Dict{String,Any}()
        sub_data["control_data"]["tapmin"] = 1.0
        sub_data["control_data"]["tapmax"] = 1.0
        sub_data["control_data"]["circuit"] = branch["CIRCUIT"]

        # Transformer control_data fields are created and set to nothing
        sub_data["control_data"]["control_type"] = nothing
        sub_data["control_data"]["constraint_type"] = nothing
        sub_data["control_data"]["controlled_bus"] = nothing
        sub_data["control_data"]["vmsp"] = nothing
        sub_data["control_data"]["vmmin"] = nothing
        sub_data["control_data"]["vmmax"] = nothing
        sub_data["control_data"]["shift_control_variable"] = nothing
        sub_data["control_data"]["shiftmin"] = nothing
        sub_data["control_data"]["shiftmax"] = nothing
        sub_data["control_data"]["valsp"] = nothing
        sub_data["control_data"]["control"] = nothing
    end

    sub_data["source_id"] = ["branch", sub_data["f_bus"], sub_data["t_bus"], "01"]
    sub_data["index"] = length(pm_data["branch"]) + 1

    sub_data["rate_a"] = pop!(branch, "NORMAL CAPACITY")
    sub_data["rate_b"] = pop!(branch, "EMERGENCY CAPACITY")
    sub_data["rate_c"] = pop!(branch, "EQUIPAMENT CAPACITY")

    if sub_data["rate_a"] >= 9999
        delete!(sub_data, "rate_a")
    end
    if sub_data["rate_b"] >= 9999
        delete!(sub_data, "rate_b")
    end
    if sub_data["rate_c"] >= 9999
        delete!(sub_data, "rate_c")
    end

    rep = findall(x -> x["f_bus"] == sub_data["f_bus"] && x["t_bus"] == sub_data["t_bus"] && x["circuit"] == sub_data["circuit"], pm_data["branch"])
    if length(rep) > 0
        @warn "Branch from $(sub_data["f_bus"]) to $(sub_data["t_bus"]) in circuit $(sub_data["circuit"]) is duplicated"
    end
    idx = string(sub_data["index"])
    pm_data["branch"][idx] = sub_data

end

function _pwf2pm_branch!(pm_data::Dict, pwf_data::Dict; add_control_data::Bool=false)
    pm_data["branch"] = Dict{String, Any}()
    if haskey(pwf_data, "DLIN")
        for (i,branch) in pwf_data["DLIN"]
            if !branch["TRANSFORMER"]
                _pwf2pm_branch!(pm_data, pwf_data, branch, add_control_data = add_control_data)
            end
        end
    end
    if haskey(pwf_data, "DCSC")
        for (i,csc) in pwf_data["DCSC"]
            _pwf2pm_DCSC_branch!(pm_data, pwf_data, csc, add_control_data = add_control_data)
        end
    end
end


using DataFrames
using HDF5

tmp = mktempdir()

n = 100
d1 = 10
d2 = 20

warn_msg = "Cannot join columns with the same name because var_names are intersecting."
warn_msg2 = "Duplicated obs_names should not be present in different modalities due to the ambiguity that leads to."

function make_ads()
    obs_names = ["obs_$i" for i in 1:n]
    ad1 = AnnData(
        X=rand(n, d1),
        obs_names=copy(obs_names),
        var=DataFrame(testcol=["ad1_$i" for i in 1:d1]),
    )

    obs_names[10] = obs_names[n] = "testobs"
    ad2 = AnnData(
        X=rand(n, d2),
        obs_names=copy(obs_names),
        layers=Dict("testlayer" => rand(UInt16, n, d2)),
        varm=Dict("test" => rand(Int8, d2, 5, 2, 10)),
    )

    obs_names[10] = "obs_10"
    ad3 = AnnData(
        X=rand(n, d1),
        obs_names=copy(obs_names),
        var=DataFrame(commoncol=(d1 + 1):(2d1), ad3col=1:d1),
    )
    return ad1, ad2, ad3
end

ad1, ad2, ad3 = make_ads()
expected_n = 102

@testset "create mudata" begin
    md_single = MuData(mod=Dict("ad1" => ad1))
    global md = (@test_logs (:warn, warn_msg) MuData(mod=Dict("ad1" => ad1, "ad2" => ad2)))
    md.var[!, :mutestcol] = rand(size(md, 2))
    md.obsm["mdtest"] = rand(size(md, 1), 5, 3)

    @test size(md, 1) == expected_n
    @test size(md, 2) == d1 + d2

    modalities = sort(collect(keys(md.mod)))
    @test modalities == ["ad1", "ad2"]

    ad1.var[!, :commoncol] = ["ad1_$i" for i in 1:size(ad1, 2)]
    ad1.var_names = ["ad1_$i" for i in 1:size(ad1, 2)]
    ad2.var[!, :commoncol] = ["ad2_$i" for i in 1:size(ad2, 2)]
    ad2.var_names = ["ad2_$i" for i in 1:size(ad2, 2)]
    md = MuData(mod=Dict("ad1" => ad1, "ad2" => ad2))
    md.var[!, :mutestcol] = rand(size(md, 2))
    md.obsm["mdtest"] = rand(size(md, 1), 5, 3)

    @test size(md, 1) == expected_n
    @test size(md, 2) == d1 + d2
    @test sort(names(md.var)) == ["ad1:testcol", "commoncol", "mutestcol"]

    a1, a2, _ = make_ads()
    a1.obs = DataFrame(demo=[1 for i in 1:size(ad1, 1)])
    a2.obs = DataFrame(demo=[2 for i in 1:size(ad2, 1)])
    a2.obs_names = copy(a1.obs_names)
    m = (@test_logs (:warn, warn_msg) MuData(mod=Dict("ad1" => a1, "ad2" => a2)))
    m.obs[!, "demo"] .= "common"
    @test_logs (:warn, warn_msg) update!(m)
    @test sort(names(m.obs)) == ["ad1:demo", "ad2:demo", "demo"]
end

@testset "insert new AnnData" begin
    ad2.var[!, :testcol] = ["ad2_$i" for i in 1:size(ad2, 2)]
    update!(md)
    @test sort(names(md.var)) == ["commoncol", "mutestcol", "testcol"]

    md["ad3"] = ad3
    @test_logs (:warn, warn_msg2) update!(md)
    @test sort(names(md.var)) ==
          ["ad1:testcol", "ad2:testcol", "ad3:ad3col", "commoncol", "mutestcol", "testcol"]
    @test !any(ismissing.(md.var.commoncol))
    @test sum(ismissing.(md.var.mutestcol[md.varm["ad1"]])) ==
          sum(ismissing.(md.var.mutestcol[md.varm["ad2"]])) ==
          0
    @test all(ismissing.(md.var.mutestcol[md.varm["ad3"]]))
end

function test_row_slice(md, i1, n, d, j=:)
    md1, md2, md3, md4 =
        md[i1, :], md[collect(i1), :], md[unique(md.obs_names[i1]), :], md[1:n .∈ (i1,), :]

    @test size(md1) == size(md2) == size(md4) == (length(i1), d) == size(@test_logs (:warn, warn_msg) MuData(mod=md1.mod))
    @test size(md3) == (length(md.obs_names[unique(md.obs_names[i1]), true]), d)
    @test md1.obs_names == md2.obs_names == md4.obs_names == md.obs_names[i1]
    @test sort(unique(md3.obs_names)) == sort(unique(md.obs_names[i1]))

    ad1_names = filter(in(md["ad1"].obs_names), md.obs_names[i1])
    ad2_names = filter(in(md["ad2"].obs_names), md.obs_names[i1])

    ad1_idx = filter(>(0x0), md.obsmap["ad1"][i1])
    ad2_idx = filter(>(0x0), md.obsmap["ad2"][i1])

    ad1, ad2 = parent(md["ad1"])[ad1_idx, j], parent(md["ad2"])[ad2_idx, j]

    @test md1["ad1"].X == md2["ad1"].X == md4["ad1"].X == ad1.X
    @test md1["ad2"].X == md2["ad2"].X == md4["ad2"].X == ad2.X
    @test md1["ad1"].obs_names == md2["ad1"].obs_names == md4["ad1"].obs_names == ad1.obs_names
    @test md1["ad2"].obs_names == md2["ad2"].obs_names == md4["ad2"].obs_names == ad2.obs_names
end

function test_md_slicing(md, n, d, j=:)
    @testset "row slice $(1:i)" for i in 2:n
        test_row_slice(md, 1:i, n, d, j)
    end
    @testset "row slice $(i:(i + 1))" for i in 1:(n - 1)
        test_row_slice(md, i:(i + 1), n, d, j)
    end
end

ad1, ad2, _ = make_ads()
md = (@test_logs (:warn, warn_msg) MuData(mod=Dict("ad1" => ad1, "ad2" => ad2)))
md.var[!, :mutestcol] = rand(size(md, 2))
md.obsm["mdtest"] = rand(size(md, 1), 5, 3)

@testset "slicing mudata" begin
    test_md_slicing(md, expected_n, d1 + d2)
end

@testset "views" begin
    i, j = (26:75, [string(i) for i in 3:7])
    mdview = @view md[i, j]
    submd = md[i, j]
    viewcp = copy(mdview)

    @test mdview.obs_names == submd.obs_names
    @test mdview.var_names == submd.var_names
    @testset "modality $mod" for (mod, ad) in submd.mod
        @test ad.X == mdview[mod].X == viewcp[mod].X
        @test ad.obs_names == mdview[mod].obs_names == viewcp[mod].obs_names
        @test ad.var_names == mdview[mod].var_names == viewcp[mod].var_names
    end
    test_md_slicing(submd, 50, 10, j)
    test_md_slicing(mdview, 50, 10, j)
end

@testset "readwrite" begin
    tempfile1 = joinpath(tmp, "tmp1.h5mu")
    writeh5mu(tempfile1, md)
    read_md = (@test_logs (:warn, warn_msg) readh5mu(tempfile1, backed=false))
    read_md_backed = (@test_logs (:warn, warn_msg) readh5mu(tempfile1, backed=true))

    @test size(md) == size(read_md) == size(read_md_backed)
    @test md.obs_names == read_md.obs_names == read_md_backed.obs_names
    @test md.var_names == read_md.var_names == read_md_backed.var_names
    @test isequal(md.var, read_md.var)
    @test isequal(md.var, read_md_backed.var)
    @test isequal(md.obs, read_md.obs)
    @test isequal(md.obs, read_md_backed.obs)
    @test md["ad1"].X == read_md["ad1"].X == read_md_backed["ad1"].X
    @test md["ad2"].X == read_md["ad2"].X == read_md_backed["ad2"].X
    @test md["ad1"].var == read_md["ad1"].var == read_md_backed["ad1"].var
    @test md["ad2"].var == read_md["ad2"].var == read_md_backed["ad2"].var
    @test md["ad2"].varm == read_md["ad2"].varm == read_md_backed["ad2"].varm
    @test md["ad2"].layers == read_md["ad2"].layers == read_md_backed["ad2"].layers

    md["ad3"] = ad3
    tempfile2 = joinpath(tmp, "tmp2.h5mu")
    writeh5mu(tempfile2, md)
    @test_logs (:warn, warn_msg2) (:warn, warn_msg) update!(md)
    read_md = (@test_logs (:warn, warn_msg2) (:warn, warn_msg) readh5mu(tempfile2, backed=false))
    read_md_backed = (@test_logs (:warn, warn_msg2) (:warn, warn_msg) readh5mu(tempfile2, backed=true))
    @test size(md) == size(read_md) == size(read_md_backed)
    @test md.obs_names == read_md.obs_names == read_md_backed.obs_names
    @test md.var_names == read_md.var_names == read_md_backed.var_names
    @test isequal(md.var, read_md.var)
    @test isequal(md.var, read_md_backed.var)
    @test isequal(md.obs, read_md.obs)
    @test isequal(md.obs, read_md_backed.obs)
end

@testset "obs_var" begin
    ad1, ad2, _ = make_ads()
    ad2.var_names = ["ad2_$i" for i in 1:size(ad2, 2)]
    md = MuData(mod=Dict("ad1" => ad1, "ad2" => ad2))
    tempfile = joinpath(tmp, "tmp.h5mu")
    @testset "obs_global_columns" begin
        for (m, mod) in md.mod
            mod.obs = DataFrame()
            mod.obs[!, "demo"] = repeat([m], size(mod, 1))
        end
        md.obs[!, "demo"] = repeat(["global"], size(md, 1))
        update!(md)
        @test sort(names(md.obs)) == sort(vcat(["$m:demo" for m in keys(md.mod)], ["demo"]))
        writeh5mu(tempfile, md)
        read_md = readh5mu(tempfile, backed=false)
        @test sort(names(read_md.obs)) == sort(vcat(["$m:demo" for m in keys(md.mod)], ["demo"]))
    end

    @testset "var_global_columns" begin
        for (m, mod) in md.mod
            mod.var = DataFrame()
            mod.var[!, "demo"] = repeat([m], size(mod, 2))
        end
        md.var[!, "global"] = repeat(["global_var"], size(md, 2))
        update!(md)
        @test sort(names(md.var)) == sort(["demo", "global"])
        select!(md.var, Not("global"))
        update!(md)
        @test names(md.var) == ["demo"]
        writeh5mu(tempfile, md)
        read_md = readh5mu(tempfile, backed=false)
        @test names(md.var) == ["demo"]
    end
end


import Base: show, ==, hash

export Qubit, Edge, Marker

import .config.get_qubit_params
import .config.get_marker_params
import .config.get_edge_params
import .config.get_instrument_params

@compat abstract type Channel end
show(io::IO, c::Channel) = print(io, c.label)

# TODO: is there a benefit to having the concrete channel types immutable?
"""
Channel representing single qubit drive.
"""
immutable Qubit <: Channel
	label::String
	awg_channel::String
	gate_channel::String
	shape_params::Dict{Symbol,Any}
	frequency::Float64
end

function Qubit(label)
	channel_params = get_qubit_params()

	if label in keys(channel_params)
		q_params = channel_params[label]["control"]
		phys_chan = get(q_params, "AWG", "")
		gate_chan = get(q_params, "gate", "")

		# pull out shape_params and convert keys to symbols for splatting into shape function
		shape_params = Dict{Symbol, Any}(Symbol(k) => v for (k,v) in q_params["pulse_params"])

		# translate pulse function in shape params from a string to a function handle and snakeify key
		shape_params[:shape_function] = getfield(QGL.PulseShapes, Symbol(pop!(shape_params, :shape_fun)))

		# snakeify some other parameters
		if :dragScaling in keys(shape_params)
			shape_params[:drag_scaling] = pop!(shape_params, :dragScaling)
		end

		Qubit(label, phys_chan, gate_chan, shape_params, q_params["frequency"])
	else
		warn("Unable to find Qubit label $label in channel json file. Creating default Qubit channel")
		Qubit(label,"", "", Dict{String,Any}(), 0.0)
	end

end

"""
Channel representing a digital output marker line.
"""
immutable Marker <: Channel
	label::String
	awg_channel::String
	shape_params::Dict{Any, Any}
end

function Marker(label)
	phys_chan = get_marker_params()[label]
	m_params = get_instrument_params()[split(phys_chan)[1]]["markers"][split(phys_chan)[2]]
	# translate pulse function in shape params from a string to a function handle and snakeify key
	shape_params =  Dict{Symbol, Any}(Symbol(k) => v for (k,v) in m_params)
	shape_params[:shape_function] = getfield(QGL.PulseShapes,  Symbol(pop!(shape_params, :shape_fun)))
	Marker(label, phys_chan, shape_params)
end

# NOTE is this mutable on purpose??
type QuadratureAWGChannel
	awg::String
	delay::Real
	mixer_correction::Matrix{Real}
end

==(a::Channel, b::Channel) = a.label == b.label
hash(c::Channel, h::UInt) = hash(c.label, h)

"""
Channel representing qubit measurement drive
"""
immutable Measurement <: Channel
	label::String
	awg_channel::String
	gate_channel::String
	trigger_channel::String
	shape_params::Dict{Any,Any}
	frequency::Real
end

"""
	measurement_channel(l:String)

Looks us the measurement channel associated with qubit.
"""
function measurement_channel(l::String)
	m_params = get_qubit_params()[l]["measure"]
	phys_chan = get(m_params, "AWG", "")
	gate_chan = get(m_params, "gate", "")
	trig_chan = get(m_params, "trigger", "")
	# pull out shape_params and convert keys to symbols for splatting into shape function
	shape_params = Dict{Symbol, Any}(Symbol(k) => v for (k,v) in m_params["pulse_params"])
	# translate pulse function in shape params from a string to a function handle and snakeify key
	shape_params[:shape_function] = getfield(QGL.PulseShapes,  Symbol(pop!(shape_params, :shape_fun)))
	# inject autodyne frequence into the `shape_params` where it should be
	shape_params[:autodyne_freq] = m_params["autodyne_freq"]
	Measurement(string("M-", l), phys_chan, gate_chan, trig_chan, shape_params, 0)
end

"""
Channel representing an interaction drive.
"""
immutable Edge <: Channel
	label::String
	source::Qubit
	target::Qubit
	awg_channel::String
	gate_channel::String
	shape_params::Dict{Symbol,Any}
	frequency::Float64
end

"""
	Edge(source::Qubit, target::Qubit)

Create the edge representing interaction drive from `source` to `target`.
"""
function Edge(source::Qubit, target::Qubit)
	# look up whether we have an edge connecting source -> target
	channel_params = get_edge_params()

	edges = filter(
		(k,v) -> v["source"] == source.label && v["target"] == target.label,
		channel_params)

	@assert length(edges) == 1 "Found $(length(edges)) matching edges for $source → $target"

	e_params = collect(values(edges))[1]

	phys_chan = get(e_params, "AWG", "")
	gate_chan = get(e_params, "gate", "")

	# pull out shape_params and convert keys to symbols for splatting into shape function
	shape_params = e_params["pulse_params"]
	shape_params = Dict{Symbol, Any}(Symbol(k) => v for (k,v) in shape_params)

	# translate pulse function in shape params from a string to a function handle and snakeify key
	shape_params[:shape_function] = getfield(QGL.PulseShapes, Symbol(pop!(shape_params, :shape_fun)))

	Edge(collect(keys(edges))[1], source, target, phys_chan, gate_chan, shape_params, e_params["frequency"])
end


export AbstractTaxType,
    AbstractTechChange,
    AbstractConstrainedType,
    AbstractLearningType,
    AbstractRoyaltyType,
    NoTaxes,
    WithTaxes,
    GathProcess,
    NoTrend,
    TimeTrend,
    TimeFE,
    Unconstrained,
    Constrained,
    DrillingRevenue,
    Learn,
    NoLearn,
    WithRoyalty,
    NoRoyalty,
    ConstrainedProblem,
    UnconstrainedProblem,
    NoLearningProblem,
    STARTING_α_ψ,
    STARTING_log_ogip,
    STARTING_α_t,
    STARTING_α_D,
    TimeTrendDgt0

# ----------------------------------------------------------------
# Drilling revenue variations
# ----------------------------------------------------------------

abstract type AbstractLearningType    <: AbstractModelVariations end
abstract type AbstractConstrainedType <: AbstractModelVariations end
abstract type AbstractTechChange      <: AbstractModelVariations end
abstract type AbstractTaxType         <: AbstractModelVariations end
abstract type AbstractRoyaltyType     <: AbstractModelVariations end

# drilling revenue
struct DrillingRevenue{
    Cn <: AbstractConstrainedType, Tech <: AbstractTechChange,
    Tax <: AbstractTaxType, Lrn <: AbstractLearningType,
     Roy <: AbstractRoyaltyType
} <: AbstractDrillingRevenue
    constr::Cn
    tech::Tech
    tax::Tax
    learn::Lrn
    royalty::Roy
end

DrillingRevenue(Cn, Tech, Tax) = DrillingRevenue(Cn, Tech, Tax, Learn(), WithRoyalty())

@inline constr( x::DrillingRevenue) = x.constr
@inline tech(   x::DrillingRevenue) = x.tech
@inline tax(    x::DrillingRevenue) = x.tax
@inline learn(  x::DrillingRevenue) = x.learn
@inline royalty(x::DrillingRevenue) = x.royalty

# Constrained parameters
# -----------------------

# Whether we've constrained coefs
struct Unconstrained <: AbstractConstrainedType end
struct Constrained   <: AbstractConstrainedType
    log_ogip::Float64
    α_ψ::Float64
    α_t::Float64
    α_D::Float64
end
Constrained(; log_ogip=STARTING_log_ogip, α_ψ = STARTING_α_ψ, α_t = STARTING_α_t, α_D = STARTING_α_D, args...) = Constrained(log_ogip, α_ψ, α_t, α_D)

@inline theta_g(x::Constrained) = x.log_ogip
@inline theta_ψ(x::Constrained) = x.α_ψ
@inline theta_t(x::Constrained) = x.α_t
@inline theta_D(x::Constrained) = x.α_D

const DrillingRevenueConstrained = DrillingRevenue{Constrained}
const DrillingRevenueUnconstrained = DrillingRevenue{Unconstrained}

# Technology
# -----------------

abstract type AbstractTimeTrend <: AbstractTechChange end

const TIME_TREND_BASE = 2008.5

struct NoTrend <: AbstractTechChange  end
struct TimeTrend <: AbstractTimeTrend
    baseyear::Float64
end
struct TimeTrendDgt0 <: AbstractTimeTrend
    baseyear::Float64
end
(::Type{T})() where {T<:AbstractTimeTrend} = T(TIME_TREND_BASE)
struct TimeFE <: AbstractTechChange
    yearrange::UnitRange{Float64}
end

# Special TimeFE stuff
yearrange(x::TimeFE) = x.yearrange
yearrange(x::AbstractTechChange) = UnitRange(2008.5, 2016.5)
start(x::TimeFE) = first(yearrange(x))
stop(x::TimeFE) = last(yearrange(x))
_nparm(x::TimeFE) = length(yearrange(x))-1
TimeFE(a,b) = TimeFE(UnitRange(a,b))
TimeFE() = TimeFE(2008.5, 2016.5)
function coefnames(x::TimeFE)
    yrng = UnitRange(start(x)+1, stop(x))
    cfnms = ["\\alpha_{" * string(yr) *"}" for yr in yrng]
    return cfnms
end

baseyear(x::AbstractTimeTrend) = x.baseyear
baseyear(x::NoTrend) = TIME_TREND_BASE
baseyear(x::TimeFE) = start(x)
baseyear(x::DrillingRevenue) = baseyear(tech(x))
baseyear(x::DrillReward) = baseyear(revenue(x))

const DrillingRevenueTimeTrend         = DrillingRevenue{Cn,TimeTrend} where {Cn}
const DrillingRevenueTimeTrendDgt0     = DrillingRevenue{Cn,TimeTrendDgt0} where {Cn}
const DrillingRevenueAbstractTimeTrend = DrillingRevenue{Cn,T} where {Cn,T<:AbstractTimeTrend}
const DrillingRevenueNoTrend           = DrillingRevenue{Cn,NoTrend} where {Cn}
const DrillingRevenueTimeFE            = DrillingRevenue{Cn,TimeFE} where {Cn}

@inline centered_time(  x::DrillingRevenue, z::Tuple) = year(z) - baseyear(tech(x))
@inline trend_component(x::DrillingRevenueTimeTrend,     θ, obs) = theta_t(x,θ) * centered_time(x, zchars(obs))
@inline trend_component(x::DrillingRevenueNoTrend,       θ, obs) = 0
@inline trend_component(x::DrillingRevenueTimeTrendDgt0, θ, obs) = theta_t(x,θ)*centered_time(x,zchars(obs)) + _Dgt0(obs)*theta_D(x,θ)

@inline function idx_year(x::DrillingRevenueTimeFE, z)
    timefe = tech(x)
    if year(z) < start(timefe)
        return 0
    elseif year(z) > stop(timefe)
        return _nparm(timefe)
    else
        t = centered_time(x,z)
        return Int(floor(t))
    end
end

function trend_component(x::DrillingRevenueTimeFE, θ, obs)
    i = idx_year(x,zchars(obs))
    if i > 0
        return theta_t(x,θ,i)
    else
        return zero(eltype(θ))
    end
end

@inline function gradt!(grad, x::DrillingRevenueNoTrend, rev, θ, obs)
    nothing
end
@inline function gradt!(grad, x::DrillingRevenueTimeTrend, rev, θ, obs)
    grad[idx_t(x)] = rev*centered_time(x, zchars(obs))
end
@inline function gradt!(grad, x::DrillingRevenueTimeTrendDgt0, rev, θ, obs)
    grad[idx_t(x)] = rev*centered_time(x, zchars(obs))
    grad[idx_D(x)] = rev*_Dgt0(obs)
end
@inline function gradt!(grad, x::DrillingRevenueTimeFE, rev, θ, obs)
    grad[idx_t(x)] .= 0
    i = idx_year(x,zchars(obs))
    if i > 0
        grad[idx_t(x,i)] = rev
    end
end

# chebshev polynomials
# See http://www.aip.de/groups/soe/local/numres/bookcpdf/c5-8.pdf
@inline checkinterval(x::Real,min::Real,max::Real) =  min <= x <= max || throw(DomainError("x = $x must be in [$min,$max]"))
@inline checkinterval(x::Real) = checkinterval(x,-1,1)

@inline cheb0(z::Real) = one(z)
@inline cheb1(z::Real) = clamp(z,-1,1)
@inline cheb2(z::Real) = (x = cheb1(z); return 2*x^2 - 1)
@inline cheb3(z::Real) = (x = cheb1(z); return 4*x^3 - 3*x)
@inline cheb4(z::Real) = (x = cheb1(z); return 8*(x^4 - x^2) + 1)

# learning
# -----------------------

struct Learn       <: AbstractLearningType end
struct NoLearn     <: AbstractLearningType end
struct PerfectInfo <: AbstractLearningType end
struct MaxLearning <: AbstractLearningType end

const DrillingRevenueLearn       = DrillingRevenue{Cn,Tech,Tax,Learn}       where {Cn,Tech,Tax}
const DrillingRevenueNoLearn     = DrillingRevenue{Cn,Tech,Tax,NoLearn}     where {Cn,Tech,Tax}
const DrillingRevenuePerfectInfo = DrillingRevenue{Cn,Tech,Tax,PerfectInfo} where {Cn,Tech,Tax}
const DrillingRevenueMaxLearning = DrillingRevenue{Cn,Tech,Tax,MaxLearning} where {Cn,Tech,Tax}

@inline _ρ(x::AbstractLearningType, σ::Number) = _ρ(σ)
@inline _ρ(x::PerfectInfo         , σ::Number) = 1
@inline _ρ(x::MaxLearning         , σ::Number) = 0
@inline _ρ(x::DrillingRevenue     , σ::Number) = _ρ(learn(x), σ)

# @deprecate _ρ(σ::Number, x::AbstractLearningType) _ρ(x,σ)
# @deprecate _ρ(σ::Number, x::PerfectInfo) _ρ(x,σ)
# @deprecate _ρ(σ::Number, x::MaxLearning) _ρ(x,σ)
# @deprecate _ρ(σ::Number, x::DrillingRevenue) _ρ(x,σ)

# parameter access
# -----------------------

@inline _nparm(x::DrillingRevenue{Constrained}) = 2
@inline _nparm(x::DrillingRevenue{Unconstrained, NoTrend}) = 4
@inline _nparm(x::DrillingRevenue{Unconstrained, TimeTrend}) = 5
@inline _nparm(x::DrillingRevenue{Unconstrained, TimeTrendDgt0}) = 6
@inline _nparm(x::DrillingRevenue{Unconstrained, TimeFE}) = 4 + _nparm(tech(x))

@inline idx_0(x::DrillingRevenue) = 1
@inline idx_g(x::DrillingRevenue) = 2
@inline idx_ψ(x::DrillingRevenue) = 3
@inline idx_t(x::DrillingRevenueAbstractTimeTrend) = 4
@inline idx_t(x::DrillingRevenueTimeFE) = 3 .+ OneTo(_nparm(tech(x)))
function idx_t(x::DrillingRevenueTimeFE, i)
    1 <= i <= _nparm(tech(x)) || throw(DomainError(i))
    return 3 + i
end
@inline idx_D(x::DrillingRevenueTimeTrendDgt0) = 5
@inline idx_ρ(x::DrillingRevenue) = _nparm(x)

@inline theta_0(x::DrillingRevenue, θ) = θ[idx_0(x)]
@inline theta_g(x::DrillingRevenue, θ) = θ[idx_g(x)]
@inline theta_ψ(x::DrillingRevenue, θ) = θ[idx_ψ(x)]
@inline theta_t(x::DrillingRevenue{Unconstrained}, θ) = θ[idx_t(x)]
@inline theta_t(x::DrillingRevenue{Unconstrained, TimeFE}, θ) = uview(θ, idx_t(x))
@inline theta_t(x::DrillingRevenue{Unconstrained, TimeFE}, θ, i) = θ[idx_t(x,i)]
@inline theta_D(x::DrillingRevenue{Unconstrained, TimeTrendDgt0},θ) = θ[idx_D(x)]
@inline theta_D(x, θ) = zero(Float64)
@inline theta_ρ(x, θ) = θ[idx_ρ(x)]
function theta_t(x::DrillingRevenue{Unconstrained, NoTrend}, θ)
    # @warn "No t in $x. setting to $STARTING_α_t"
    return STARTING_α_t
end
@inline theta_g(x::DrillingRevenue{Constrained}, θ) = theta_g(constr(x))
@inline theta_ψ(x::DrillingRevenue{Constrained}, θ) = theta_ψ(constr(x))
@inline theta_t(x::DrillingRevenue{Constrained}, θ) = theta_t(constr(x))

function ConstrainedCoefs(x::DrillingRevenueUnconstrained, theta)
    @assert length(theta) == _nparm(x)
    g   = theta_g(x, theta)
    psi = theta_ψ(x, theta)
    t   = theta_t(x, theta)
    D   = theta_D(x, theta)
    out = (log_ogip=g, α_ψ=psi, α_t=t, α_D=D)
    return out
end

ConstrainedCoefs(x::DrillReward, θ) = ConstrainedCoefs(revenue(x), vw_revenue(x, θ))
ConstrainedIdx(x::AbstractDrillingRevenue) = idx_g(x), idx_ψ(x), idx_t(x)
ConstrainedIdx(x::DrillingRevenueNoTrend) = idx_g(x), idx_ψ(x)
ConstrainedIdx(x::DrillingRevenueTimeTrendDgt0) = idx_g(x), idx_ψ(x), idx_t(x), idx_D(x)

function UnconstrainedFmConstrainedIdx(x::DrillingRevenueConstrained)
    xu = UnconstrainedProblem(x)
    idxu = [idx_0(xu), idx_ρ(xu)]
end

# @deprecate α_0(     x::DrillingRevenue, θ) theta_0(x, θ)
# @deprecate _σ(      x::DrillingRevenue, θ) theta_ρ(x, θ)
# @deprecate log_ogip(x::DrillingRevenue, θ) theta_g(x, θ)
# @deprecate α_ψ(     x::DrillingRevenue, θ) theta_ψ(x, θ)
# @deprecate α_t(     x::DrillingRevenue, θ) theta_t(x, θ)

# taxes
# -----------------------

# From Gulen et al (2015) "Production scenarios for the Haynesville..."
const GATH_COMP_TRTMT_PER_MCF   = 0.42 + 0.07
const MARGINAL_TAX_RATE = 0.402

# other calculations
const REAL_DISCOUNT_AND_DECLINE = 0x1.8e06b611ed4d8p-1  #  0.777394952476835= sum( β^((t+5-1)/12) q(t)/Q(240) for t = 1:240)

const STARTING_α_ψ      = 0x1.7587cc6793516p-2 # 0.365
const STARTING_log_ogip = 0x1.401755c339009p-1 # 0.625
const STARTING_α_t      = 0.01948
const STARTING_α_D      = 0.0

struct NoTaxes <: AbstractTaxType end

struct WithTaxes <: AbstractTaxType
    one_minus_mgl_tax_rate::Float64
    cost_per_mcf::Float64
end

struct GathProcess <: AbstractTaxType
    cost_per_mcf::Float64
end

const DrillingRevenueNoTaxes     = DrillingRevenue{Cn,Tech,NoTaxes}     where {Cn,Tech}
const DrillingRevenueWithTaxes   = DrillingRevenue{Cn,Tech,WithTaxes}   where {Cn,Tech}
const DrillingRevenueGathProcess = DrillingRevenue{Cn,Tech,GathProcess} where {Cn,Tech}

WithTaxes(;mgl_tax_rate = 1-MARGINAL_TAX_RATE, cost_per_mcf = GATH_COMP_TRTMT_PER_MCF * REAL_DISCOUNT_AND_DECLINE) = WithTaxes(mgl_tax_rate, cost_per_mcf)
GathProcess(;cost_per_mcf = GATH_COMP_TRTMT_PER_MCF * REAL_DISCOUNT_AND_DECLINE) = GathProcess(cost_per_mcf)

@inline one_minus_mgl_tax_rate(x::AbstractTaxType) = 1
@inline one_minus_mgl_tax_rate(x::WithTaxes) = x.one_minus_mgl_tax_rate
@inline one_minus_mgl_tax_rate(x) = one_minus_mgl_tax_rate(tax(x))

@inline cost_per_mcf(x::NoTaxes) = 0
@inline cost_per_mcf(x::AbstractTaxType) = x.cost_per_mcf
@inline cost_per_mcf(x) = cost_per_mcf(tax(x))

# royalty
# -----------------------

struct WithRoyalty <: AbstractRoyaltyType end
struct NoRoyalty   <: AbstractRoyaltyType end

@inline (::WithRoyalty)(roy) = roy
@inline (::NoRoyalty)(roy) = 0
@inline oneminusroyalty(x::DrillingRevenue, roy) = 1-royalty(x)(roy)
@inline d_tax_royalty(x::DrillingRevenue, d::Integer, roy::Real) = d * oneminusroyalty(x, roy) * one_minus_mgl_tax_rate(x)

# base functions
# -------------------------------------------

@inline function Eexpψ(x::DrillingRevenueLearn, d, obs, θ, sim)
    ψ = _ψ(obs, sim)
    θ4 = theta_ψ(x,θ)

    if _Dgt0(obs)
        return θ4*ψ
    else
        σ = theta_ρ(x,θ)
        ρ = _ρ(learn(x), σ)
        return θ4*(ψ*ρ + θ4*0.5*(1-ρ^2))
    end
end

@inline function dEexpψdα_ψ(x::DrillingRevenueLearn, d, obs, θ, sim)
    ψ = _ψ(obs, sim)
    if _Dgt0(obs)
        return ψ
    else
        σ = theta_ρ(x,θ)
        ρ = _ρ(x, σ)
        return ψ*ρ + theta_ψ(x,θ) * (1-ρ^2)
    end
end

@inline function Eexpψ(x::DrillingRevenueNoLearn, d, obs, θ, sim)
    ψ = _ψ(obs, sim)
    θ4= theta_ψ(x,θ)
    σ = theta_ρ(x,θ)
    ρ = _ρ(σ)
    return θ4*(ψ*ρ + θ4*0.5*(1-ρ^2))
end

@inline function Eexpψ(x::DrillingRevenueMaxLearning, d, obs, θ, sim)
    θ4= theta_ψ(x,θ)
    if _Dgt0(obs)
        return θ4 * _ψ(obs, sim)
    else
        return 0.5 * θ4^2
    end
end

@inline Eexpψ(x::DrillingRevenuePerfectInfo, d, obs, θ, sim) = theta_ψ(x,θ) * _ψ(obs, sim)

@inline function dEexpψdσ(x::DrillingRevenueLearn, d, obs, θ, sim)
    T = eltype(θ)
    ψ = _ψ(obs, sim)
    if !_Dgt0(obs) && d > 0
        θ4= theta_ψ(x,θ)
        σ = theta_ρ(x,θ)
        out = θ4 *(ψ - θ4*_ρ(σ)) * _dρdθρ(σ)
    else
        out = zero(T)
    end
    return out::T
end

# quantity (for post-estimation SIMULATIONS)
# -----------
function eur_kernel(x::DrillingRevenue, d, obs::ObservationDrill, θ, sim)
    @assert length(θ) == _nparm(x)
    z = zchars(obs)
    logQ = theta_g(x,θ)*geology(obs) + theta_ψ(x,θ)*_ψ2(sim) + trend_component(x, θ, obs)
    return exp(logQ)
end

# ----------------------------------------------------------------
# flow revenue
# ----------------------------------------------------------------

function flow(x::DrillingRevenueNoTaxes, d::Integer, obs::ObservationDrill, θ::AbstractVector, sim)
    z = zchars(obs)
    u = d_tax_royalty(x, d, royalty(obs)) *
    exp(
        theta_0(x,θ) + logprice(obs) + theta_g(x,θ)*geology(obs) +
        Eexpψ(x, d, obs, θ, sim) +
        trend_component(x, θ, obs)
    )
    return u
end

function flow(x::DrillingRevenue, d::Integer, obs::ObservationDrill, θ::AbstractVector, sim)
    z = zchars(obs)
    u = d_tax_royalty(x, d, royalty(obs)) *
    exp(
        theta_0(x,θ) + theta_g(x,θ)*geology(obs) +
        Eexpψ(x, d, obs, θ, sim) +
        trend_component(x, θ, obs)
    ) * (
        price(obs) - cost_per_mcf(x)
    )
    return u
end

# Gradient
# ------------------------------

@inline function flowzero!(grad, θ, dograd)
    dograd && fill!(grad, 0)
    return azero(θ)
end

function flow!(grad, x::DrillingRevenue{Unconstrained}, d, obs, θ, sim, dograd::Bool)
    d == 0 && return flowzero!(grad,θ,dograd)
    z = zchars(obs)
    rev = flow(x, d, obs, θ, sim)
    if dograd
        grad[idx_0(x)] = rev
        grad[idx_g(x)] = rev*geology(obs)
        grad[idx_ψ(x)] = rev*dEexpψdα_ψ(x, d, obs, θ, sim)
        gradt!(grad, x, rev, θ, obs)
        grad[idx_ρ(x)] = rev*dEexpψdσ(x, d, obs, θ, sim)
    end
    return rev
end

# Constrained gradient
# ------------------------------

function flow!(grad, x::DrillingRevenue{Constrained}, d, obs, θ, sim, dograd::Bool)
    d == 0 && return flowzero!(grad,θ,dograd)
    rev = flow(x, d, obs, θ, sim)
    if dograd
        grad[idx_0(x)] = rev
        grad[idx_ρ(x)] = rev*dEexpψdσ(x, d, obs, θ, sim)
    end
    return rev
end

coefnames(x::DrillingRevenue{Unconstrained, NoTrend}) =
    ["\\alpha_0", "\\alpha_g", "\\alpha_\\psi", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Unconstrained, TimeTrend}) =
    ["\\alpha_0", "\\alpha_g", "\\alpha_\\psi", "\\alpha_t", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Unconstrained, TimeTrendDgt0}) =
    ["\\alpha_0", "\\alpha_g", "\\alpha_\\psi", "\\alpha_t", "\\alpha_D", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Constrained}) =
    ["\\alpha_0", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Unconstrained, TimeFE}) = vcat(
    "\\alpha_0", "\\alpha_g", "\\alpha_\\psi", coefnames(tech(x)), "\\theta_\\rho")

# ----------------------------------------------------------------
# dψ is the same across many functions
# ----------------------------------------------------------------

function flowdψ(rev::Real, x::DrillingRevenue, d, obs, θ, sim)
    dψ = rev * theta_ψ(x,θ)
    if _Dgt0(obs)
        return dψ
    else
        return dψ * _ρ(x,theta_ρ(x,θ))
    end
end

function flowdψ(x::DrillingRevenue, d, obs, θ, sim)
    d == 0 && return azero(θ)
    rev = flow(x, d, obs, θ, sim)
    return flowdψ(rev, x, d, obs, θ, sim)
end

# Conversion between model types
# ----------------------------------------------------------------

# Constrained / Unconstrained
ConstrainedProblem(  x::AbstractPayoffComponent; kwargs...) = x
UnconstrainedProblem(x::AbstractPayoffComponent; kwargs...) = x

UnconstrainedProblem(x::DrillingRevenue; kwargs...) = DrillingRevenue(Unconstrained(;kwargs...), tech(x), tax(x), learn(x), royalty(x))
ConstrainedProblem(  x::DrillingRevenue; kwargs...) = DrillingRevenue(Constrained(;kwargs...),   tech(x), tax(x), learn(x), royalty(x))

ConstrainedProblem(  x::DrillReward; kwargs...) = DrillReward(ConstrainedProblem(revenue(x); kwargs...), ConstrainedProblem(drill(x)), ConstrainedProblem(extend(x)))
UnconstrainedProblem(x::DrillReward; kwargs...) = DrillReward(UnconstrainedProblem(revenue(x); kwargs...), UnconstrainedProblem(drill(x)), UnconstrainedProblem(extend(x)))
ConstrainedProblem(x::DrillReward, theta::AbstractVector) = ConstrainedProblem(x; ConstrainedCoefs(x,theta)...)

# Learning models
NoLearningProblem(x::AbstractPayoffComponent, args...) = x
LearningProblem(  x::AbstractPayoffComponent, args...) = x

NoLearningProblem(x::DrillingRevenue, args...) = DrillingRevenue(constr(x), tech(x), tax(x), NoLearn(),     royalty(x))
LearningProblem(  x::DrillingRevenue, args...) = DrillingRevenue(constr(x), tech(x), tax(x), Learn(),       royalty(x))
PerfectInfo(      x::DrillingRevenue, args...) = DrillingRevenue(constr(x), tech(x), tax(x), PerfectInfo(), royalty(x))

NoLearningProblem(x::DrillReward, args...) = DrillReward(NoLearningProblem(revenue(x), args...), drill(x), extend(x))
LearningProblem(  x::DrillReward, args...) = DrillReward(  LearningProblem(revenue(x), args...), drill(x), extend(x))

# Royalty
NoRoyaltyProblem(  x::AbstractPayoffComponent, args...) = x
WithRoyaltyProblem(x::AbstractPayoffComponent, args...) = x

NoRoyaltyProblem(  x::DrillingRevenue, args...)      = DrillingRevenue(constr(x), tech(x), tax(x), learn(x), NoRoyalty())
WithRoyaltyProblem(x::DrillingRevenue, args...)      = DrillingRevenue(constr(x), tech(x), tax(x), learn(x), WithRoyalty())

NoRoyaltyProblem(  x::DrillReward, args...) = DrillReward(  NoRoyaltyProblem(revenue(x), args...), drill(x), extend(x))
WithRoyaltyProblem(x::DrillReward, args...) = DrillReward(WithRoyaltyProblem(revenue(x), args...), drill(x), extend(x))


export @fermiview

macro fermiview()
    quote
        try
            let mol = Fermi.Options.get("molstring")
                E = Ensamble()
                Alchemy.xyz_to!(E, mol)
                Alchemy.find_bonds!(E)
                Alchemy.run(E)
            end
        catch UndefVarError
            @error "Fermi.jl could not be found, try running `using Fermi` first."
        end
    end |> esc
end


using AlphaGo
using AlphaGo: go
using Flux
using BSON: @load

# Note:
# do Pkg.checkout("WebIO", "s/asset-registry")
# and Pkg.checkout("Mux", "s/asset-registry")
include("../src/interface.jl")
include("./helpers.jl")

set_all_params(9)

assets = "$(@__DIR__)/assets"
files=[
    "$(@__DIR__)/../lib/interface.js"
    "$(@__DIR__)/../lib/interface.css"
    "$(assets)/css/go.css"
    "$(assets)/css/wgo.player.css"
    "$(assets)/js/wgo.min.js"
    "$(assets)/js/go.js"]

gw = GameWindow(files; config=Dict("boardSize"=>9))
setup(gw) # default setup

# agz_nn = loadNeuralNet()
agz_nn = NeuralNet(;tower_height=1)
agz = MCTSPlayer(agz_nn, num_readouts = 64, two_player_mode = true)

playing = false
isPlaying() = playing
setPlaying(p=true) = (playing = p)

# game play
oninput(gw) do x
    println("js sent $x")
    is_done(agz) && ( return gw.msg[] = "Game over! " * agz.result_string )
    isPlaying() && return gw.msg[] = "computer's turn!"
    isIllegalMove(x) && return gw.msg[] = "Illegal Move"
    setPlaying()

    move = playMove(x) # user action
    println("user action...")
    updateBoard(x)   # this happens kinda late ( along with the next update )
    move = playMove()  # computer action
    println("computer action...")
    updateBoard(stringToAction(go.to_kgs(move), -1))

    setPlaying(false)
    is_done(agz) && ( return gw.msg[] = "Game over! " * agz.result_string )
end

function startGame()
    initialize_game!(agz)
    gw(dom"div#demo_wrapper"(
        dom"div#controls"(
            dom"h1"("Go game"),
            dom"button.pass"("Pass")
        ),
        dom"div#playground"()))
end

startGame()