###############################################################################
#
# Calculate gas exchange from flow rate only for required parameters
#
###############################################################################
"""
    leaf_gas_exchange_nonopt!(
                node::SPACSimple{FT},
                photo_set::AbstractPhotoModelParaSet{FT},
                flow::FT,
                par::FT,
                rad::FT,
                la::FT,
                container::SPACContainer1L{FT}
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange_nonopt!(
                node::SPACSimple{FT},
                photo_set::AbstractPhotoModelParaSet{FT},
                flow::FT
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange_nonopt!(
                node::SPACSimple{FT},
                photo_set::AbstractPhotoModelParaSet{FT},
                f_sl::FT,
                f_sh::FT
    ) where {FT<:AbstractFloat}

Simulate leaf level gas exchange and fill it into the `container` for 1-layer
    or 2-layer canopy, given
- `node` [`SPACSimple`] type struct
- `photo_set` [`AbstractPhotoModelParaSet`] type struct
- `flow` Flow rate per basal area into the leaves (e.g., for sunlit leaves)
- `f_sl` Flow rate per basal area into the sunlit leaves
- `f_sh` Flow rate per basal area into the shaded leaves
- `par` Leaf-level photosynthetic active radiation
- `rad` Leaf-level absorbed radiative energy
- `la` Leaf area of the leaves (total or each layer)
- `container` [`SPACContainer1L`] type container
"""
function leaf_gas_exchange_nonopt!(
            node::SPACSimple{FT},
            photo_set::AbstractPhotoModelParaSet{FT},
            flow::FT,
            par::FT,
            rad::FT,
            la::FT,
            container::SPACContainer1L{FT}
) where {FT<:AbstractFloat}
    # if flow < 0
    if flow < 0
        container.an = FT(-Inf);

    # if flow >= 0
    else
        # 0. unpack required variables
        @unpack envir, g_max, width = node;
        @unpack p_atm, p_H₂O, t_air, wind = envir;

        # 1. calculate leaf temperature from the flow rate
        t_leaf = max(200, leaf_temperature(node, rad, flow));

        # 2. update leaf photosynthetic variables and leaf-to-air VPD
        node.ps.APAR = par;
        leaf_temperature_dependence!(photo_set, node.ps, envir, t_leaf);
        d_leaf = node.ps.p_sat - p_H₂O;

        # 3. update f_vis and f_st in leaf and calculate water potentials
        # TODO do not account for temperature effects for flow now
        # More reasonable functions need to be added
        # vc_temperature_effects!(node.hs.leaf, t_leaf);

        # if flow == 0
        if flow == 0
            g_lw = FT(0);
            g_lc = FT(1e-6);
            leaf_photosynthesis!(photo_set, node.ps, envir, GCO₂Mode(), g_lc);
            container.an = node.ps.An;

        # if flow > 0 and reasonable
        elseif (d_leaf > 0) && (t_leaf > K_0(FT))
            t_cor = relative_diffusive_coefficient( (t_leaf+t_air)/2 );
            g_lw  = flow / la / d_leaf * p_atm;
            g_bw  = boundary_layer_conductance(wind, width);
            g_bc  = g_bw / FT(1.35);
            g_sw  = 1 / (1/g_lw - 1/(g_bw*t_cor));
            g_sc  = g_sw / FT(1.6);
            g_lc  = max(FT(1e-6), 1 / (1/g_bc + 1/g_sc));
            g_lim = 1 / (1/g_bw + 1/g_max);
            if g_lw < g_lim * t_cor
                leaf_photosynthesis!(photo_set, node.ps, envir, GCO₂Mode(),
                                     g_lc);
                container.an = node.ps.An;
            else
                container.an = FT(-Inf);
            end

        # if flow > 0 and unrealistic
        else
            container.an = FT(-Inf);
        end
    end

    return nothing
end

function leaf_gas_exchange_nonopt!(
            node::SPACSimple{FT},
            photo_set::AbstractPhotoModelParaSet{FT},
            flow::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack frac_sh, frac_sl, par_sh, par_sl, rad_sh,
            rad_sl = node.container2L;

    # calculate mean par and rad per leaf area, then gas exchange rate
    par_mean = par_sl * frac_sl + par_sh * frac_sh;
    rad_mean = rad_sl * frac_sl + rad_sh * frac_sh;
    leaf_gas_exchange_nonopt!(node, photo_set, flow, par_mean, rad_mean,
                              node.laba, node.container1L);

    node.containerOP = (node.ec - flow) * (node.container1L).an;

    return nothing
end

function leaf_gas_exchange_nonopt!(
            node::SPACSimple{FT},
            photo_set::AbstractPhotoModelParaSet{FT},
            f_sl::FT,
            f_sh::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack frac_sh, frac_sl, la_sh, la_sl, par_sh, par_sl, rad_sh,
            rad_sl = node.container2L;

    # calculate gas exchangr for sunlit and shaded layers
    leaf_gas_exchange_nonopt!(node, photo_set, f_sl, par_sl, rad_sl, la_sl,
                              node.container2L.cont_sl);
    leaf_gas_exchange_nonopt!(node, photo_set, f_sh, par_sh, rad_sh, la_sh,
                              node.container2L.cont_sh);

    a_sum = frac_sl * node.container2L.cont_sl.an +
            frac_sh * node.container2L.cont_sh.an;
    e_sum = f_sl + f_sh;
    node.containerOP = (node.ec - e_sum) * a_sum;

    return nothing
end

###############################################################################
#
# Calculate gas exchange from flow rate considering temperature dynamics
#
###############################################################################
"""
    leaf_gas_exchange!(
                node::SPACSimple{FT},
                photo_set::AbstractPhotoModelParaSet{FT},
                flow::FT,
                par::FT,
                rad::FT,
                la::FT,
                container::SPACContainer1L{FT}
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange!(
                node::SPACSimple{FT},
                photo_set::AbstractPhotoModelParaSet{FT},
                flow::FT
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange!(
                node::SPACSimple{FT},
                photo_set::AbstractPhotoModelParaSet{FT},
                f_sl::FT,
                f_sh::FT
    ) where {FT<:AbstractFloat}

Simulate leaf level gas exchange and fill it into the `container` for 1-layer
    or 2-layer canopy, given
- `node` [`SPACSimple`] type struct
- `photo_set` [`AbstractPhotoModelParaSet`] type struct
- `flow` Flow rate per basal area into the leaves (e.g., for sunlit leaves)
- `f_sl` Flow rate per basal area into the sunlit leaves
- `f_sh` Flow rate per basal area into the shaded leaves
- `par` Leaf-level photosynthetic active radiation
- `rad` Leaf-level absorbed radiative energy
- `la` Leaf area of the leaves (total or each layer)
- `container` [`SPACContainer1L`] type container
"""
function leaf_gas_exchange!(
            node::SPACSimple{FT},
            photo_set::AbstractPhotoModelParaSet{FT},
            flow::FT,
            par::FT,
            rad::FT,
            la::FT,
            container::SPACContainer1L{FT}
) where {FT<:AbstractFloat}
    # 0. unpack required variables
    @unpack envir = node;
    @unpack p_atm, p_H₂O = envir;

    # 1. calculate leaf temperature from the flow rate
    t_leaf = max(200, leaf_temperature(node, rad, flow));

    # 2. update leaf photosynthetic variables and leaf-to-air VPD
    node.ps.APAR = par;
    leaf_temperature_dependence!(photo_set, node.ps, envir, t_leaf);
    d_leaf = node.ps.p_sat - p_H₂O;

    # 3. update f_vis and f_st in leaf and calculate water potentials
    # TODO do not account for temperature effects for flow now
    # More reasonable functions need to be added
    # vc_temperature_effects!(node.hs.leaf, t_leaf);
    p_leaf = end_pressure(node.hs, flow);

    # 4. calculate photosynthesis
    g_lw = flow / la / d_leaf * p_atm;
    g_lc = max(FT(1e-6), g_lw / FT(1.6));
    leaf_photosynthesis!(photo_set, node.ps, envir, GCO₂Mode(), g_lc);
    container.ag = node.ps.Ag;
    container.an = node.ps.An;
    container.c  = node.ps.p_i;
    container.e  = flow;
    container.gh = g_lw;
    container.p  = p_leaf;
    container.t  = t_leaf;

    return nothing
end

function leaf_gas_exchange!(
            node::SPACSimple{FT},
            photo_set::AbstractPhotoModelParaSet{FT},
            flow::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack container2L, laba = node;
    @unpack frac_sh, frac_sl, par_sh, par_sl, rad_sh, rad_sl = container2L;

    # calculate mean par and rad per leaf area, then gas exchange rate
    par_mean = par_sl * frac_sl + par_sh * frac_sh;
    rad_mean = rad_sl * frac_sl + rad_sh * frac_sh;
    leaf_gas_exchange!(node, photo_set, flow, par_mean, rad_mean, laba,
                       node.container1L);

    return nothing
end

function leaf_gas_exchange!(
            node::SPACSimple{FT},
            photo_set::AbstractPhotoModelParaSet{FT},
            f_sl::FT,
            f_sh::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack frac_sh, frac_sl, la_sh, la_sl, par_sh, par_sl, rad_sh,
            rad_sl = node.container2L;

    # calculate gas exchangr for sunlit and shaded layers
    leaf_gas_exchange!(node, photo_set, f_sl, par_sl, rad_sl, la_sl,
                       node.container2L.cont_sl);
    leaf_gas_exchange!(node, photo_set, f_sh, par_sh, rad_sh, la_sh,
                       node.container2L.cont_sh);

    return nothing
end


# __BEGIN_LICENSE__
#
# WaveOptics.jl
#
# Copyright (c) 2015, Stanford University
#
# All rights reserved.
#
# Redistribution and use in source and binary forms for academic and other
# non-commercial purposes with or without modification, are permitted provided
# that the following conditions are met:
#
# * Redistributions of source code, including modified source code, must retain
#   the above copyright notice, this list of conditions and the following
#   disclaimer.
#
# * Redistributions in binary form or a modified form of the source code must
#   reproduce the above copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other materials provided with
#   the distribution.
#
# * Neither the name of The Leland Stanford Junior University, any of its
#   trademarks, the names of its employees, nor contributors to the source code
#   may be used to endorse or promote products derived from this software without
#   specific prior written permission.
#
# * Where a modified version of the source code is redistributed publicly in
#   source or binary forms, the modified source code must be published in a freely
#   accessible manner, or otherwise redistributed at no charge to anyone
#   requesting a copy of the modified source code, subject to the same terms as
#   this agreement.
#
# THIS SOFTWARE IS PROVIDED BY THE TRUSTEES OF THE LELAND STANFORD JUNIOR
# UNIVERSITY "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE LELAND STANFORD JUNIOR
# UNIVERSITY OR ITS TRUSTEES BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
# OF SUCH DAMAGE.
#
# __END_LICENSE__

abstract type OpticalRecipe end

# Aperture code types (instances are defined in the 'aperture' sub-module)
abstract type CodedAperture end

const MaskParameters = Dict{String,Float64}
const MaskOptions = Dict{String,Any}

# -------------------------------------------------------------------------------------------
#                                 Null Coded Aperture
# -------------------------------------------------------------------------------------------

mutable struct NullCodedAperture <: CodedAperture
    mask_name::String
    mask_parameters::MaskParameters
    mask_options::MaskOptions
end

function NullCodedAperture()
    return NullCodedAperture("null_phase", MaskParameters(), MaskOptions())
end

function NullCodedAperture(mask_parameters::MaskParameters, mask_options)
    return NullCodedAperture("null_phase", MaskParameters(), MaskOptions())
end

function mask(
    ap::NullCodedAperture,
    coordinates::OpticalCoordinates,
    parameters = Dict{String,Float64}(),
)
    return ones(Complex64, size(coordinates.rho))
end

function jacobian(
    ap::NullCodedAperture,
    coordinates::OpticalCoordinates,
    parameters = Dict{String,Float64}(),
)
    error("No Jacobian defined for NullCodedAperture")
end

# ------------------------------------------------------------------------------------------------------
#                                   Wide Field Optical Recipe
# ------------------------------------------------------------------------------------------------------

mutable struct WideFieldOpticalRecipe <: OpticalRecipe
    objective_mag::Float64
    objective_na::Float64
    wavelength::Float64
    medium_index::Float64
    f_tubelens::Float64
    f_objective::Float64
    static_aperture_mask::Any
    dynamic_aperture_mask::Any
    normalize_intensity::Bool
end

"""
Generic Optical Recipe for Widefield PSF Models

**Parameters**

    objective_mag - magnification of the objective
    objective_na - numerical aperture of the objective
    wavelength - emission wavelength of the light being imaged (all simulations are monochromatic)
    medium_index - index of refraction of the sample medium
    f_tubelens - tube lens focal length (e.g. 200mm for Nikon, 180mm for Olympus, etc.)
    static_aperture_mask - Supply a back aperture mask that is a subclass of CodedAperture, or
               'NullCodedAperture()' if there is no aperture code.  The static mask is computed
               when the PSF model is constructed, and is therefore more efficient than the
               dynamic_aperture_mask.
    dynamic_aperture_mask - Supply a back aperture mask that is a subclass of CodedAperture, or
               'NullCodedAperture()' if there is no aperture code.  The dynamic_aperture_mask
               is recomputed each time apsf() is called, but this flexibility incurs additional
               performance overhead.
    normalize_intensity - Set max PSF value @ z = 0 to 1.0.  This is an arbitrary choice, but
               one that keeps PSF values relatively well scaled and makes it easy to adjust the
               "exposure" of the simulated microscope by multipyling the PSF to get some number of
               photons that is >= 1.  However, be careful when mixing optical models with different
               recipes, because their absolute intensities cannot be directly compared with this
               arbitrary scaling in effect.
"""
function WideFieldOpticalRecipe(
    objective_mag::Number,
    objective_na::Float64,
    wavelength::Float64,
    medium_index::Float64,
    f_tubelens::Float64;
    static_aperture_mask::CodedAperture = NullCodedAperture(),
    dynamic_aperture_mask::CodedAperture = NullCodedAperture(),
    normalize_intensity::Bool = false,
)

    f_objective = f_tubelens / objective_mag
    return WideFieldOpticalRecipe(
        objective_mag,
        objective_na,
        wavelength,
        medium_index,
        f_tubelens,
        f_objective,
        static_aperture_mask,
        dynamic_aperture_mask,
        normalize_intensity,
    )
end


function initialize_converter!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::DynamicWrapper{
        PSY.DynamicInverter{PSY.AverageConverter, O, IC, DC, P, F},
    },
    inner_vars::AbstractVector,
) where {
    O <: PSY.OuterControl,
    IC <: PSY.InnerControl,
    DC <: PSY.DCSource,
    P <: PSY.FrequencyEstimator,
    F <: PSY.Filter,
} end

function initialize_converter!(
    device_states,
    static::PSY.StaticInjection,
    dynamic_device::DynamicWrapper{
        PSY.DynamicInverter{PSY.RenewableEnergyConverterTypeA, O, IC, DC, P, F},
    },
    inner_vars::AbstractVector,
) where {
    O <: PSY.OuterControl,
    IC <: PSY.InnerControl,
    DC <: PSY.DCSource,
    P <: PSY.FrequencyEstimator,
    F <: PSY.Filter,
}
    #Get inner vars
    V_R = inner_vars[Vr_inv_var]
    V_I = inner_vars[Vi_inv_var]
    θ = atan(V_I / V_R)
    V_t = sqrt(V_R^2 + V_I^2)
    Iq_external = inner_vars[Ii_cnv_var]
    Ip_external = inner_vars[Ir_cnv_var]
    #Reference Transformation
    Ip = Ip_external * cos(-θ) - Iq_external * sin(-θ)
    Iq = Ip_external * sin(-θ) + Iq_external * cos(-θ)
    converter = PSY.get_converter(dynamic_device)
    Io_lim = PSY.get_Io_lim(converter)
    Vo_lim = PSY.get_Vo_lim(converter)

    # Lv_pnt0 is unused in the initialization
    _, Lv_pnt1 = PSY.get_Lv_pnts(converter)

    if (Iq < Io_lim) || (V_t > Vo_lim) || (V_t < Lv_pnt1)
        error("Power flow solution outside of inverter limits. Update parameters.")
    end

    #Update converter states
    converter_ix = get_local_state_ix(dynamic_device, PSY.RenewableEnergyConverterTypeA)
    converter_states = @view device_states[converter_ix]
    converter_states[1] = Ip #Ip_cnv
    converter_states[2] = -Iq #Iq_cnv
    converter_states[3] = V_t #Vmeas

    #Update entry to converter currents (inner currents)
    inner_vars[Id_ic_var] = Ip
    # Multiplied by -1 because PSS/e doesn't do a conjugate.
    inner_vars[Iq_ic_var] = -Iq
end


using Test, MakieCore

# Main tests live in Makie.jl, but we should write some unit tests going forward!
using MakieCore: @recipe, Attributes, Plot
import MakieCore: plot!, convert_arguments, used_attributes, plot
import MakieCore: Observable, PointBased

struct AbstractTimeseriesSolution
    results::Vector{Float32}
end

function plot!(plot::Plot(AbstractTimeseriesSolution))
    # plot contains any keyword arguments that you pass to plot(series; kw...)
    var = get(plot, :var, Observable(5))
    density!(plot, map((v, r)-> v .* r.results, var, plot[1]))
end

struct Test2
    series::Any
end

struct Solution
    data::Any
end

function plot!(plot::Plot(Test2))
    arg1 = plot[1]
    scatter!(plot, arg1[].series)
    ser = AbstractTimeseriesSolution(arg1[].series)
    sol = Solution(arg1[].series)
    plot!(plot, ser, var = 10)
    scatter!(plot, sol, attribute = 3, color=:red)
end

used_attributes(::Any, x::Solution) = (:attribute,)

# Convert for all point based types (lines, scatter)
function convert_arguments(p::MakieCore.PointBased, x::Solution; attribute = 1.0)
    return convert_arguments(p, x.data .* attribute)
end


module Xkcd221

export Xkcd221RNG

using Random: Random, AbstractRNG, Sampler, SamplerType, LessThan

import Random: rand, seed!

mutable struct Xkcd221RNG <: AbstractRNG
    seed::UInt64
    Xkcd221RNG(seed::Integer = 4) = new(seed)
end

function Base.copy!(dst::Xkcd221RNG, src::Xkcd221RNG)
    dst.seed = src.seed
    dst
end

Base.copy(src::Xkcd221RNG) = Xkcd221RNG(src.seed)

Base.:(==)(x::Xkcd221RNG, y::Xkcd221RNG) = x.seed == y.seed

Base.hash(rng::Xkcd221RNG, h::UInt) = hash(rng.seed, 0x2bad40c9493e0718 % UInt ⊻ h)

## Sampling

rand(rng::Xkcd221RNG, ::SamplerType{UInt64}) = rng.seed

for T = [Bool, Base.BitInteger_types...]
    T === UInt64 && continue
    @eval rand(rng::Xkcd221RNG, ::SamplerType{$T}) = rand(rng, UInt64) % $T
end

Random.rng_native_52(::Xkcd221RNG) = UInt64

function rand(rng::Xkcd221RNG, sp::LessThan)
    x = rand(rng, sp.s)
    sp.sup < rng.seed ? x % sp.sup  : x
end

## Range Sampling

struct RangeSampler{T} <: Sampler{T}
    s::T
    m::UInt64
end

function RangeSampler(r::AbstractUnitRange{T}) where {T}
    isempty(r) && throw(ArgumentError("range must be non-empty"))
    m = (last(r) - first(r)) % unsigned(T) % UInt64
    RangeSampler(first(r), m)
end

function rand(rng::Xkcd221RNG, sp::RangeSampler{T}) where {T}
    x = !iszero(sp.m) ? ((sp.m>rng.seed ? rng.seed : (rng.seed%sp.m))) : 0
    (sp.s % UInt64 + x) % T
end

for T in Base.BitInteger_types
    @eval Sampler(::Type{Xkcd221RNG}, rg::AbstractUnitRange{$T}, ::Random.Repetition) = RangeSampler(rg)
end

end


## Morikawa and Fujiwara 2013 GMPE tests
@testset "Morikawa and Fujiwara 2013 GMPE PGA" begin
  # init model parameters
  include("../examples/morikawa-fujiwara-2013.conf")
  ## PGA at epicenter on grid M7.0, Dl = 250 (constant)
  @test gmpe_mf2013(eq_7,config_mf2013_crustal_pga,grid_epicenter)[1].pga == 53.28
  ## PGA,PGV,PSA on grid, M6, ASID false, Dl - constant
  S_c = gmpe_mf2013(eq_6,config_mf2013_crustal_pga,grid)
  @test length(S_c) == TEST_GRID_SIZE
  @test typeof(S_c) == Array{GroundMotion.Point_pga_out,1}
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 3.4
  S_c = gmpe_mf2013(eq_6,config_mf2013_crustal_pgv,grid)
  @test typeof(S_c) == Array{GroundMotion.Point_pgv_out,1}
  @test round(sum([S_c[i].pgv for i=1:length(S_c)]),digits=2) == 4.61
  S_c = gmpe_mf2013(eq_6,config_mf2013_crustal_psa_03,grid)
  @test typeof(S_c) == Array{GroundMotion.Point_psa_out,1}
  S_c = gmpe_mf2013(eq_6,config_mf2013_crustal_psa_10,grid)
  @test round(sum([S_c[i].psa for i=1:length(S_c)]),digits=2) == 5.45
  S_c = gmpe_mf2013(eq_6,config_mf2013_crustal_psa_30,grid)
  @test round(sum([S_c[i].psa for i=1:length(S_c)]),digits=2) == 1.43
  ## PGA modeling on grid M8.5, withoit minpga, ASID false, Dl - constant
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_pga,grid)
  @test length(S_c) == TEST_GRID_SIZE
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 37.6
  ## PGA modeling on grid M8.5, withoit minpga, ASID true, Dl - constant
  S_c = gmpe_mf2013(eq_85,config_mf2013_intraplate_pga_asid,grid)
  @test length(S_c) == TEST_GRID_SIZE
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 74.4
  ## run PGA modeling on grid M6, withoit minpga, ASID false, Dl on GRID
  S_c = gmpe_mf2013(eq_6,config_mf2013_crustal_pga,grid_dl)
  @test length(S_c) == TEST_GRID_SIZE
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 506.54
  ## run PGA modeling on grid M8.5, withoit minpga, ASID true, Dl on GRID
  S_c = gmpe_mf2013(eq_85,config_mf2013_intraplate_pga_asid,grid_dl,Xvf=40)
  @test length(S_c) == TEST_GRID_SIZE
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 4110.09
  ## PGV,PSA modeling on grid M6, withoit minpga, ASID false, Dl on GRID
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_pgv,grid_dl)
  @test length(S_c) == TEST_GRID_SIZE
  @test typeof(S_c) == Array{GroundMotion.Point_pgv_out,1}
  @test round(sum([S_c[i].pgv for i=1:length(S_c)]),digits=2) == 2989.47
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_psa_03,grid_dl)
  @test length(S_c) == TEST_GRID_SIZE
  @test typeof(S_c) == Array{GroundMotion.Point_psa_out,1}
  @test round(sum([S_c[i].psa for i=1:length(S_c)]),digits=2) == 4177.5
  ## PGA,PGV,PSA without grid Dl and VS30 by default, ASID=false
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_pga)
  @test length(S_c) == SIMULATION_ARRAY_SIZE
  @test round(sum(S_c),digits=2) == 5468.53
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_pgv)
  @test round(sum(S_c),digits=2) == 7710.7
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_psa_03)
  @test round(sum(S_c),digits=2) == 13116.13
  ## PGA without grid Dl=500 VS30 by default, ASID=false
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_pga,Dl=500)
  @test round(sum(S_c),digits=2) == 5725.89
  ## PGA without grid Dl by default VS30=500, ASID=false
  S_c = gmpe_mf2013(eq_85,config_mf2013_crustal_pga,VS30=500)
  @test round(sum(S_c),digits=2) == 4790.89
  ## ## PGA without grid Dl and VS30 by default, ASID=true, Xvf=40
  S_c = gmpe_mf2013(eq_85,config_mf2013_intraplate_pga_asid,Xvf=40)
  @test round(sum(S_c),digits=2) == 12845.95
end


#
println("Ck) Test of the MultipolePolarizibility module with ASF from an internally generated multiplet.")
#
setDefaults("print summary: open", "zzz-MultipolePolarizibility.sum")
wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(26.),
                        properties=[Polarity()],
                        configs=[Configuration("[Ne] 3s^2 3p^5"), Configuration("[Ne] 3s 3p^6")],
                        polaritySettings=MultipolePolarizibility.Settings(EmMultipole[], 0, 0, Float64[], false, false, Int64[]) )

wb = perform(wa)
setDefaults("print summary: close", "")


function construct_plankton(arch::Architecture, sp::Int64, params::Dict, maxN)
    rawdata = StructArray(x    = zeros(maxN), y    = zeros(maxN), z    = zeros(maxN),
                          xi   = zeros(Int,maxN), yi  = zeros(Int,maxN), zi  = zeros(Int,maxN),
                          CH   = zeros(maxN), NST  = zeros(maxN), PST  = zeros(maxN),
                          PRO  = zeros(maxN), DNA  = zeros(maxN), RNA  = zeros(maxN),
                          Chl  = zeros(maxN), gen  = zeros(maxN), age  = zeros(maxN),
                          ac   = zeros(maxN), idx  = zeros(maxN),
                          PS   = zeros(maxN), VDOC = zeros(maxN), VNH4 = zeros(maxN),
                          VNO3 = zeros(maxN), VPO4 = zeros(maxN), ρChl = zeros(maxN),
                          resp = zeros(maxN), S_PRO= zeros(maxN), S_DNA= zeros(maxN),
                          S_RNA= zeros(maxN), exu  = zeros(maxN),
                          graz = zeros(maxN), mort = zeros(maxN), dvid = zeros(maxN)
                          )
    data = replace_storage(array_type(arch), rawdata)

    param_names=(:Nsuper, :C_DNA, :mean, :var, :PRO2DNA, :RNA2DNA, :CH2DNA, :Chl2DNA, :α, :Φ, :T⁺, :Ea,
                 :PCmax, :VDOCmax, :VNO3max, :VNH4max, :VPO4max, :KsatDOC, :KsatNH4, :KsatNO3,
                 :KsatPO4, :CHmax, :CHmin, :NSTmax, :NSTmin, :PSTmax, :PSTmin,
                 :Chl2N, :R_NC_PRO, :R_NC_DNA, :R_NC_RNA, :R_PC_DNA, :R_PC_RNA, :respir,
                 :k_pro, :k_sat_pro, :k_rna, :k_sat_rna, :k_dna, :k_sat_dna,
                 :dvid_P, :grz_P, :mort_P, :mort_reg, :grazFracC, :grazFracN, :grazFracP, :mortFracC, :mortFracN, :mortFracP)

    pkeys = collect(keys(params))
    tmp = zeros(length(param_names))
    for i in 1:length(param_names)
        if length(findall(x->x==string(param_names[i]),pkeys))==0
            throw(ArgumentError("PARAM: parameter not found $(param_names[i])"))
        else
            tmp[i] = params[string(param_names[i])][sp]
        end
    end
    p = NamedTuple{param_names}(tmp)
    return plankton(data, p)
end

function generate_plankton!(plank, N::Int64, g::AbstractGrid, arch::Architecture)
    mean = plank.p.mean
    var = plank.p.var
    C_DNA = plank.p.C_DNA
    Nsuper = plank.p.Nsuper
    RNA2DNA = plank.p.RNA2DNA
    PRO2DNA = plank.p.PRO2DNA
    CH2DNA  = plank.p.CH2DNA
    Chl2DNA = plank.p.Chl2DNA

    plank.data.ac[1:N]  .= 1.0                                                                             # activity
    plank.data.gen[1:N] .= 1.0                                                                             # generation
    plank.data.age[1:N] .= 0.0                                                                             # age

    rand!(rng_type(arch), plank.data.DNA)
    rand!(rng_type(arch), plank.data.RNA)
    rand!(rng_type(arch), plank.data.PRO)
    rand!(rng_type(arch), plank.data.x)
    rand!(rng_type(arch), plank.data.y)
    rand!(rng_type(arch), plank.data.z)
    rand!(rng_type(arch), plank.data.CH)
    rand!(rng_type(arch), plank.data.NST)
    rand!(rng_type(arch), plank.data.PST)

    plank.data.DNA .= plank.data.DNA .* var .+ 1.0                # range: (1.0,1.0+var)
    plank.data.RNA .= plank.data.RNA .* var .* 2.0 .+ 1.0 .- var  # range: (1.0-var,1.0+var)
    plank.data.PRO .= plank.data.PRO .* var .* 2.0 .+ 1.0 .- var  # range: (1.0-var,1.0+var)
    plank.data.CH  .= plank.data.CH  .* var .* 2.0 .+ 1.0 .- var  # range: (1.0-var,1.0+var)
    plank.data.NST .= plank.data.NST .* var .* 2.0 .+ 1.0 .- var  # range: (1.0-var,1.0+var)
    plank.data.PST .= plank.data.PST .* var .* 2.0 .+ 1.0 .- var  # range: (1.0-var,1.0+var)

    plank.data.x   .=(plank.data.x .* g.Nx) .* plank.data.ac                                         # x, unit: grid spacing, starting from 0
    plank.data.y   .=(plank.data.y .* g.Ny) .* plank.data.ac                                         # y, unit: grid spacing, starting from 0
    plank.data.z   .=(plank.data.z .* g.Nz) .* plank.data.ac                                         # z, unit: grid spacing, starting from 0
    plank.data.DNA .= plank.data.DNA .* C_DNA .* Nsuper .* plank.data.ac                             # DNA mmolC/individual
    plank.data.RNA .= plank.data.RNA .* C_DNA .* Nsuper .* plank.data.ac .* RNA2DNA                  # RNA mmolC/individual
    plank.data.PRO .= plank.data.PRO .* C_DNA .* Nsuper .* plank.data.ac .* PRO2DNA                  # PRO mmolC/individual
    plank.data.CH  .= plank.data.CH  .* plank.data.DNA .* CH2DNA                                     # CH  mmolC/individual
    plank.data.NST .= plank.data.NST .* plank.data.CH .* 16 ./ 106                                   # NST mmolN/individual
    plank.data.PST .= plank.data.PST .* plank.data.CH ./ 106                                         # PST mmolP/individual
    plank.data.Chl .= plank.data.DNA .* Chl2DNA * 893.49 / 55.0                                      # Chl mgChl/individual

    mask_individuals!(plank.data, g, N, arch)
end

@inline function total_C_biomass(PRO, DNA, RNA, CH, Chl)
    C_tot = PRO + DNA + RNA + CH + Chl / 893.49 * 55.0
    return C_tot
end
@inline function total_N_biomass(PRO, DNA, RNA, NST, Chl, p)
    N_tot = PRO * p.R_NC_PRO + DNA * p.R_NC_DNA + RNA * p.R_NC_RNA + NST + Chl / 893.49 * 4.0
    return N_tot
end
@inline function total_P_biomass(DNA, RNA, PST, p)
    P_tot = DNA * p.R_PC_DNA + RNA * p.R_PC_RNA + PST
    return P_tot
end


# # Dry Atmosphere GCM diagnostics
#
# This file computes selected diagnostics for the GCM and outputs them
# on the spherical interpolated diagnostic grid.
#
# Use it by calling `Diagnostics.setup_atmos_default_diagnostics()`.
#
# TODO:
# - enable zonal means and calculation of covariances using those means
#   - ds.T_zm = mean(.*1., ds.T; dims = 3)
#   - ds.u_zm = mean((ds.u); dims = 3 )
#   - v_zm = mean(ds.v; dims = 3)
#   - w_zm = mean(ds.w; dims = 3)
#   - ds.uvcovariance = (ds.u .- ds.u_zm) * (ds.v .- v_zm)
#   - ds.vTcovariance = (ds.v .- v_zm) * (ds.T .- ds.T_zm)
# - add more variables, including horiz streamfunction from laplacial of vorticity (LN)
# - density weighting
# - maybe change thermo/dyn separation to local/nonlocal vars?

using LinearAlgebra
using Printf
using Statistics

using ..Atmos
using ..Atmos: thermo_state, turbulence_tensors

include("diagnostic_fields.jl")

function atmos_gcm_default_init(dgngrp::DiagnosticsGroup, currtime)
    if !(dgngrp.interpol isa InterpolationCubedSphere)
        @warn """
            Diagnostics ($dgngrp.name): currently requires `InterpolationCubedSphere`!
            """
    end
end

# 3D variables
function vars_atmos_gcm_default_simple_3d(atmos::AtmosModel, FT)
    @vars begin
        u::FT
        v::FT
        w::FT
        rho::FT
        temp::FT
        pres::FT
        thd::FT                 # θ_dry
        et::FT                  # e_tot
        ei::FT                  # e_int
        ht::FT
        hi::FT
        vort::FT                # Ω₃

        moisture::vars_atmos_gcm_default_simple_3d(atmos.moisture, FT)
    end
end
vars_atmos_gcm_default_simple_3d(::MoistureModel, FT) = @vars()
function vars_atmos_gcm_default_simple_3d(m::EquilMoist, FT)
    @vars begin
        qt::FT                  # q_tot
        ql::FT                  # q_liq
        qv::FT                  # q_vap
        qi::FT                  # q_ice
        thv::FT                 # θ_vir
        thl::FT                 # θ_liq

    end
end
num_atmos_gcm_default_simple_3d_vars(m, FT) =
    varsize(vars_atmos_gcm_default_simple_3d(m, FT))
atmos_gcm_default_simple_3d_vars(m, array) =
    Vars{vars_atmos_gcm_default_simple_3d(m, eltype(array))}(array)

function atmos_gcm_default_simple_3d_vars!(
    atmos::AtmosModel,
    state_conservative,
    thermo,
    dyni,
    vars,
)
    vars.u = state_conservative.ρu[1] / state_conservative.ρ
    vars.v = state_conservative.ρu[2] / state_conservative.ρ
    vars.w = state_conservative.ρu[3] / state_conservative.ρ
    vars.rho = state_conservative.ρ
    vars.temp = thermo.temp
    vars.pres = thermo.pres
    vars.thd = thermo.θ_dry
    vars.et = state_conservative.ρe / state_conservative.ρ
    vars.ei = thermo.e_int
    vars.ht = thermo.h_tot
    vars.hi = thermo.h_int

    vars.vort = dyni.Ω₃

    atmos_gcm_default_simple_3d_vars!(
        atmos.moisture,
        state_conservative,
        thermo,
        vars,
    )

    return nothing
end
function atmos_gcm_default_simple_3d_vars!(
    ::MoistureModel,
    state_conservative,
    thermo,
    vars,
)
    return nothing
end
function atmos_gcm_default_simple_3d_vars!(
    moist::EquilMoist,
    state_conservative,
    thermo,
    vars,
)
    vars.moisture.qt = state_conservative.moisture.ρq_tot / state_conservative.ρ
    vars.moisture.ql = thermo.moisture.q_liq
    vars.moisture.qv = thermo.moisture.q_vap
    vars.moisture.qi = thermo.moisture.q_ice
    vars.moisture.thv = thermo.moisutre.θ_vir
    vars.moisture.thl = thermo.moisture.θ_liq_ice

    return nothing
end

# Dynamic variables
function vars_dyn(FT)
    @vars begin
        Ω₁::FT
        Ω₂::FT
        Ω₃::FT
    end
end
dyn_vars(array) = Vars{vars_dyn(eltype(array))}(array)

"""
    atmos_gcm_default_collect(bl, currtime)

    Master function that performs a global grid traversal to compute various
    diagnostics using the above functions.
"""
function atmos_gcm_default_collect(dgngrp::DiagnosticsGroup, currtime)
    interpol = dgngrp.interpol
    if !(interpol isa InterpolationCubedSphere)
        @warn """
            Diagnostics ($dgngrp.name): currently requires `InterpolationCubedSphere`!
            """
        return nothing
    end

    mpicomm = Settings.mpicomm
    dg = Settings.dg
    Q = Settings.Q
    mpirank = MPI.Comm_rank(mpicomm)
    bl = dg.balance_law
    grid = dg.grid
    topology = grid.topology
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dimensionality(grid) == 2 ? 1 : Nq
    npoints = Nq * Nq * Nqk
    nrealelem = length(topology.realelems)
    nvertelem = topology.stacksize
    nhorzelem = div(nrealelem, nvertelem)

    # get needed arrays onto the CPU
    device = array_device(Q)
    if device isa CPU
        ArrayType = Array
        state_data = Q.realdata
        aux_data = dg.state_auxiliary.realdata
    else
        ArrayType = CuArray
        state_data = Array(Q.realdata)
        aux_data = Array(dg.state_auxiliary.realdata)
    end
    FT = eltype(state_data)

    # TODO: can this be done in one pass?
    #
    # Non-local vars, e.g. relative vorticity
    vgrad = VectorGradients(dg, Q)
    vort = Vorticity(dg, vgrad)

    # Compute thermo variables
    thermo_array = Array{FT}(undef, npoints, num_thermo(bl, FT), nrealelem)
    @visitQ nhorzelem nvertelem Nqk Nq begin
        state_conservative = extract_state_conservative(dg, state_data, ijk, e)
        state_auxiliary = extract_state_auxiliary(dg, aux_data, ijk, e)

        thermo = thermo_vars(bl, view(thermo_array, ijk, :, e))
        compute_thermo!(bl, state_conservative, state_auxiliary, thermo)
    end

    # Interpolate the state, thermo and dyn vars to sphere (u and vorticity
    # need projection to zonal, merid). All this may happen on the GPU.
    istate =
        ArrayType{FT}(undef, interpol.Npl, number_state_conservative(bl, FT))
    interpolate_local!(interpol, Q.realdata, istate)

    ithermo = ArrayType{FT}(undef, interpol.Npl, num_thermo(bl, FT))
    interpolate_local!(interpol, ArrayType(thermo_array), ithermo)

    idyn = ArrayType{FT}(undef, interpol.Npl, size(vort.data, 2))
    interpolate_local!(interpol, vort.data, idyn)

    # TODO: get indices here without hard-coding them
    _ρu, _ρv, _ρw = 2, 3, 4
    project_cubed_sphere!(interpol, istate, (_ρu, _ρv, _ρw))
    _Ω₁, _Ω₂, _Ω₃ = 1, 2, 3
    project_cubed_sphere!(interpol, idyn, (_Ω₁, _Ω₂, _Ω₃))

    # FIXME: accumulating to rank 0 is not scalable
    all_state_data = accumulate_interpolated_data(mpicomm, interpol, istate)
    all_thermo_data = accumulate_interpolated_data(mpicomm, interpol, ithermo)
    all_dyn_data = accumulate_interpolated_data(mpicomm, interpol, idyn)

    if mpirank == 0
        # get dimensions for the interpolated grid
        dims = dimensions(dgngrp.interpol)

        # adjust the level dimension for `planet_radius`
        level_val = dims["level"]
        dims["level"] =
            (level_val[1] .- FT(planet_radius(param_set)), level_val[2])

        # set up the array for the diagnostic variables based on the interpolated grid
        nlong = length(dims["long"][1])
        nlat = length(dims["lat"][1])
        nlevel = length(dims["level"][1])

        simple_3d_vars_array = Array{FT}(
            undef,
            1,
            nlong,
            nlat,
            nlevel,
            num_atmos_gcm_default_simple_3d_vars(bl, FT),
        )

        @visitI nlong nlat nlevel begin
            statei = Vars{vars_state_conservative(bl, FT)}(view(
                all_state_data,
                lo,
                la,
                le,
                :,
            ))
            thermoi = thermo_vars(bl, view(all_thermo_data, lo, la, le, :))
            dyni = dyn_vars(view(all_dyn_data, lo, la, le, :))
            simple_3d_vars = atmos_gcm_default_simple_3d_vars(
                bl,
                view(simple_3d_vars_array, 1, lo, la, le, :),
            )

            atmos_gcm_default_simple_3d_vars!(
                bl,
                statei,
                thermoi,
                dyni,
                simple_3d_vars,
            )
        end

        # assemble the diagnostics for writing
        dim_names = tuple("time", collect(keys(dims))...)
        varvals = OrderedDict()
        varnames = map(
            s -> startswith(s, "moisture.") ? s[10:end] : s,
            flattenednames(vars_atmos_gcm_default_simple_3d(bl, FT)),
        )
        for (vari, varname) in enumerate(varnames)
            var = Variables[varname]
            varvals[varname] = (
                dim_names,
                simple_3d_vars_array[:, :, :, :, vari],
                OrderedDict(
                    "units" => var.units,
                    "long_name" => var.long,
                    "standard_name" => var.standard,
                ),
            )
        end

        # write output
        dprefix = @sprintf(
            "%s_%s_%s_num%04d",
            dgngrp.out_prefix,
            dgngrp.name,
            Settings.starttime,
            dgngrp.num
        )
        dfilename = joinpath(Settings.output_dir, dprefix)
        write_data(dgngrp.writer, dfilename, dims, varvals, currtime)
    end

    MPI.Barrier(mpicomm)
    return nothing
end # function collect

function atmos_gcm_default_fini(dgngrp::DiagnosticsGroup, currtime) end


module TestFoldl

using Transducers: Map, ProductRF, TeeRF, __foldl__

using ..Utils: check_rrule

function test_simple()
    check_rrule(__foldl__, +, 0.0, 1:10)
    check_rrule(__foldl__, *, 1.0, 1:5)
    check_rrule(__foldl__, max, 0, [2, 6, 4, 9, 2, 7, 9, 5, 4, 9])
end

function test_tee()
    check_rrule(__foldl__, TeeRF(+, *), (0.0, 1.0), 1:5)
    check_rrule(__foldl__, TeeRF(+, *, max), (0.0, 1.0, 0.0), [2, 9, 6, 4])
    # TODO: Make sure that it can have non-differentiable parts in the accumulator?
    #=
    check_rrule(
        __foldl__,
        TeeRF(+, *, Map(string)'(*), max),
        (0.0, 1.0, "", 0),
        [2, 6, 4, 9, 2, 7, 9, 5, 4, 9],
    )
    =#
end

dup3(x) = (x, x, x)

function test_product()
    check_rrule(__foldl__, Map(dup3)'(ProductRF(+, *, *)), (0.0, 1.0, 0.0), [2, 9, 6, 4])
    # TODO: support zip-of-arrays
    #=
    check_rrule(__foldl__, ProductRF(+, *, *), (0.0, 1.0, 0.0), zip(1:5, 1:5, 'a':'e'))
    =#
end

end  # module


using MeshCat, MeshCatMechanisms, RigidBodyDynamics
using FileIO, MeshIO, GeometryBasics, CoordinateTransformations, Rotations, Meshing
using Plots

include_ddp()

contact_control_path = "/home/taylor/Research/ContactControl.jl/src"

using Parameters
# Utilities
include(joinpath(contact_control_path, "utils.jl"))

# Solver
include(joinpath(contact_control_path, "solver/cones.jl"))
include(joinpath(contact_control_path, "solver/interior_point.jl"))
include(joinpath(contact_control_path, "solver/lu.jl"))

# Environment
include(joinpath(contact_control_path, "simulator/environment.jl"))

# Dynamics
include(joinpath(contact_control_path, "dynamics/model.jl"))

# Simulator
include(joinpath(contact_control_path, "simulation/contact_methods.jl"))
include(joinpath(contact_control_path, "simulation/simulation.jl"))
include(joinpath(contact_control_path, "simulator/trajectory.jl"))

include(joinpath(contact_control_path, "dynamics/code_gen_dynamics.jl"))
include(joinpath(contact_control_path, "dynamics/fast_methods_dynamics.jl"))

# Models
include(joinpath(contact_control_path, "dynamics/quaternions.jl"))
include(joinpath(contact_control_path, "dynamics/mrp.jl"))
include(joinpath(contact_control_path, "dynamics/euler.jl"))

# Simulation
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))

include(joinpath(contact_control_path, "simulation/residual_approx.jl"))
include(joinpath(contact_control_path, "simulation/code_gen_simulation.jl"))

# Visuals
using MeshCatMechanisms
include(joinpath(contact_control_path, "dynamics/visuals.jl"))
include(joinpath(contact_control_path, "dynamics/visual_utils.jl"))

include(joinpath(pwd(),"models/kuka/kuka_utils.jl"))

urdf_path = joinpath(pwd(), "models/kuka/temp/kuka.urdf")

kuka = MeshCatMechanisms.parse_urdf(urdf_path, remove_fixed_tree_joints = true)
kuka_visuals = MeshCatMechanisms.URDFVisuals(urdf_path)

state = MechanismState(kuka)
state_cache = StateCache(kuka)
result = DynamicsResult(kuka)
result_cache = DynamicsResultCache(kuka)

vis = Visualizer()
mvis = MechanismVisualizer(kuka, kuka_visuals, vis[:base])
render(vis)
default_background!(vis)

nq = num_positions(kuka)
q_init = zeros(nq)
q_init[1] = 0
q_init[2] = pi/4

q_init[4] = -pi/2
q_init[5] = 0
q_init[6] = -pi/4
q_init[7] = 0
set_configuration!(state, q_init)
set_configuration!(mvis, q_init)
ee = findbody(kuka, "iiwa_link_7")
ee_point = Point3D(default_frame(ee), 0.0, 0.0, 0.0)
setelement!(mvis, ee_point, 0.05)
ee_jacobian_frame = ee_point.frame
ee_jacobian_path = path(kuka, root_body(kuka), ee)

world = root_frame(kuka)
ee_in_world = transform(state, ee_point, world).v
desired = Point3D(world, 0.5, 0.0, 0.0)
q_res1 = jacobian_transpose_ik!(state, ee, ee_point, desired,
    visualize = true, mvis = mvis)
q_res1 = Array(q_res1)
set_configuration!(mvis, q_res1)
set_configuration!(state, q_res1)
ee_in_world = transform(state, ee_point, world).v

desired = Point3D(world, 0.75, 0.0, 0.0)
q_res2 = jacobian_transpose_ik!(state, ee, ee_point, desired,
    visualize = true, mvis = mvis)
q_res2 = Array(q_res2)
set_configuration!(mvis, q_res2)
set_configuration!(state, q_res2)
ee_in_world = transform(state, ee_point, world).v

# Kuka iiwa arm parsed from URDF using RigidBodyDynamics.jl
# + 3D particle
struct KukaParticle <: ContactModel
	dim::Dimensions

	# particle
	mp

    state_cache1
    state_cache2
    state_cache3

    result_cache1
    result_cache2
    result_cache3

    world

	ee
    ee_point
    ee_jacobian_frame
    ee_jacobian_path
end

# Dimensions
nq = 7 + 3 # configuration dim
nu = 7 # control dim
nc = 1 # number of contact points
nf = 0 # number of faces for friction cone
nb = 0
ns = 0

n = 2 * nq
m_contact = 0 #nc + nb + nc + nb + 1
m = nu #nu + m_contact
d = 0

mp = 0.01

results_cache1 = DynamicsResultCache(kuka)
results_cache2 = DynamicsResultCache(kuka)
results_cache3 = DynamicsResultCache(kuka)

state_cache1 = StateCache(kuka)
state_cache2 = StateCache(kuka)
state_cache3 = StateCache(kuka)

function kuka_q(q)
	view(q, 1:7)
end

function particle_q(q)
	view(q, 8:10)
end

# Methods
function M_func(m::KukaParticle, q::AbstractVector{T}) where T
    # return Diagonal(m.mp * ones(3))
    state = m.state_cache3[T]
    result = m.result_cache3[T]
    set_configuration!(state, kuka_q(q))
    mass_matrix!(result.massmatrix, state)
    return SMatrix{10, 10}(cat(result.massmatrix, Diagonal(m.mp * ones(3)), dims = (1, 2)))
end

function C_func(m::KukaParticle, q, q̇)
    # return SVector{3}([0.0, 0.0, m.mp * 9.81])
    if eltype(q) <: ForwardDiff.Dual
		return SVector{10}([dynamics_bias_q(m, kuka_q(q), kuka_q(q̇))..., 0.0, 0.0, m.mp * 9.81])
	elseif eltype(q̇) <: ForwardDiff.Dual
		return SVector{10}([dynamics_bias_q̇(m, kuka_q(q), kuka_q(q̇))..., 0.0 ,0.0, m.mp * 9.81])
	else
		return SVector{10}([_dynamics_bias(m, kuka_q(q), kuka_q(q̇))..., 0.0, 0.0, m.mp * 9.81])
	end
end

function _dynamics_bias(m::KukaParticle, q, q̇)
	T = eltype(m.mp)
	state = m.state_cache1[T]
	result = m.result_cache1[T]
	set_configuration!(state, q)
	set_velocity!(state, q̇)
	dynamics_bias!(result, state)
    return result.dynamicsbias
end

function dynamics_bias_q(m::KukaParticle, q::AbstractVector{T}, q̇) where T
    state = m.state_cache2[T]
    result = m.result_cache2[T]
    set_configuration!(state, q)
    set_velocity!(state,q̇)
    dynamics_bias!(result, state)
    return result.dynamicsbias
end

function dynamics_bias_q̇(m::KukaParticle, q, q̇::AbstractVector{T}) where T
    state = m.state_cache3[T]
    result = m.result_cache3[T]
    set_configuration!(state, q)
    set_velocity!(state, q̇)
    dynamics_bias!(result, state)
    return result.dynamicsbias
end

function B_func(m::KukaParticle, q)
    # return zeros(3, 3)
    SMatrix{10, 7}([Diagonal(ones(7)); zeros(3, 7)])
end

function kinematics_ee(m::KukaParticle, q::AbstractVector{T}) where T
    state = m.state_cache3[T]
    set_configuration!(state, kuka_q(q))
    return transform(state, m.ee_point, m.world).v
end

function ϕ_func(m::KukaParticle, q::AbstractVector{T}) where T
	p_ee = kinematics_ee(m, q)
	p_p = q[8:10]
	diff = (p_ee - p_p)
	d_ee_p = sqrt(diff' * diff)
    SVector{1}([0.5 - d_ee_p])
end

function P_func(m::KukaParticle, q::AbstractVector{T}) where T
    state = m.state_cache3[T]

    set_configuration!(state, kuka_q(q))

    pj1 = PointJacobian(transform(state, m.ee_point, m.world).frame,
		zeros(T, 3, 7))

    ee_in_world = transform(state, m.ee_point, m.world)

    point_jacobian!(pj1, state, m.ee_jacobian_path,ee_in_world) #TODO confirm this is correct

	p_ee = ee_in_world.v
	p_p = q[8:10]
	diff = (p_ee - p_p)

    return Array(-1.0 * 2.0 * diff' * [pj1.J[1:3,:] -1.0 * Diagonal(ones(3))])
end

ϕ(z) = ϕ_func(_model, z)
norm(ForwardDiff.jacobian(ϕ, q0) - P_func(_model, q1))

nq = 7 + 3
nu = 7
nw = 0
nc = 1

h = 0.1

_model = KukaParticle(
	Dimensions(nq, nu, nw, nc),
	mp,
    state_cache1, state_cache2, state_cache3,
    results_cache1, results_cache2, results_cache3,
    world,
	ee,
	ee_point,
	ee_jacobian_frame,
	ee_jacobian_path)

q0 = zeros(_model.dim.q)
q0[1] = 0
q0[3] = 0
q0[4] = -pi/2
q0[5] = 0.0
p_pos0 = Array(kinematics_ee(_model, q0))
p_pos0[1] += 0.0
p_pos0[3] -= 0.5
q0[8:10] = p_pos0

# visualize!(mvis, _model, [q0], Δt = h)

q1 = copy(q0)

qN = zeros(_model.dim.q)
qN[1] = 0
qN[3] = 0
qN[4] = -pi/2
qN[5] = 0.

ee_posN = Array(kinematics_ee(_model, qN))
p_posN = ee_posN
p_posN[1] += 0.0
p_posN[3] += 0.1
qN[8:10] = p_posN

qD = zeros(_model.dim.q)
qD[1] = 0
qD[3] = 0
qD[4] = -pi/2
qD[5] = 0.

ee_posD = Array(kinematics_ee(_model, qD))
p_posD = ee_posD
# p_posD[1] += 0.5
qD[8:10] = p_posD

visualize!(mvis, _model, [qD], Δt = h)

# Visualization
function visualize!(mvis, model::KukaParticle, q;
		verbose = false, r_ball = 0.035, Δt = 0.1)

	setobject!(vis["ball"], Sphere(Point3f0(0),
				convert(Float32, r_ball)),
				MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0,1.0)))

	settransform!(vis["ball"], compose(Translation(0.66,3.0,0.0)))

	state = model.state_cache1[Float64]

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q)
    for t = 1:T
        q_kuka = kuka_q(q[t])
		q_particle = particle_q(q[t])
		set_configuration!(state, kuka_q(q[t]))

        MeshCat.atframe(anim,t) do
			set_configuration!(mvis,q_kuka)
            settransform!(vis["ball"], compose(Translation(q_particle), LinearMap(RotZ(0))))
		end
    end

    MeshCat.setanimation!(vis, anim)
end

include(joinpath(contact_control_path, "simulator/environment.jl"))
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))

env = flat_3D_lc

function lagrangian_derivatives(model::KukaParticle, q, v)
	D1L = -1.0 * C_func(model, q, v)
    D2L = M_func(model, q) * v
	return D1L, D2L
end

function gravity_compensation(model::KukaParticle, h, q0, q1)
    q2 = q1
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	# return 0.0
	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2)[1:7]
end

function dynamics(model::KukaParticle, h, q0, q1, u1, λ1, q2)
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

    u_gravity = gravity_compensation(model, h, q0, q1)

	# return 0.0
	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * (u1 - u_gravity)
        + transpose(P_func(model, q2)) * λ1)
        # - h[1] * 0.5 .* vm2)
end

function residual(model::KukaParticle, env::Environment{<:World,LinearizedCone}, z, θ, κ)
	# nc = model.dim.c
	# nb = nc * friction_dim(env)
	# nf = Int(nb / nc)
	# np = dim(env)

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]
    λ1 = z[nq .+ (1:nc)]
    s1 = z[nq + nc .+ (1:nc)]

	# q0, q1, u1, w1, μ, h = unpack_θ(model, θ)
	# q2, γ1, b1, ψ1, η1, s1, s2 = unpack_z(model, env, z)

	# ϕ = ϕ_func(model, env, q2)
    #
	# k = kinematics(model, q2)
	# λ1 = contact_forces(model, env, γ1, b1, q2, k)
	# Λ1 = transpose(J_func(model, env, q2)) * λ1 #@@@@ maybe need to use J_fast
	# vT_stack = velocity_stack(model, env, q1, q2, k, h)
	# ψ_stack = transpose(E_func(model, env)) * ψ1
    # return dynamics(model, h, q0, q1, u1, zeros(nc), q2)

    [
     dynamics(model, h, q0, q1, u1, λ1, q2);
     s1 .- ϕ_func(model, q2);
     λ1 .* s1 .- κ;
    ]
	# [
    #  dynamics(model, h, q0, q1, u1, q2);
	#  # s1 - ϕ;
	#  # vT_stack + ψ_stack - η1;
	#  # s2 .- (μ[1] * γ1 .- E_func(model, env) * b1);
	#  # γ1 .* s1 .- κ;
	#  # b1 .* η1 .- κ;
	#  # ψ1 .* s2 .- κ
    #  ]
end

nz = nq + nc + nc
nθ = nq + nq + nu + 1

u0 = zeros(_model.dim.u)
z0 = copy([q1; 0.1 * ones(2 * nc)])
θ0 = copy([q0; q1; u0; h])

residual(_model, env, z0, θ0, [1.0])

function r_func(r, z, θ, κ)
    r .= residual(_model, env, z, θ, κ)
end

function rz_func(rz, z, θ)
    r(a) = residual(_model, env, a, θ, 0.0)
    rz .= ForwardDiff.jacobian(r, z)
end

function rθ_func(rθ, z, θ)
    r(a) = residual(_model, env, z, a, 0.0)
    rθ .= ForwardDiff.jacobian(r, θ)
end

r0 = zeros(nz)
rz0 = zeros(nz, nz)
rθ0 = zeros(nz, nθ)
r_func(r0, z0, θ0, [1.0])
rz_func(rz0, z0, θ0)
rθ_func(rθ0, z0, θ0)

# options
opts = InteriorPointOptions(
   κ_init = 0.1,
   κ_tol = 1.0e-4,
   r_tol = 1.0e-8,
   diff_sol = true)

idx_ineq = collect([nq .+ (1:(nc + nc))]...)

# solver
ip = interior_point(z0, θ0,
   r! = r_func, rz! = rz_func,
   rz = similar(rz0, Float64),
   rθ! = rθ_func,
   rθ = similar(rθ0, Float64),
   idx_ineq = idx_ineq,
   opts = opts)

# simulate
T = 26
q_hist = [q0, q0]
u_hist = []
h = 0.1
for t = 1:T-1
    u_gravity = zeros(_model.dim.u)#-1.0 * gravity_compensation(_model, h, q_hist[end-1], q_hist[end])
    push!(u_hist, 0.0 * u_gravity)
    # if t == 5
    #     u_gravity[6] = -1.0
    # end
    # u_gravity = -1.0 * gravity_compensation(_model, h, q0, q1)

    ip.z .= copy([q_hist[end]; 0.1 * ones(nc + nc)])
    ip.θ .= copy([q_hist[end-1]; q_hist[end];
        # u0;
        u_hist[end];
        h])
    status = interior_point_solve!(ip)

    if status
        push!(q_hist, ip.z[1:nq])
    else
        println("dynamics failure t = $t")
        println("res norm = $(norm(ip.r, Inf))")
        break
    end
end

vz_particle = [(((q_hist[t+2] - q_hist[t+1]) / h - (q_hist[t+1] - q_hist[t]) / h) / h)[10] for t = 1:T-2]

visualize!(mvis, _model, q_hist, Δt = h)

s = Simulation(_model, env)
s.res.r! = r_func
s.res.rz! = rz_func
s.res.rθ! = rθ_func
s.rz = ip.rz
s.rθ = ip.rθ

struct Dynamics{T}
	s::Simulation
	ip_dyn::InteriorPoint
	ip_jac::InteriorPoint
	h::T
end

function gen_dynamics(s::Simulation, h;
		dyn_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true),
		jac_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true))

	z = zeros(nz)
	θ = zeros(nθ)

	ip_dyn = interior_point(z, θ,
		idx_ineq = idx_ineq,
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = dyn_opts)

	ip_dyn.opts.diff_sol = false

	ip_jac = interior_point(z, θ,
		idx_ineq = idx_ineq,
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = jac_opts)

	ip_jac.opts.diff_sol = true

	Dynamics(s, ip_dyn, ip_jac, h)
end

d = gen_dynamics(s, h,
	dyn_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-4, κ_init = 0.1),
	jac_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-2, κ_init = 0.1))

function f!(d::Dynamics, q0, q1, u1, mode = :dynamics)
	s = d.s
	ip = (mode == :dynamics ? d.ip_dyn : d.ip_jac)
	h = d.h

    ip.z .= copy([q1; 0.1 * ones(nc + nc)])
	ip.θ .= copy([q0; q1; u1; h])

	status = interior_point_solve!(ip)

	!status && (@warn "dynamics failure")
end

function f(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :dynamics)
	return copy(d.ip_dyn.z[1:d.s.model.dim.q])
end

f(d, q0, q1, zeros(m))

function fq0(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:d.s.model.dim.q])
end

fq0(d, q0, q1, zeros(m))

function fq1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, d.s.model.dim.q .+ (1:d.s.model.dim.q)])
end

fq1(d, q0, q1, zeros(m))

function fx1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:(2 * d.s.model.dim.q)])
end

fx1(d, q0, q1, zeros(m))

function fu1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 2 * d.s.model.dim.q .+ (1:d.s.model.dim.u)])
end

fu1(d, q0, q1, zeros(m))

struct KukaParticleI{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int
	dynamics::Dynamics
end

model = KukaParticleI{Midpoint, FixedTime}(2 * s.model.dim.q, s.model.dim.u, 0, d)

function fd(model::KukaParticleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]

	q2 = f(model.dynamics, q0, q1, u)

	return [q1; q2]
end

fd(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

function fdx(model::KukaParticleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2dx1 = fx1(model.dynamics, q0, q1, u)

	return [zeros(nq, nq) I; dq2dx1]
end

fdx(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

function fdu(model::KukaParticleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2du1 = fu1(model.dynamics, q0, q1, u)
	return [zeros(nq, model.m); dq2du1]
end

fdu(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

T = 21
x1 = [q1; q1]
xT = [qN; qN]
ū = [0.1 * randn(_model.dim.u) for t = 1:T-1]#[t == 1 ? [0.0, 10.0, 0.0, 0.0, 0.0, -0.0, 0.0] : t == 2 ? [0.0, -10.0, 0.0, -0.0, 0.0, 0.0, 0.0] : [0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0] for t = 1:T-1] #u_hist #[-1.0 * gravity_compensation(_model, h, q1, q1) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)
q̄ = state_to_configuration(x̄)
visualize!(mvis, _model, q̄, Δt = h)

# visualize!(mvis, _model, [x_track_extend[1:nq]], Δt = h)

# Objective
V = Diagonal(ones(s.model.dim.q))
Q_velocity = 1.0e-3 * [V -V; -V V] ./ h^2.0
Q_track = 1.0 * Diagonal(ones(model.n))
Q_track_extend = copy(Q_track)
Q_track_extend[8, 8] = 10.0
Q_track_extend[nq + 8, nq + 8] = 10.0
x_track = [qD; qD]
x_track_extend = copy(x_track)
x_track_extend[8] += 0.5
x_track_extend[nq + 8] += 0.5
Q = [t < T ? Q_velocity + (t == 11 ? Q_track_extend : 1.0 * Q_track) : Q_velocity + 1.0 * Q_track for t = 1:T]
q = [t < T ? -2.0 * (t == 11 ? Q_track_extend * x_track_extend : 1.0 * Q_track * x_track) : -2.0 * 1.0 * Q_track * x_track for t = 1:T]
R = [Diagonal(1.0e-3 * ones(model.m)) for t = 1:T-1]
r = [zeros(model.m) for t = 1:T-1]

obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
ul = -100.0 * ones(model.m)
uu = 100.0 * ones(model.m)
p = [t < T ? (t == 11 ? 2 * m + 0 * nq : 2 * m) : n for t = 1:T]
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m))
info_m = Dict(:qD => x_track_extend[1:nq], :ul => ul, :uu => uu, :inequality => (1:2 * m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? (t == 11 ? info_m : info_t) : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c[1:2 * m] .= [ul - u; u - uu]
        # if t == 11
        #     c[2 * m .+ (1:nq)] = x[1:nq] - cons.con[t].info[:qD]
        # end
	else
		c[1:n] .= x - cons.con[T].info[:xT]
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T,
	analytical_dynamics_derivatives = true)

ū = [t == 1 ? [0.0, 6.5, 0.0, 0.0, 0.0, -0.0, 0.0] : t == 2 ? [0.0, -6.5, 0.0, -0.0, 0.0, 0.0, 0.0] : [0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0] for t = 1:T-1] #u_hist #[-1.0 * gravity_compensation(_model, h, q1, q1) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)
q̄ = state_to_configuration(x̄)
visualize!(mvis, _model, q̄, Δt = h)
open(mvis)
# Solve
@time constrained_ddp_solve!(prob,
	max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0,
	con_tol = 0.005)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

q̄ = state_to_configuration(x̄)
visualize!(mvis, _model, q̄, Δt = h)


function optimize!(model::Model)
  @assert is_valid(model.env)
  stat = (if model.has_int
    @cpx_ccall_intercept(model, mipopt, Cint, (Ptr{Void}, Ptr{Void}), model.env.ptr, model.lp)
  elseif model.has_qc
    @cpx_ccall_intercept(model, qpopt, Cint, (Ptr{Void}, Ptr{Void}), model.env.ptr, model.lp)
  else
    @cpx_ccall_intercept(model, lpopt, Cint, (Ptr{Void}, Ptr{Void}), model.env.ptr, model.lp)
    end)
  if stat != 0
    throw(CplexError(model.env, stat))
  end
end

set_branching_priority(model::Model, priority) =
    set_branching_priority(model, Cint[1:num_var(model)], priority)

set_branching_priority(model::Model, indices, priority) =
    set_branching_priority(model, indices, priority, C_NULL)

set_branching_priority(model, indices, priority, direction) =
    set_branching_priority(model, convert(Vector{Cint},indices), convert(Vector{Cint},priority), direction)

function set_branching_priority(model::Model, indices::Vector{Cint}, priority::Vector{Cint}, direction)
    @assert (cnt = length(indices)) == length(priority)
    isa(direction,Vector) && @assert cnt == length(direction)
    stat = @cpx_ccall(copyorder, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Cint,
                      Ptr{Cint},
                      Ptr{Cint},
                      Ptr{Cint}),
                      model.env.ptr, model.lp, cnt, indices.-1, priority, direction)
    stat == 0 || throw(CplexError(model.env, stat))
    return nothing
end

function get_objval(model::Model)
  objval = Array(Cdouble, 1)
  stat = @cpx_ccall(getobjval, Cint, (
                    Ptr{Void},
                    Ptr{Void},
                    Ptr{Cdouble}
                    ),
                    model.env.ptr, model.lp, objval)
  if stat != 0
    throw(CplexError(model.env, stat))
  end
  return objval[1]
end

function get_solution(model::Model)
  nvars = num_var(model)
  x = Array(Cdouble, nvars)
  stat = @cpx_ccall(getx, Cint, (
                    Ptr{Void},
                    Ptr{Void},
                    Ptr{Cdouble},
                    Cint,
                    Cint
                    ),
                    model.env.ptr, model.lp, x, 0, nvars-1)
  if stat != 0
    throw(CplexError(model.env, stat))
  end
  return x
end

function get_reduced_costs(model::Model)
    nvars = num_var(model)
    p = Array(Cdouble, nvars)
    status = Array(Cint, 1)
    stat = @cpx_ccall(getdj, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cdouble},
                      Cint,
                      Cint
                      ),
                      model.env.ptr, model.lp, p, 0, nvars-1)
    if stat != 0
       throw(CplexError(model.env, stat))
   end
   return p
end

function get_constr_duals(model::Model)
    ncons = num_constr(model)
    p = Array(Cdouble, ncons)
    status = Array(Cint, 1)
    stat = @cpx_ccall(getpi, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cdouble},
                      Cint,
                      Cint
                      ),
                      model.env.ptr, model.lp, p, 0, ncons-1)
    if stat != 0
       throw(CplexError(model.env, stat))
   end
   return p
end

function get_constr_solution(model::Model)
  ncons = num_constr(model)
  Ax = Array(Cdouble, ncons)
  stat = @cpx_ccall(getax, Cint, (
                    Ptr{Void},
                    Ptr{Void},
                    Ptr{Cdouble},
                    Cint,
                    Cint
                    ),
                    model.env.ptr, model.lp, Ax, 0, ncons-1)
  if stat != 0
    throw(CplexError(model.env, stat))
  end
  return Ax
end

function get_infeasibility_ray(model::Model)
  ncons = num_constr(model)
  y = Array(Cdouble, ncons)
  proof_p = Array(Cdouble, 1)
  stat = @cpx_ccall(dualfarkas, Cint, (
                    Ptr{Void},
                    Ptr{Void},
                    Ptr{Cdouble},
                    Ptr{Cdouble}
                    ),
                    model.env.ptr, model.lp, y, proof_p)
  if stat != 0
    error("CPLEX is unable to grab infeasible ray; consider resolving with presolve turned off")
    throw(CplexError(model.env, stat))
  end
  return y
end

function get_unbounded_ray(model::Model)
  solve_stat = get_status(model)
  if solve_stat == :CPX_STAT_UNBOUNDED
    n = num_var(model)
    z = Array(Cdouble, n)
    stat = @cpx_ccall(getray, Cint, (
                      Ptr{Void},
                      Ptr{Void},
                      Ptr{Cdouble}
                      ),
                      model.env.ptr, model.lp, z)
    if stat != 0
      throw(CplexError(model.env, stat))
    end
    return z
  else
    error("CPLEX is unable to grab unbounded ray; consider resolving with presolve turned off")
  end
end

const varmap = Dict(
    0 => :NonbasicAtLower,
    1 => :Basic,
    2 => :NonbasicAtUpper,
    3 => :Free
)

const conmap = Dict(
    0 => :NonbasicAtLower,
    1 => :Basic,
    2 => :NonbasicAtUpper
)

function get_basis(model::Model)
    cval = Array(Cint, num_var(model))
    rval = Array(Cint, num_constr(model))
    stat = @cpx_ccall(getbase, Cint, (Ptr{Void},Ptr{Void},Ptr{Cint},Ptr{Cint}),
                      model.env.ptr, model.lp, cval, rval)
    stat != 0 && throw(CplexError(model.env, stat))

    csense = get_constr_senses(model)
    cbasis = Array(Symbol, num_var(model))
    rbasis = Array(Symbol, num_constr(model))
    for it in 1:num_var(model)
        cbasis[it] = varmap[cval[it]]
    end
    for it in 1:num_constr(model)
        rbasis[it] = conmap[rval[it]]
        if (rbasis[it] == :NonbasicAtLower) && (csense[it] == convert(Cchar,'L'))
            rbasis[it] = :NonbasicAtUpper
        end
    end
    return cbasis, rbasis
end

get_node_count(model::Model) = @cpx_ccall(getnodecnt, Cint, (Ptr{Void},Ptr{Void}), model.env.ptr, model.lp)

function get_rel_gap(model::Model)
  ret = Array(Cdouble,1)
  stat = @cpx_ccall(getmiprelgap, Cint, (Ptr{Void},Ptr{Void},Ptr{Cdouble}), model.env.ptr, model.lp, ret)
  if stat != 0
    throw(CplexError(model.env, stat))
  end
  ret[1]
end

function get_num_cuts(model::Model,cuttype)
    cutcount = Array(Cint,1)

    stat = @cpx_ccall(getnumcuts, Cint, (Ptr{Void},Ptr{Void},Cint,Ptr{Void}), model.env.ptr , model.lp, cuttype, cutcount)
    if stat != 0
        error(CplexError(model.inner.env, stat).msg)
    end
    return cutcount[1]
end

const status_symbols = Dict(
    1   => :CPX_STAT_OPTIMAL,
    2   => :CPX_STAT_UNBOUNDED,
    3   => :CPX_STAT_INFEASIBLE,
    4   => :CPX_STAT_INForUNBD,
    5   => :CPX_STAT_OPTIMAL_INFEAS,
    6   => :CPX_STAT_NUM_BEST,
    7   => :CPX_STAT_FEASIBLE_RELAXED,
    8   => :CPX_STAT_OPTIMAL_RELAXED,
    10  => :CPX_STAT_ABORT_IT_LIM,
    11  => :CPX_STAT_ABORT_TIME_LIM,
    12  => :CPX_STAT_ABORT_OBJ_LIM,
    13  => :CPX_STAT_ABORT_USER,
    20  => :CPX_STAT_OPTIMAL_FACE_UNBOUNDED,
    21  => :CPX_STAT_ABORT_PRIM_OBJ_LIM,
    22  => :CPX_STAT_ABORT_DUAL_OBJ_LIM,
    101 => :CPXMIP_OPTIMAL,
    102 => :CPXMIP_OPTIMAL_TOL,
    103 => :CPXMIP_INFEASIBLE,
    104 => :CPXMIP_SOL_LIM,
    105 => :CPXMIP_NODE_LIM_FEAS,
    106 => :CPXMIP_NODE_LIM_INFEAS,
    107 => :CPXMIP_TIME_LIM_FEAS,
    108 => :CPXMIP_TIME_LIM_INFEAS,
    109 => :CPXMIP_FAIL_FEAS,
    110 => :CPXMIP_FAIL_INFEAS,
    111 => :CPXMIP_MEM_LIM_FEAS,
    112 => :CPXMIP_MEM_LIM_INFEAS,
    113 => :CPXMIP_ABORT_FEAS,
    114 => :CPXMIP_ABORT_INFEAS,
    115 => :CPXMIP_OPTIMAL_INFEAS,
    116 => :CPXMIP_FAIL_FEAS_NO_TREE,
    117 => :CPXMIP_FAIL_INFEAS_NO_TREE,
    118 => :CPXMIP_UNBOUNDED,
    119 => :CPXMIP_INForUNBD,
    120 => :CPXMIP_FEASIBLE_RELAXED,
    121 => :CPXMIP_OPTIMAL_RELAXED
)

get_status(model::Model) = status_symbols[Int(get_status_code(model))]::Symbol
get_status_code(model::Model) = @cpx_ccall(getstat, Cint, (Ptr{Void}, Ptr{Void}), model.env.ptr, model.lp)


struct BandedBlockBandedSizes <: AbstractBlockSizes{2}
    block_sizes::DefaultBlockSizes{2}
    data_block_sizes::DefaultBlockSizes{2}
    l::Int
    u::Int
    λ::Int
    μ::Int
end

function BandedBlockBandedSizes(bs::BlockSizes{2}, l, u, λ, μ)
    # data matrix has row blocks all the same size but same column blocks
    # we access the cumul vec directly to reuse bs.cumul_sizes[2]
    d_bs = BlockSizes((BlockArrays._cumul_vec(fill(max(λ+μ+1,0), max(l+u+1,0))),bs.cumul_sizes[2]))
    BandedBlockBandedSizes(bs, d_bs, l, u, λ, μ)
end

BandedBlockBandedSizes(bs::AbstractBlockSizes{2}, l, u, λ, μ) =
    BandedBlockBandedSizes(convert(BlockSizes{2},bs), l, u, λ, μ)

BandedBlockBandedSizes(rows::AbstractVector{Int}, cols::AbstractVector{Int}, l, u, λ, μ) =
    BandedBlockBandedSizes(BlockSizes(rows,cols), l, u, λ, μ)

cumulsizes(B::BandedBlockBandedSizes) = cumulsizes(B.block_sizes)

convert(::Type{BlockSkylineSizes}, B::BandedBlockBandedSizes) =
    BlockBandedSizes(B.block_sizes, B.l, B.u)
convert(::Type{BlockBandedSizes}, B::BandedBlockBandedSizes) =
    BlockBandedSizes(B.block_sizes, B.l, B.u)

convert(::Type{BlockSizes}, B::BandedBlockBandedSizes) = B.block_sizes
convert(::Type{BlockSizes{2}}, B::BandedBlockBandedSizes) = B.block_sizes
convert(::Type{DefaultBlockSizes{2}}, B::BandedBlockBandedSizes) = B.block_sizes

const BandedOrBlockBandedSizes = Union{BandedBlockBandedSizes,BlockBandedSizes}

==(A::BandedOrBlockBandedSizes, B::BandedOrBlockBandedSizes) = A.block_sizes == B.block_sizes
==(A::BandedOrBlockBandedSizes, B::BlockSkylineSizes) = A.block_sizes == B.block_sizes
==(A::BlockSkylineSizes, B::BandedOrBlockBandedSizes) = A.block_sizes == B.block_sizes
==(A::BandedOrBlockBandedSizes, B::AbstractBlockSizes) = A.block_sizes == B
==(A::AbstractBlockSizes, B::BandedOrBlockBandedSizes) = A == B.block_sizes

BlockBandedSizes(B::BandedBlockBandedSizes) = convert(BlockBandedSizes, B)
BlockSkylineSizes(B::BandedBlockBandedSizes) = convert(BlockSkylineSizes, B)

function check_data_sizes(data::AbstractBlockMatrix, B::BandedBlockBandedSizes)
    bs = data.block_sizes
    c_rows, c_cols = bs.cumul_sizes
    if length(c_rows) ≠ B.l + B.u + 2 && !(length(c_rows) == 1 && -B.l > B.u)
        throw(ArgumentError("Data matrix must have number of row blocks equal to number of block bands"))
    end
    for k = 1:length(c_rows)-1
        if c_rows[k+1]-c_rows[k] ≠ B.λ + B.μ + 1 && !(c_rows[k+1] == c_rows[k]  && -B.λ > B.μ)
            throw(ArgumentError("Data matrix must have row block sizes equal to number of subblock bands"))
        end
    end
    if c_cols ≠ B.block_sizes.cumul_sizes[2]
        throw(ArgumentError("Data matrix must have same column blocks as matrix"))
    end
end

function _BandedBlockBandedMatrix end

# Represents a block banded matrix with banded blocks
#   similar to BandedMatrix{BandedMatrix{T,Matrix{T}}}
# Here the data is stored by blocks, in a way that is consistent with
# BandedMatrix
#

struct BandedBlockBandedMatrix{T, BLOCKS} <: AbstractBlockBandedMatrix{T}
    data::BLOCKS
    block_sizes::BandedBlockBandedSizes

    l::Int  # block lower bandwidth
    u::Int  # block upper bandwidth
    λ::Int  # sub lower bandwidth
    μ::Int  # sub upper bandwidth

    global function _BandedBlockBandedMatrix(data::AbstractBlockMatrix,
                                             block_sizes::BandedBlockBandedSizes)
      @boundscheck check_data_sizes(data, block_sizes)
      new{eltype(data), typeof(data)}(data, block_sizes, block_sizes.l,
                                      block_sizes.u, block_sizes.λ,
                                      block_sizes.μ)
    end
end

const DefaultBandedBlockBandedMatrix{T} = BandedBlockBandedMatrix{T, PseudoBlockMatrix{T, Matrix{T}}}

@inline _BandedBlockBandedMatrix(data::AbstractMatrix, block_sizes::BandedBlockBandedSizes) =
    _BandedBlockBandedMatrix(PseudoBlockArray(data, block_sizes.data_block_sizes), block_sizes)

BandedBlockBandedMatrix{T, B}(::UndefInitializer, block_sizes::BandedBlockBandedSizes) where {T, B} =
    _BandedBlockBandedMatrix(B(undef, block_sizes.data_block_sizes), block_sizes)

"""
    BandedBlockBandedMatrix{T}(undef, (rows, cols), (l, u), (λ, μ))

returns an undef `sum(rows)`×`sum(cols)` banded-block-banded matrix `A`
of type `T` with block-bandwidths `(l,u)` and where `A[Block(K,J)]`
is a `BandedMatrix{T}` of size `rows[K]`×`cols[J]` with bandwidths `(λ,μ)`.
"""
BandedBlockBandedMatrix

BandedBlockBandedMatrix{T, B}(::UndefInitializer, dims::NTuple{2, AbstractVector{Int}},
                              lu::NTuple{2, Int}, λμ::NTuple{2, Int}) where {T, B} =
    BandedBlockBandedMatrix{T, B}(undef, BandedBlockBandedSizes(dims..., lu..., λμ...))

BandedBlockBandedMatrix{T, B}(::UndefInitializer, bs::AbstractBlockSizes,
                              lu::NTuple{2, Int}, λμ::NTuple{2, Int}) where {T, B} =
    BandedBlockBandedMatrix{T, B}(undef, BandedBlockBandedSizes(bs, lu..., λμ...))

BandedBlockBandedMatrix{T}(::UndefInitializer, args...) where T =
    DefaultBandedBlockBandedMatrix{T}(undef, args...)

# Auxiliary outer constructors
@inline _BandedBlockBandedMatrix(data::AbstractMatrix, dims::NTuple{2, AbstractVector{Int}},
                                         lu::NTuple{2, Int}, λμ::NTuple{2, Int}) =
    _BandedBlockBandedMatrix(data, BandedBlockBandedSizes(dims..., lu..., λμ...))

function convert(::Type{<:BandedBlockBandedMatrix}, B::BandedMatrix)
    if isdiag(B)
        _BandedBlockBandedMatrix(copy(B.data), (fill(1,size(B,1)),fill(1,size(B,2))), (0,0), (0,0))
    else
        _BandedBlockBandedMatrix(copy(B.data), [size(B,1)], [size(B,2)], (0,0), (B.l,B.u))
    end
end

function BandedBlockBandedMatrix{T, B}(Z::Zeros, dims::NTuple{2,AbstractVector{Int}},
                                       lu::NTuple{2,Int}, λμ::NTuple{2,Int}) where {T, B}
   if size(Z) ≠ sum.(dims)
       throw(DimensionMismatch())
   end

   bs = BandedBlockBandedSizes(dims..., lu..., λμ...)
   d_bs = bs.data_block_sizes
   blocks = fill!(B(undef, d_bs), zero(T))
   _BandedBlockBandedMatrix(blocks, bs)
end

function BandedBlockBandedMatrix{T, B}(E::Eye, dims::NTuple{2,AbstractVector{Int}},
                                       lu::NTuple{2,Int}, λμ::NTuple{2,Int}) where {T, B}
    if size(E) ≠ sum.(dims)
        throw(DimensionMismatch())
    end
    ret = BandedBlockBandedMatrix{T, B}(Zeros{T}(E), dims, lu, λμ)
    ret[diagind(ret)] .= one(T)
    ret
end

function BandedBlockBandedMatrix{T, B}(A::UniformScaling, dims::NTuple{2, AbstractVector{Int}},
                                       lu::NTuple{2,Int}, λμ::NTuple{2,Int}) where {T, B}
    ret = BandedBlockBandedMatrix{T, B}(Zeros{T}(sum.(dims)), dims, lu, λμ)
    ret[diagind(ret)] .= convert(T, A.λ)
    ret
end

BandedBlockBandedMatrix{T}(m::Union{Zeros, Eye, UniformScaling},
                           dims::NTuple{2,AbstractVector{Int}},
                           lu::NTuple{2,Int}, λμ::NTuple{2,Int}) where T =
  DefaultBandedBlockBandedMatrix{T}(m, dims, lu, λμ)

function BandedBlockBandedMatrix{T, B}(Z::Zeros, block_sizes::BandedBlockBandedSizes) where {T, B}
   if size(Z) ≠ size(block_sizes)
       throw(DimensionMismatch())
   end

   d_bs = block_sizes.data_block_sizes
   data = fill!(B(undef, d_bs), zero(T))
   _BandedBlockBandedMatrix(data, block_sizes)
end

BandedBlockBandedMatrix{T}(Z::Zeros, block_sizes::BandedBlockBandedSizes,
                           lu::NTuple{2,Int}, λμ::NTuple{2,Int}) where T =
  DefaultBandedBlockBandedMatrix{T}(Z, block_sizes)

function BandedBlockBandedMatrix{T, Blocks}(A::AbstractMatrix, block_sizes::BandedBlockBandedSizes) where {T, Blocks}
    ret = BandedBlockBandedMatrix{T, Blocks}(Zeros{T}(size(A)), block_sizes)
    L,M = block_sizes.λ, block_sizes.μ
    for J = Block.(1:nblocks(ret, 2)), K = blockcolrange(ret, Int(J))
        kr, jr = globalrange(block_sizes, (Int(K), Int(J)))

        # We have the correct block - now we need to only add the entries from
        # the correct bands
        B = view(A, kr, jr)
        R = view(ret, K, J)
        rows, cols = size(B)
        for λ = 1:L, j = 1:min(rows-λ, cols)
            view(R, j+λ, j) .= view(B, j+λ, j)
        end
        for i = 1:min(rows, cols)
            view(R, i, i) .= view(B, i, i)
        end
        for μ = 1:M, k = 1:min(rows, cols-μ)
            view(R, k, k+μ) .= view(B, k, k+μ)
        end
    end
    ret
end

BandedBlockBandedMatrix{T}(A::AbstractMatrix, block_sizes::BandedBlockBandedSizes) where T =
    DefaultBandedBlockBandedMatrix{T}(A, block_sizes)

BandedBlockBandedMatrix{T}(A::AbstractMatrix, block_sizes::BandedBlockBandedSizes,
                           lu::NTuple{2,Int}, λμ::NTuple{2,Int}) where T =
  DefaultBandedBlockBandedMatrix{T}(A, block_sizes, lu, λμ)

BandedBlockBandedMatrix(A::Union{AbstractMatrix,UniformScaling},
                        block_sizes::BandedBlockBandedSizes) =
    BandedBlockBandedMatrix{eltype(A)}(A, block_sizes)

BandedBlockBandedMatrix(A::Union{AbstractMatrix,UniformScaling},
                        block_sizes::AbstractBlockSizes, lu::NTuple{2,Int},
                        λμ::NTuple{2,Int}) =
                        BandedBlockBandedMatrix(A, BandedBlockBandedSizes(block_sizes, lu..., λμ...))

BandedBlockBandedMatrix{T}(A::Union{AbstractMatrix,UniformScaling},
                           dims::NTuple{2,AbstractVector{Int}}, lu::NTuple{2,Int},
                           λμ::NTuple{2,Int}) where T =
    BandedBlockBandedMatrix{T}(A, BandedBlockBandedSizes(dims..., lu..., λμ...))

BandedBlockBandedMatrix{T,B}(A::Union{AbstractMatrix,UniformScaling},
                           dims::NTuple{2,AbstractVector{Int}}, lu::NTuple{2,Int},
                           λμ::NTuple{2,Int}) where {T,B} =
    BandedBlockBandedMatrix{T,B}(A, BandedBlockBandedSizes(dims..., lu..., λμ...))

BandedBlockBandedMatrix(A::Union{AbstractMatrix,UniformScaling},
                        dims::NTuple{2, AbstractVector{Int}}, lu::NTuple{2,Int},
                        λμ::NTuple{2,Int}) =
    BandedBlockBandedMatrix{eltype(A)}(A, dims, lu, λμ)

BandedBlockBandedMatrix(A::AbstractMatrix) =
    BandedBlockBandedMatrix(A, blocksizes(A), blockbandwidths(A), subblockbandwidths(A))

similar(A::BandedBlockBandedMatrix, T::Type=eltype(A), bs::BandedBlockBandedSizes=blocksizes(A)) =
      BandedBlockBandedMatrix{T}(undef, bs)

function sparse(A::BandedBlockBandedMatrix)
    i = Vector{Int}()
    j = Vector{Int}()
    z = Vector{eltype(A)}()
    for J = Block.(1:nblocks(A,2)), K = blockcolrange(A, J)
        B = view(A, K, J)
        ĩ = _banded_rowval(B)
        j̃ = _banded_colval(B)
        z̃ = _banded_nzval(B)
        ĩ .+= cumulsizes(A, 1, Int(K))-1
        j̃ .+= cumulsizes(A, 2, Int(J))-1
        append!(i, ĩ)
        append!(j, j̃)
        append!(z, z̃)
    end
    sparse(i, j, z, size(A)...)
end

################################
# BandedBlockBandedMatrix Interface #
################################

MemoryLayout(::Type{<:BandedBlockBandedMatrix}) = BandedBlockBandedColumnMajor()
BroadcastStyle(::Type{<:BandedBlockBandedMatrix}) = BandedBlockBandedStyle()

isbandedblockbanded(_) = false
isbandedblockbanded(::BandedBlockBandedMatrix) = true

blockbandwidths(A::BandedBlockBandedMatrix) = (A.l, A.u)

"""
    subblockbandwidths(A)

returns the sub-block bandwidths of `A`, where `A` is a banded-block-banded
matrix. In other words, `A[Block(K,J)]` will return a `BandedMatrix` with
bandwidths given by `subblockbandwidths(A)`.
"""
subblockbandwidths(A::BandedBlockBandedMatrix) = (A.λ, A.μ)

# default is to use bandwidth
subblockbandwidths(A::AbstractMatrix) = bandwidths(A)

"""
    subblockbandwidth(A, i)

returns the sub-block lower (`i == 1`) or upper (`i == 2`) bandwidth of `A`,
where `A` is a banded-block-banded matrix. In other words, `A[Block(K,J)]` will
return a `BandedMatrix` with the returned lower/upper bandwidth.
"""
subblockbandwidth(A::AbstractArray, k::Integer) = subblockbandwidths(A)[k]

isdiag(A::BandedBlockBandedMatrix) = A.λ == A.μ == A.l == A.u

################################
# AbstractBlockArray Interface #
################################

@inline blocksizes(block_array::BandedBlockBandedMatrix) = block_array.block_sizes

zeroblock(A::BandedBlockBandedMatrix, K::Int, J::Int) =
    BandedMatrix(Zeros{eltype(A)}(blocksize(A, (K, J))), (A.λ, A.μ))

@inline function getblock(A::BandedBlockBandedMatrix, K::Int, J::Int)
    @boundscheck blockcheckbounds(A, K, J)
    if -A.l ≤ J - K ≤ A.u
        convert(BandedMatrix, view(A, Block(K, J)))
    else
        zeroblock(A, K, J)
    end
end

# @inline function Base.getindex(block_arr::BlockArray{T,N}, blockindex::BlockIndex{N}) where {T,N}
#     @boundscheck checkbounds(block_arr.blocks, blockindex.I...)
#     @inbounds block = block_arr.blocks[blockindex.I...]
#     @boundscheck checkbounds(block, blockindex.α...)
#     @inbounds v = block[blockindex.α...]
#     return v
# end

###########################
# AbstractArray Interface #
###########################

# @inline function Base.similar(block_array::BandedBlockBandedMatrix{T}, ::Type{T2}) where {T,N,T2}
#     BlockArray(similar(block_array.blocks, Array{T2, N}), copy(block_array.block_sizes))
# end

Base.size(arr::BandedBlockBandedMatrix) =
    @inbounds return size(arr.block_sizes)

@inline function getindex(A::BandedBlockBandedMatrix{T}, i::Int, j::Int) where T
    @boundscheck checkbounds(A, i, j)
    bi = global2blockindex(A.block_sizes, (i, j))
    @inbounds v = view(A, Block(bi.I))[bi.α...]
    return v
end

@inline function setindex!(A::BandedBlockBandedMatrix{T}, v, i::Int, j::Int) where T
    @boundscheck checkbounds(A, i, j)
    bi = global2blockindex(A.block_sizes, (i, j))
    if -A.l ≤ bi.I[2] - bi.I[1] ≤ A.u
        V = view(A, Block(bi.I))
        @inbounds V[bi.α...] = convert(T, v)::T
    elseif !iszero(v)
        throw(BandError(A))
    end
    return v
end

## structured matrix methods ##
function Base.replace_in_print_matrix(A::BandedBlockBandedMatrix, i::Integer, j::Integer, s::AbstractString)
    bi = global2blockindex(A.block_sizes, (i, j))
    I,J = bi.I
    i,j = bi.α
    -A.l ≤ J-I ≤ A.u && -A.λ ≤ j-i ≤ A.μ ? s : Base.replace_with_centered_mark(s)
end

############
# Indexing #
############

# function _check_setblock!(block_arr::BlockArray{T, N}, v, block::NTuple{N, Int}) where {T,N}
#     for i in 1:N
#         if size(v, i) != blocksize(block_arr.block_sizes, (i, block[i]))
#             throw(DimensionMismatch(string("tried to assign $(size(v)) array to ", blocksize(block_arr, block), " block")))
#         end
#     end
# end
#
#
@inline function setblock!(A::BandedBlockBandedMatrix, v, K::Int, J::Int)
    @boundscheck blockcheckbounds(A, K, J)

    @boundscheck (bandwidth(v, 1) > A.λ || bandwidth(v, 2) > A.μ) && throw(BandError())
    V = view(A, Block(K), Block(J))
    V .= v
    return A
end
#
# @propagate_inbounds function Base.setindex!(block_array::BlockArray{T, N}, v, block_index::BlockIndex{N}) where {T,N}
#     getblock(block_array, block_index.I...)[block_index.α...] = v
# end

########
# Misc #
########

# @generated function Base.Array(block_array::BlockArray{T, N, R}) where {T,N,R}
#     # TODO: This will fail for empty block array
#     return quote
#         block_sizes = block_array.block_sizes
#         arr = similar(block_array.blocks[1], size(block_array)...)
#         @nloops $N i i->(1:nblocks(block_sizes, i)) begin
#             block_index = @ntuple $N i
#             indices = globalrange(block_sizes, block_index)
#             arr[indices...] = getblock(block_array, block_index...)
#         end
#
#         return arr
#     end
# end
#
# @generated function Base.copyto!(block_array::BlockArray{T, N, R}, arr::R) where {T,N,R <: AbstractArray}
#     return quote
#         block_sizes = block_array.block_sizes
#
#         @nloops $N i i->(1:nblocks(block_sizes, i)) begin
#             block_index = @ntuple $N i
#             indices = globalrange(block_sizes, block_index)
#             copyto!(getblock(block_array, block_index...), arr[indices...])
#         end
#
#         return block_array
#     end
# end
#
# function Base.fill!(block_array::BlockArray, v)
#     for block in block_array.blocks
#         fill!(block, v)
#     end
# end

##################
# BandedBlockBandedBlock
#
#   views of the blocks satisfy the BandedMatrix interface, and are memory-compatible
#   with BLASBandedMatrix.
##################

const BandedBlockBandedBlock{T, BLOCKS} = SubArray{T,2,BandedBlockBandedMatrix{T, BLOCKS},Tuple{BlockSlice1,BlockSlice1},false}

isbanded(::BandedBlockBandedBlock) = true
MemoryLayout(::BandedBlockBandedBlock) = BandedColumnMajor()
BroadcastStyle(::Type{<: BandedBlockBandedBlock}) = BandedStyle()

function inblockbands(V::BandedBlockBandedBlock)
    A = parent(V)
    K_sl, J_sl = parentindices(V)
    K, J = K_sl.block, J_sl.block
    -A.l ≤ Int(J-K) ≤ A.u
end

######################################
# BandedMatrix interface  for Blocks #
######################################
@inline function bandwidths(V::BandedBlockBandedBlock)
    inblockbands(V) && return subblockbandwidths(parent(V))
    (-720,-720)
end

# gives the columns of parent(V).data that encode the block
blocks(V::BandedBlockBandedBlock)::Tuple{Int,Int} = Int(first(parentindices(V)).block),
                                                    Int(last(parentindices(V)).block)

function bandeddata(V::BandedBlockBandedBlock{T}) where T
    inblockbands(V) || return Array{T}(undef, 0, size(V,2))
    A = parent(V)
    u = A.u
    K_sl, J_sl = parentindices(V)
    K, J = K_sl.block, J_sl.block
    view(A.data, u + K - J + 1, J)
end

@inline function inbands_getindex(V::BandedBlockBandedBlock, k::Int, j::Int)
    A = parent(V)
    banded_getindex(bandeddata(V), A.λ, A.μ, k, j)
end

@inline function inbands_setindex!(V::BandedBlockBandedBlock, v, k::Int, j::Int)
    A = parent(V)
    banded_setindex!(bandeddata(V), A.λ, A.μ, v, k, j)
end

# we need getindex as we use block-wise get index above
@propagate_inbounds function getindex(V::BandedBlockBandedBlock, k::Int, j::Int)
    @boundscheck checkbounds(V, k, j)
    A = parent(V)
    K,J = blocks(V)
    if -A.l ≤ J-K ≤ A.u && -A.λ ≤ j-k ≤ A.μ
        inbands_getindex(V, k, j)
    else
        zero(eltype(V))
    end
end

@propagate_inbounds function setindex!(V::BandedBlockBandedBlock, v, k::Int, j::Int)
    @boundscheck checkbounds(V, k, j)
    A = parent(V)
    K,J = blocks(V)
    if -A.l ≤ J-K ≤ A.u && -A.λ ≤ j-k ≤ A.μ
        inbands_setindex!(V, v, k, j)
    elseif iszero(v) # allow setindex for 0 datya
        v
    else
        throw(BandError(V, J-K))
    end
end

BLAS.axpy!(a::T, A::BandedBlockBandedMatrix{T}, B::BandedBlockBandedMatrix{T}) where T =
    B .= a .* A .+ B

lmul!(x::Number, A::BandedBlockBandedMatrix) = (lmul!(x, A.data); A)
rmul!(A::BandedBlockBandedMatrix, x::Number) = (rmul!(A.data, x); A)
*(x::Number, A::BandedBlockBandedMatrix) = _BandedBlockBandedMatrix(x*A.data, A.block_sizes)
*(A::BandedBlockBandedMatrix, x::Number) = _BandedBlockBandedMatrix(A.data*x, A.block_sizes)