export MultiConductorValue, MultiConductorVector, MultiConductorMatrix, conductors

# "a data structure for working with multiconductor datasets"
if VERSION < v"0.7.0-"
    abstract type MultiConductorValue{T,N} end
else
    abstract type MultiConductorValue{T,N} <: AbstractArray{T,N} end
end

"a data structure for working with multiconductor datasets"
mutable struct MultiConductorVector{T} <: MultiConductorValue{T,1}
    values::Vector{T}
end

MultiConductorVector(value::T, conductors::Int) where T = MultiConductorVector([value for i in 1:conductors])
Base.map(f, a::MultiConductorVector{T}) where T = MultiConductorVector{T}(map(f, a.values))
Base.map(f, a::MultiConductorVector{T}, b::MultiConductorVector{T}) where T = MultiConductorVector{T}(map(f, a.values, b.values))
conductors(mcv::MultiConductorVector) = length(mcv.values)

MultiConductorVector(value::Array{T,2}) where T = MultiConductorMatrix{T}(value)

""
function Base.setindex!(mcv::MultiConductorVector{T}, v::T, i::Int) where T
    mcv.values[i] = v
end

""
mutable struct MultiConductorMatrix{T} <: MultiConductorValue{T,2}
    values::Matrix{T}
end

MultiConductorMatrix(value::T, conductors::Int) where T = MultiConductorMatrix(value*Matrix{Float64}(LinearAlgebra.I, conductors, conductors))
Base.map(f, a::MultiConductorMatrix{T}) where T = MultiConductorMatrix{T}(map(f, a.values))
Base.map(f, a::MultiConductorMatrix{T}, b::MultiConductorMatrix{T}) where T = MultiConductorMatrix{T}(map(f, a.values, b.values))
conductors(mcv::MultiConductorMatrix) = size(mcv.values, 1)

""
function Base.setindex!(mcv::MultiConductorMatrix{T}, v::T, i::Int, j::Int) where T
    mcv.values[i,j] = v
end

if VERSION < v"0.7.0-"
    Base.start(mcv::MultiConductorValue) = start(mcv.values)
    Base.next(mcv::MultiConductorValue, state) = next(mcv.values, state)
    Base.done(mcv::MultiConductorValue, state) = done(mcv.values, state)
else
    iterate(mcv::MultiConductorValue, kwargs...) = iterate(mcv.values, kwargs...)
end

Base.length(mcv::MultiConductorValue) = length(mcv.values)
Base.size(mcv::MultiConductorValue, a...) = size(mcv.values, a...)
Base.getindex(mcv::MultiConductorValue, args...) = mcv.values[args...]

Base.show(io::IO, mcv::MultiConductorValue) = Base.show(io, mcv.values)

Base.broadcast(f::Any, a::Any, b::MultiConductorValue) = broadcast(f, a, b.values)
Base.broadcast(f::Any, a::MultiConductorValue, b::Any) = broadcast(f, a.values, b)
Base.broadcast(f::Any, a::MultiConductorValue, b::MultiConductorValue) = broadcast(f, a.values, b.values)

# Broadcast implementation for Julia v0.7
if VERSION > v"0.7.0-"
    Base.BroadcastStyle(::Type{<:MultiConductorVector}) = Broadcast.ArrayStyle{MultiConductorVector}()
    Base.BroadcastStyle(::Type{<:MultiConductorMatrix}) = Broadcast.ArrayStyle{MultiConductorMatrix}()

    function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MultiConductorVector}}, ::Type{ElType}) where ElType
        A = find_mcv(bc)
        return MultiConductorVector(similar(Array{ElType}, axes(bc)))
    end

    "`A = find_mcv(As)` returns the first MultiConductorVector among the arguments."
    find_mcv(bc::Base.Broadcast.Broadcasted) = find_mcv(bc.args)
    find_mcv(args::Tuple) = find_mcv(find_mcv(args[1]), Base.tail(args))
    find_mcv(x) = x
    find_mcv(a::MultiConductorVector, rest) = a
    find_mcv(::Any, rest) = find_mcv(rest)

    function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{MultiConductorMatrix}}, ::Type{ElType}) where ElType
        A = find_mcm(bc)
        return MultiConductorMatrix(similar(Array{ElType}, axes(bc)))
    end

    "`A = find_mcm(As)` returns the first MultiConductorMatrix among the arguments."
    find_mcm(bc::Base.Broadcast.Broadcasted) = find_mcm(bc.args)
    find_mcm(args::Tuple) = find_mcm(find_mcm(args[1]), Base.tail(args))
    find_mcm(x) = x
    find_mcm(a::MultiConductorMatrix, rest) = a
    find_mcm(::Any, rest) = find_mcm(rest)
end

# Vectors
Base.:+(a::MultiConductorVector) = MultiConductorVector(+(a.values))
Base.:+(a::MultiConductorVector, b::Union{Array,Number}) = MultiConductorVector(+(a.values, b))
Base.:+(a::Union{Array,Number}, b::MultiConductorVector) = MultiConductorVector(+(a, b.values))
Base.:+(a::MultiConductorVector, b::MultiConductorVector) = MultiConductorVector(+(a.values, b.values))

Base.:-(a::MultiConductorVector) = MultiConductorVector(-(a.values))
Base.:-(a::MultiConductorVector, b::Union{Array,Number}) = MultiConductorVector(-(a.values, b))
Base.:-(a::Union{Array,Number}, b::MultiConductorVector) = MultiConductorVector(-(a, b.values))
Base.:-(a::MultiConductorVector, b::MultiConductorVector) = MultiConductorVector(-(a.values, b.values))

Base.:*(a::Number, b::MultiConductorVector) = MultiConductorVector(*(a, b.values))
Base.:*(a::MultiConductorVector, b::Number) = MultiConductorVector(*(a.values, b))
Base.:*(a::Array, b::MultiConductorVector) = MultiConductorVector(Base.broadcast(*, a, b.values))
Base.:*(a::MultiConductorVector, b::Array) = MultiConductorVector(Base.broadcast(*, a.values, b))
Base.:*(a::MultiConductorVector, b::MultiConductorVector) = MultiConductorVector(Base.broadcast(*, a.values, b.values))

Base.:/(a::MultiConductorVector, b::Number) = MultiConductorVector(/(a.values, b))
Base.:/(a::Union{Array,Number}, b::MultiConductorVector) = MultiConductorVector(Base.broadcast(/, a, b.values))
Base.:/(a::MultiConductorVector, b::MultiConductorVector) = MultiConductorVector(Base.broadcast(/, a.values, b.values))

if VERSION < v"0.7.0-"
    Base.:*(a::MultiConductorVector, b::RowVector) = MultiConductorMatrix(Base.broadcast(*, a.values, b))
    Base.:*(a::RowVector, b::MultiConductorVector) = MultiConductorMatrix(Base.broadcast(*, a, b.values))
else
    Base.:*(a::MultiConductorVector, b::LinearAlgebra.Adjoint) = MultiConductorMatrix(Base.broadcast(*, a.values, b))
    Base.:*(a::LinearAlgebra.Adjoint, b::MultiConductorVector) = MultiConductorMatrix(Base.broadcast(*, a, b.values))
end

# Matrices
Base.:+(a::MultiConductorMatrix) = MultiConductorMatrix(+(a.values))
Base.:+(a::MultiConductorMatrix, b::Union{Array,Number}) = MultiConductorMatrix(+(a.values, b))
Base.:+(a::Union{Array,Number}, b::MultiConductorMatrix) = MultiConductorMatrix(+(a, b.values))
Base.:+(a::MultiConductorMatrix, b::MultiConductorMatrix) = MultiConductorMatrix(+(a.values, b.values))

Base.:-(a::MultiConductorMatrix) = MultiConductorMatrix(-(a.values))
Base.:-(a::MultiConductorMatrix, b::Union{Array,Number}) = MultiConductorMatrix(-(a.values, b))
Base.:-(a::Union{Array,Number}, b::MultiConductorMatrix) = MultiConductorMatrix(-(a, b.values))
Base.:-(a::MultiConductorMatrix, b::MultiConductorMatrix) = MultiConductorMatrix(-(a.values, b.values))

Base.:*(a::MultiConductorMatrix, b::Number) = MultiConductorMatrix(*(a.values, b))
Base.:*(a::Number, b::MultiConductorMatrix) = MultiConductorMatrix(*(a, b.values))
Base.:*(a::MultiConductorMatrix, b::Array) = MultiConductorMatrix(*(a.values, b))
Base.:*(a::Array, b::MultiConductorMatrix) = MultiConductorMatrix(*(a, b.values))
Base.:*(a::MultiConductorMatrix, b::MultiConductorMatrix) = MultiConductorMatrix(*(a.values, b.values))

Base.:/(a::MultiConductorMatrix, b::Union{Array,Number}) = MultiConductorMatrix(/(a.values, b))
Base.:/(a::Union{Array,Number}, b::MultiConductorMatrix) = MultiConductorMatrix(/(a, b.values))
Base.:/(a::MultiConductorMatrix, b::MultiConductorMatrix) = MultiConductorMatrix(/(a.values, b.values))

Base.:*(a::MultiConductorMatrix, b::MultiConductorVector) = MultiConductorVector(*(a.values, b.values))

if VERSION < v"0.7.0-"
    Base.:/(a::MultiConductorMatrix, b::RowVector) = MultiConductorVector(squeeze(/(a.values, b), 2))
else
    Base.:/(a::MultiConductorMatrix, b::LinearAlgebra.Adjoint) = MultiConductorVector(squeeze(/(a.values, b), 2))
end

Base.:^(a::MultiConductorVector, b::Complex) = MultiConductorVector(Base.broadcast(^, a.values, b))
Base.:^(a::MultiConductorVector, b::Integer) = MultiConductorVector(Base.broadcast(^, a.values, b))
Base.:^(a::MultiConductorVector, b::AbstractFloat) = MultiConductorVector(Base.broadcast(^, a.values, b))
Base.:^(a::MultiConductorMatrix, b::Complex) = MultiConductorMatrix(a.values ^ b)
Base.:^(a::MultiConductorMatrix, b::Integer) = MultiConductorMatrix(a.values ^ b)
Base.:^(a::MultiConductorMatrix, b::AbstractFloat) = MultiConductorMatrix(a.values ^ b)

LinearAlgebra.inv(a::MultiConductorMatrix) = MultiConductorMatrix(inv(a.values))
LinearAlgebra.pinv(a::MultiConductorMatrix) = MultiConductorMatrix(LinearAlgebra.pinv(a.values))

Base.real(a::MultiConductorVector) = MultiConductorVector(real(a.values))
Base.real(a::MultiConductorMatrix) = MultiConductorMatrix(real(a.values))
Base.imag(a::MultiConductorVector) = MultiConductorVector(imag(a.values))
Base.imag(a::MultiConductorMatrix) = MultiConductorMatrix(imag(a.values))

LinearAlgebra.transpose(a::MultiConductorVector) = a.values'
LinearAlgebra.transpose(a::MultiConductorMatrix) = MultiConductorMatrix(a.values')

LinearAlgebra.diag(a::MultiConductorMatrix) = MultiConductorVector(diag(a.values))
LinearAlgebra.diagm(p::Pair{<:Integer, MultiConductorVector{S}}) where S = MultiConductorMatrix(diagm(p.first => p.second.values))

if VERSION <= v"0.7.0-"
    LinearAlgebra.diagm(a::MultiConductorVector{S}) where S = LinearAlgebra.diagm(0 => a)
end

Base.rad2deg(a::MultiConductorVector) = MultiConductorVector(map(rad2deg, a.values))
Base.rad2deg(a::MultiConductorMatrix) = MultiConductorMatrix(map(rad2deg, a.values))

Base.deg2rad(a::MultiConductorVector) = MultiConductorVector(map(deg2rad, a.values))
Base.deg2rad(a::MultiConductorMatrix) = MultiConductorMatrix(map(deg2rad, a.values))

JSON.lower(mcv::MultiConductorValue) = mcv.values

"converts a MultiConductorValue value to a string in summary"
function _value2string(mcv::MultiConductorValue, float_precision::Int)
    a = join([_value2string(v, float_precision) for v in mcv.values], ", ")
    return "[$(a)]"
end

""
function Base.isapprox(a::MultiConductorValue, b::MultiConductorValue; kwargs...)
    if length(a) == length(b)
        return all( isapprox(a[i], b[i]; kwargs...) for i in 1:length(a))
    end
    return false
end

getmcv(value::Any, conductor::Int) = value
getmcv(value::Any, conductor_i::Int, conductor_j::Int) = value
getmcv(value::MultiConductorVector, conductor::Int) = value[conductor]
getmcv(value::MultiConductorMatrix{T}, conductor::Int) where T = MultiConductorVector{T}(value[conductor])
getmcv(value::MultiConductorMatrix, conductor_i::Int, conductor_j::Int) = value[conductor_i, conductor_j]


"""
    generate_data(params::Dict)

Generates DFT/reference data based on a particular potential.
"""
function generate_data(model::String, params::Dict)
    # Load the potential model (E.g. GaN)
    p_model = Symbol(params[model]["potential"])
    p = @eval $p_model($(params[model]))
    # Load parameters
    fit_forces = params["global"]["fit_forces"]
    a = params["global"]["no_train_atomic_conf"]
    b = params["global"]["no_atomic_conf"]
    positions_per_conf = params["global"]["positions_per_conf"]
    rcutfac = params["global"]["rcutfac"]
    # Calculate DFT training data
    training_data = generate_data_aux(p, positions_per_conf, 1, a, rcutfac, fit_forces)
    # Calculate DFT validation data
    validation_data = generate_data_aux(p, positions_per_conf, a+1, b, rcutfac, fit_forces)

    return training_data, validation_data
end

"""
    generate_data_aux(p::Potential, positions_per_conf::Vector,
                      a::Int64, b::Int64, rcutfac::Float64, fit_forces::Bool)

Auxiliar function. See `generate_data`.
"""
function generate_data_aux(p::Potential, positions_per_conf::Vector,
                           a::Int64, b::Int64, rcutfac::Float64, fit_forces::Bool)
    potentials  = [potential_energy(p, positions_per_conf[j], rcutfac) for j = a:b]
    lin_forces = fit_forces ? linearize([forces(p, positions_per_conf[j], rcutfac)
                                         for j = a:b]) : Vector{Float64}()
    return [potentials; lin_forces]
end

"""
    linearize(data::Vector{Force})

Linearizes a vector of forces.
E.g. [Force(1.0,2.0,3.0), Force(4.0,5.0,6.0)] => [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
"""
function linearize(data::Vector{Vector{Force}})
    lin_data = Vector{Float64}()
    for i = 1:length(data), j = 1:length(data[i]), k = 1:3
        push!(lin_data, data[i][j][k])
    end
    return lin_data
end

function get_dft_data(params::Dict)
    path = params["global"]["path"]
    fit_forces = params["global"]["fit_forces"]
    n = params["global"]["no_train_atomic_conf"]
    m = params["global"]["no_atoms_per_conf"]

    dirs = readdir(joinpath(path, "DATA"))
    dft_energies = Vector{Float64}()
    for d in dirs
        open(joinpath(path, "DATA", d, "ENERGY") , "r") do file
            e = parse(Float64, readline(file))
            push!(dft_energies, e)
        end
    end
    training_data = dft_energies[1:n]
    validation_data = dft_energies[n+1:end]

    if fit_forces
        dft_forces = Vector{Float64}()
        for d in dirs
            open(joinpath(path, "DATA", d, "FORCES") , "r") do file
                while !eof(f)
                    f = parse.(Float64, split(readline(file)))
                    push!(dft_forces, f[1])
                    push!(dft_forces, f[2])
                    push!(dft_forces, f[3])
                end
            end
        end
        training_data = [training_data; dft_forces[1:3*n*m]]
        validation_data = [training_data; dft_forces[3*n*m+1:end]]
    end
    return training_data, validation_data
end


const FluxOptimizer = Union{
    Descent, Momentum, Nesterov, RMSProp,
    ADAM, AdaMax, ADAGrad, ADADelta, RADAM,
    AMSGrad, NADAM
}

abstract type OptimParam end
struct LR <: OptimParam end
struct Mom <: OptimParam end

optimparamfield(::O, op::Type{P}) where {O, P<:OptimParam} = error(
    "No parameter $(P) on optimizer $(O)."
)
optimparamfield(::FluxOptimizer, ::Type{LR}) = :eta

function setoptimparam!(opt, optimparam::Type{<:OptimParam}, x)
    setfield!(opt, optimparamfield(opt, optimparam), x)
end
getoptimparam(opt, optimparam::Type{<:OptimParam}) = getfield(opt, optimparamfield(opt, optimparam))

# TODO: add momentum, weight decay and other...
getoptimparams(opt) = Dict(LR => getoptimparam(opt, LR))


using IonLinearAlgebra
using Test

@testset "IonLinearAlgebra.jl" begin
    # Write your own tests here.
end


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

# What is included in this file:
# - the `fmi2ComponentState`--enum which mirrors the internal FMU state (state-machine, not the system state)
# - the `FMU2ComponentEnvironment`- and `FMU2Component`-struct
# - the `FMU2`-struct
# - string/enum-converters for FMI-attribute-structs (e.g. `fmi2StatusToString`, ...)

# this is a custom type to catch the internal mode of the component
@enum fmi2ComponentState begin
    fmi2ComponentStateInstantiated       # after instantiation
    fmi2ComponentStateInitializationMode # after finishing initialization
    fmi2ComponentStateEventMode
    fmi2ComponentStateContinuousTimeMode
    fmi2ComponentStateTerminated
    fmi2ComponentStateError
    fmi2ComponentStateFatal
end

"""
ToDo.
"""
mutable struct FMU2ComponentEnvironment
    logStatusOK::Bool
    logStatusWarning::Bool
    logStatusDiscard::Bool
    logStatusError::Bool
    logStatusFatal::Bool
    logStatusPending::Bool

    function FMU2ComponentEnvironment()
        inst = new()
        inst.logStatusOK = true
        inst.logStatusWarning = true
        inst.logStatusDiscard = true
        inst.logStatusError = true
        inst.logStatusFatal = true
        inst.logStatusPending = true
        return inst
    end
end

"""
The mutable struct represents an instantiated instance of an FMU in the FMI 2.0.2 Standard.
"""
mutable struct FMU2Component
    compAddr::Ptr{Nothing}
    fmu # ::FMU2
    state::fmi2ComponentState
    componentEnvironment::FMU2ComponentEnvironment

    loggingOn::Bool
    callbackFunctions::fmi2CallbackFunctions
    instanceName::String
    continuousStatesChanged::fmi2Boolean
    eventInfo::fmi2EventInfo

    t::fmi2Real             # the system time
    t_offset::fmi2Real      # time offset between simulation environment and FMU
    x::Union{Array{fmi2Real, 1}, Nothing}   # the system states (or sometimes u)
    ẋ::Union{Array{fmi2Real, 1}, Nothing}   # the system state derivative (or sometimes u̇)
    ẍ::Union{Array{fmi2Real, 1}, Nothing}   # the system state second derivative
    #u::Array{fmi2Real, 1}   # the system inputs
    #y::Array{fmi2Real, 1}   # the system outputs
    #p::Array{fmi2Real, 1}   # the system parameters
    z::Array{fmi2Real, 1}   # the system event indicators
    z_prev::Union{Nothing, Array{fmi2Real, 1}}   # the last system event indicators

    realValues::Dict

    x_vrs::Array{fmi2ValueReference, 1}   # the system state value references
    ẋ_vrs::Array{fmi2ValueReference, 1}   # the system state derivative value references
    u_vrs::Array{fmi2ValueReference, 1}   # the system input value references
    y_vrs::Array{fmi2ValueReference, 1}   # the system output value references
    p_vrs::Array{fmi2ValueReference, 1}   # the system parameter value references

    # deprecated
    jac_ẋy_x::Matrix{fmi2Real}
    jac_ẋy_u::Matrix{fmi2Real}
    jac_x::Array{fmi2Real}
    jac_u::Union{Array{fmi2Real}, Nothing}
    jac_t::fmi2Real

    # linearization jacobians
    A::Union{Matrix{fmi2Real}, Nothing}
    B::Union{Matrix{fmi2Real}, Nothing}
    C::Union{Matrix{fmi2Real}, Nothing}
    D::Union{Matrix{fmi2Real}, Nothing}

    jacobianUpdate!             # function for a custom jacobian constructor (optimization)
    skipNextDoStep::Bool    # allows skipping the next `fmi2DoStep` like it is not called
    senseFunc::Symbol       # :auto, :full, :sample, :directionalDerivatives, :adjointDerivatives

    # constructor

    function FMU2Component()
        inst = new()
        inst.state = fmi2ComponentStateInstantiated
        inst.t = -Inf
        inst.t_offset = 0.0
        inst.eventInfo = fmi2EventInfo()

        inst.senseFunc = :auto

        inst.x = nothing
        inst.ẋ = nothing
        inst.ẍ = nothing
        inst.z_prev = nothing

        inst.realValues = Dict()
        inst.x_vrs = Array{fmi2ValueReference, 1}()
        inst.ẋ_vrs = Array{fmi2ValueReference, 1}()
        inst.u_vrs = Array{fmi2ValueReference, 1}()
        inst.y_vrs = Array{fmi2ValueReference, 1}()
        inst.p_vrs = Array{fmi2ValueReference, 1}()

        # initialize further variables
        inst.jac_x = Array{fmi2Real, 1}()
        inst.jac_u = nothing
        inst.jac_t = -1.0
        inst.jac_ẋy_x = zeros(fmi2Real, 0, 0)
        inst.jac_ẋy_u = zeros(fmi2Real, 0, 0)
        inst.skipNextDoStep = false

        inst.A = nothing
        inst.B = nothing
        inst.C = nothing
        inst.D = nothing

        return inst
    end

    function FMU2Component(compAddr, fmu)
        inst = FMU2Component()
        inst.compAddr = compAddr
        inst.fmu = fmu
        return inst
    end
end

"""
Overload the Base.show() function for custom printing of the FMU2Component.
"""
Base.show(io::IO, c::FMU2Component) = print(io,
"FMU:            $(c.fmu.modelDescription.modelName)
InstanceName:   $(isdefined(c, :instanceName) ? c.instanceName : "[not defined]")
Address:        $(c.compAddr)
State:          $(c.state)
Logging:        $(c.loggingOn)
FMU time:       $(c.t)
FMU states:     $(c.x)"
)

"""
A mutable struct representing the excution configuration of a FMU.
For FMUs that have issues with calls like `fmi2Reset` or `fmi2FreeInstance`, this is pretty useful.
"""
mutable struct FMU2ExecutionConfiguration
    terminate::Bool     # call fmi2Terminate before every training step / simulation
    reset::Bool         # call fmi2Reset before every training step / simulation
    setup::Bool         # call setup functions before every training step / simulation
    instantiate::Bool   # call fmi2Instantiate before every training step / simulation
    freeInstance::Bool  # call fmi2FreeInstance after every training step / simulation

    loggingOn::Bool
    externalCallbacks::Bool

    handleStateEvents::Bool                 # handle state events during simulation/training
    handleTimeEvents::Bool                  # handle time events during simulation/training

    assertOnError::Bool                     # wheter an exception is thrown if a fmi2XXX-command fails (>= fmi2StatusError)
    assertOnWarning::Bool                   # wheter an exception is thrown if a fmi2XXX-command warns (>= fmi2StatusWarning)

    autoTimeShift::Bool                     # wheter to shift all time-related functions for simulation intervals not starting at 0.0

    sensealg                                # algorithm for sensitivity estimation over solve call
    useComponentShadow::Bool                # whether FMU outputs/derivatives/jacobians should be cached for frule/rrule (useful for ForwardDiff)
    rootSearchInterpolationPoints::UInt     # number of root search interpolation points
    inPlace::Bool                           # whether faster in-place-fx should be used
    useVectorCallbacks::Bool                # whether to vector (faster) or scalar (slower) callbacks

    maxNewDiscreteStateCalls::UInt          # max calls for fmi2NewDiscreteStates before throwing an exception

    function FMU2ExecutionConfiguration()
        inst = new()

        inst.terminate = true
        inst.reset = true
        inst.setup = true
        inst.instantiate = false
        inst.freeInstance = false

        inst.loggingOn = false
        inst.externalCallbacks = false

        inst.handleStateEvents = true
        inst.handleTimeEvents = true

        inst.assertOnError = false
        inst.assertOnWarning = false

        inst.autoTimeShift = true

        inst.sensealg = nothing # auto
        inst.useComponentShadow = false
        inst.rootSearchInterpolationPoints = 100
        inst.inPlace = true
        inst.useVectorCallbacks = true

        inst.maxNewDiscreteStateCalls = 100

        return inst
    end
end

# default for a "healthy" FMU - this is the fastetst
FMU_EXECUTION_CONFIGURATION_RESET = FMU2ExecutionConfiguration()
FMU_EXECUTION_CONFIGURATION_RESET.terminate = true
FMU_EXECUTION_CONFIGURATION_RESET.reset = true
FMU_EXECUTION_CONFIGURATION_RESET.instantiate = false
FMU_EXECUTION_CONFIGURATION_RESET.freeInstance = false

# if your FMU has a problem with "fmi2Reset" - this is default
FMU_EXECUTION_CONFIGURATION_NO_RESET = FMU2ExecutionConfiguration()
FMU_EXECUTION_CONFIGURATION_NO_RESET.terminate = false
FMU_EXECUTION_CONFIGURATION_NO_RESET.reset = false
FMU_EXECUTION_CONFIGURATION_NO_RESET.instantiate = true
FMU_EXECUTION_CONFIGURATION_NO_RESET.freeInstance = true

# if your FMU has a problem with "fmi2Reset" and "fmi2FreeInstance" - this is for weak FMUs (but slower)
FMU_EXECUTION_CONFIGURATION_NO_FREEING = FMU2ExecutionConfiguration()
FMU_EXECUTION_CONFIGURATION_NO_FREEING.terminate = false
FMU_EXECUTION_CONFIGURATION_NO_FREEING.reset = false
FMU_EXECUTION_CONFIGURATION_NO_FREEING.instantiate = true
FMU_EXECUTION_CONFIGURATION_NO_FREEING.freeInstance = false

"""
ToDo
"""
struct FMU2Event
    t::Union{Float32, Float64}
    indicator::UInt

    x_left::Union{Array{Float64, 1}, Array{Float32, 1}, Nothing}
    x_right::Union{Array{Float64, 1}, Array{Float32, 1}, Nothing}

    function FMU2Event(t::Union{Float32, Float64}, indicator::UInt = 0, x_left::Union{Array{Float64, 1}, Array{Float32, 1}, Nothing} = nothing, x_right::Union{Array{Float64, 1}, Array{Float32, 1}, Nothing} = nothing)
        inst = new(t, indicator, x_left, x_right)
        return inst
    end
end

"""
Overload the Base.show() function for custom printing of the FMU2.
"""
Base.show(io::IO, e::FMU2Event) = print(io, e.indicator == 0 ? "Time-Event @ $(round(e.t; digits=4))s" : "State-Event #$(e.indicator) @ $(round(e.t; digits=4))s")

"""
ToDo
"""
mutable struct FMU2Solution <: FMUSolution
    fmu                                             # FMU2
    success::Bool

    states                                          # ODESolution

    values
    valueReferences::Union{Array, Nothing}          # Array{fmi2ValueReference}

    events::Array{FMU2Event, 1}

    function FMU2Solution(fmu)
        inst = new()

        inst.fmu = fmu
        inst.success = false
        inst.states = nothing
        inst.values = nothing
        inst.valueReferences = nothing

        inst.events = []

        return inst
    end
end

"""
Overload the Base.show() function for custom printing of the FMU2.
"""
function Base.show(io::IO, sol::FMU2Solution)
    print(io, "Model name:\n\t$(sol.fmu.modelDescription.modelName)\nSuccess:\n\t$(sol.success)\n")

    if sol.states != nothing
        print(io, "States [$(length(sol.states))]:\n")
        if length(sol.states.u) > 10
            for i in 1:9
                print(io, "\t$(sol.states.t[i])\t$(sol.states.u[i])\n")
            end
            print(io, "\t...\n\t$(sol.states.t[end])\t$(sol.states.u[end])\n")
        else
            for i in 1:length(sol.states)
                print(io, "\t$(sol.states.t[i])\t$(sol.states.u[i])\n")
            end
        end
    end

    if sol.values != nothing
        print(io, "Values [$(length(sol.values.saveval))]:\n")
        if length(sol.values.saveval) > 10
            for i in 1:9
                print(io, "\t$(sol.values.t[i])\t$(sol.values.saveval[i])\n")
            end
            print(io, "\t...\n\t$(sol.values.t[end])\t$(sol.values.saveval[end])\n")
        else
            for i in 1:length(sol.values.saveval)
                print(io, "\t$(sol.values.t[i])\t$(sol.values.saveval[i])\n")
            end
        end
    end

    if sol.events != nothing
        print(io, "Events [$(length(sol.events))]:\n")
        if length(sol.events) > 10
            for i in 1:9
                print(io, "\t$(sol.events[i])\n")
            end
            print(io, "\t...\n\t$(sol.events[end])\n")
        else
            for i in 1:length(sol.events)
                print(io, "\t$(sol.events[i])\n")
            end
        end
    end

end

"""
The mutable struct representing a FMU and all it instantiated instances in the FMI 2.0.2 Standard.
Also contains the paths to the FMU and ZIP folder as well als all the FMI 2.0.2 function pointers.
"""
mutable struct FMU2 <: FMU
    modelName::String
    fmuResourceLocation::String

    modelDescription::fmi2ModelDescription

    type::fmi2Type
    components::Array # {fmi2Component}

    # c-functions
    cInstantiate::Ptr{Cvoid}
    cGetTypesPlatform::Ptr{Cvoid}
    cGetVersion::Ptr{Cvoid}
    cFreeInstance::Ptr{Cvoid}
    cSetDebugLogging::Ptr{Cvoid}
    cSetupExperiment::Ptr{Cvoid}
    cEnterInitializationMode::Ptr{Cvoid}
    cExitInitializationMode::Ptr{Cvoid}
    cTerminate::Ptr{Cvoid}
    cReset::Ptr{Cvoid}
    cGetReal::Ptr{Cvoid}
    cSetReal::Ptr{Cvoid}
    cGetInteger::Ptr{Cvoid}
    cSetInteger::Ptr{Cvoid}
    cGetBoolean::Ptr{Cvoid}
    cSetBoolean::Ptr{Cvoid}
    cGetString::Ptr{Cvoid}
    cSetString::Ptr{Cvoid}
    cGetFMUstate::Ptr{Cvoid}
    cSetFMUstate::Ptr{Cvoid}
    cFreeFMUstate::Ptr{Cvoid}
    cSerializedFMUstateSize::Ptr{Cvoid}
    cSerializeFMUstate::Ptr{Cvoid}
    cDeSerializeFMUstate::Ptr{Cvoid}
    cGetDirectionalDerivative::Ptr{Cvoid}

    # Co Simulation function calls
    cSetRealInputDerivatives::Ptr{Cvoid}
    cGetRealOutputDerivatives::Ptr{Cvoid}
    cDoStep::Ptr{Cvoid}
    cCancelStep::Ptr{Cvoid}
    cGetStatus::Ptr{Cvoid}
    cGetRealStatus::Ptr{Cvoid}
    cGetIntegerStatus::Ptr{Cvoid}
    cGetBooleanStatus::Ptr{Cvoid}
    cGetStringStatus::Ptr{Cvoid}

    # Model Exchange function calls
    cEnterContinuousTimeMode::Ptr{Cvoid}
    cGetContinuousStates::Ptr{Cvoid}
    cGetDerivatives::Ptr{Cvoid}
    cSetTime::Ptr{Cvoid}
    cSetContinuousStates::Ptr{Cvoid}
    cCompletedIntegratorStep::Ptr{Cvoid}
    cEnterEventMode::Ptr{Cvoid}
    cNewDiscreteStates::Ptr{Cvoid}
    cGetEventIndicators::Ptr{Cvoid}
    cGetNominalsOfContinuousStates::Ptr{Cvoid}

    # paths of ziped and unziped FMU folders
    path::String
    binaryPath::String
    zipPath::String

    # execution configuration
    executionConfig::FMU2ExecutionConfiguration
    hasStateEvents::Union{Bool, Nothing}
    hasTimeEvents::Union{Bool, Nothing}

    # c-libraries
    libHandle::Ptr{Nothing}
    callbackLibHandle::Ptr{Nothing}

    # START: experimental section (to FMIFlux.jl)
    dependencies::Matrix{Union{fmi2DependencyKind, Nothing}}

    # END: experimental section

    # DEPRECATED
    instanceName::String

    # Constructor
    function FMU2()
        inst = new()
        inst.components = []
        inst.callbackLibHandle = C_NULL

        inst.hasStateEvents = nothing
        inst.hasTimeEvents = nothing

        inst.executionConfig = FMU_EXECUTION_CONFIGURATION_NO_RESET

        return inst
    end
end

"""
Overload the Base.show() function for custom printing of the FMU2.
"""
Base.show(io::IO, fmu::FMU2) = print(io,
"Model name:        $(fmu.modelDescription.modelName)
Type:              $(fmu.type)"
)

"""
Formats a fmi2Status to String.
"""
function fmi2StatusToString(status::fmi2Status)
    if status == fmi2StatusOK
        return "OK"
    elseif status == fmi2StatusWarning
        return "Warning"
    elseif status == fmi2StatusDiscard
        return "Discard"
    elseif status == fmi2StatusError
        return "Error"
    elseif status == fmi2StatusFatal
        return "Fatal"
    elseif status == fmi2StatusPending
        return "Pending"
    else
        return "Unknown"
    end
end

"""
Formats a fmi2Status to Integer.
"""
function fmi2StatusToString(status::Integer)
    if status == fmi2StatusOK
        return "OK"
    elseif status == fmi2StatusWarning
        return "Warning"
    elseif status == fmi2StatusDiscard
        return "Discard"
    elseif status == fmi2StatusError
        return "Error"
    elseif status == fmi2StatusFatal
        return "Fatal"
    elseif status == fmi2StatusPending
        return "Pending"
    else
        return "Unknown"
    end
end

function fmi2CausalityToString(c::fmi2Causality)
    if c == fmi2CausalityParameter
        return "parameter"
    elseif c == fmi2CausalityCalculatedParameter
        return "calculatedParameter"
    elseif c == fmi2CausalityInput
        return "input"
    elseif c == fmi2CausalityOutput
        return "output"
    elseif c == fmi2CausalityLocal
        return "local"
    elseif c == fmi2CausalityIndependent
        return "independent"
    else
        @assert false "fmi2CausalityToString(...): Unknown causality."
    end
end

function fmi2StringToCausality(s::AbstractString)
    if s == "parameter"
        return fmi2CausalityParameter
    elseif s == "calculatedParameter"
        return fmi2CausalityCalculatedParameter
    elseif s == "input"
        return fmi2CausalityInput
    elseif s == "output"
        return fmi2CausalityOutput
    elseif s == "local"
        return fmi2CausalityLocal
    elseif s == "independent"
        return fmi2CausalityIndependent
    else
        @assert false "fmi2StringToCausality($(s)): Unknown causality."
    end
end

function fmi2VariabilityToString(c::fmi2Variability)
    if c == fmi2VariabilityConstant
        return "constant"
    elseif c == fmi2VariabilityFixed
        return "fixed"
    elseif c == fmi2VariabilityTunable
        return "tunable"
    elseif c == fmi2VariabilityDiscrete
        return "discrete"
    elseif c == fmi2VariabilityContinuous
        return "continuous"
    else
        @assert false "fmi2VariabilityToString(...): Unknown variability."
    end
end

function fmi2StringToVariability(s::AbstractString)
    if s == "constant"
        return fmi2VariabilityConstant
    elseif s == "fixed"
        return fmi2VariabilityFixed
    elseif s == "tunable"
        return fmi2VariabilityTunable
    elseif s == "discrete"
        return fmi2VariabilityDiscrete
    elseif s == "continuous"
        return fmi2VariabilityContinuous
    else
        @assert false "fmi2StringToVariability($(s)): Unknown variability."
    end
end

function fmi2InitialToString(c::fmi2Initial)
    if c == fmi2InitialApprox
        return "approx"
    elseif c == fmi2InitialExact
        return "exact"
    elseif c == fmi2InitialCalculated
        return "calculated"
    else
        @assert false "fmi2InitialToString(...): Unknown initial."
    end
end

function fmi2StringToInitial(s::AbstractString)
    if s == "approx"
        return fmi2InitialApprox
    elseif s == "exact"
        return fmi2InitialExact
    elseif s == "calculated"
        return fmi2InitialCalculated
    else
        @assert false "fmi2StringToInitial($(s)): Unknown initial."
    end
end

function fmi2DependencyKindToString(c::fmi2DependencyKind)
    if c == fmi2DependencyKindDependent
        return "dependent"
    elseif c == fmi2DependencyKindConstant
        return "constant"
    elseif c == fmi2DependencyKindFixed
        return "fixed"
    elseif c == fmi2DependencyKindTunable
        return "tunable"
    elseif c == fmi2DependencyKindDiscrete
        return "discrete"
    else
        @assert false "fmi2DependencyKindToString(...): Unknown dependency kind."
    end
end

function fmi2StringToDependencyKind(s::AbstractString)
    if s == "dependent"
        return fmi2DependencyKindDependent
    elseif s == "exact"
        return fmi2DependencyKindConstant
    elseif s == "fixed"
        return fmi2DependencyKindFixed
    elseif s == "tunable"
        return fmi2DependencyKindTunable
    elseif s == "discrete"
        return fmi2DependencyKindDiscrete
    else
        @assert false "fmi2StringToDependencyKind($(s)): Unknown dependency kind."
    end
end


#############################################################################
# Joulia
# A Large-Scale Spatial Power System Model for Julia
# See https://github.com/JuliaEnergy/Joulia.jl
#############################################################################
# This file contains functions and structs for renewables

"""
    get_res_dict(res_df::DataFrame)

Load renewables.
"""
function get_res_dict(res_df::DataFrame)
	stacked_res = melt(res_df, :Node, variable_name=:Technology,
			           value_name=:Capacity)

	stacked_res[:id] = Symbol.(string.(stacked_res[:Technology]) .* "_"
		                      .* string.(stacked_res[:Node]))

	stacked_res = stacked_res[[:id, :Node, :Technology, :Capacity]]
	stacked_res = stacked_res[stacked_res[:Capacity] .> 0, :]
	res_dict = df_to_dict_with_id(stacked_res)

	return stacked_res[:id], res_dict
end

"""
    get_avail_dict(avail::Dict, res_dict)

Determine renewables availablity.
"""
function get_avail_dict(avail::Dict, res_dict)

	tech = res_dict[:Technology]
	node = res_dict[:Node]

	avail_dict = Dict{Symbol,Array{Float64, 1}}()
	for id in keys(tech)

		t = tech[id]
		n = Symbol(node[id])

		if haskey(avail, t)
			if Symbol(node[id]) in names(avail[t])
				avail_dict[id] = avail[t][n]
			else
				avail_dict[id] = fill(0.0, 8760)
			end
		else
			avail_dict[id] = avail[:global][t]
		end
	end

	return avail_dict
end

"""
    calc_infeed(cap_dict, avail_dict)

Determine renewables infeed.
"""
function calc_infeed(cap_dict, avail_dict)

	infeed = map(collect(keys(cap_dict))) do id
		cap = cap_dict[id]
		avail = avail_dict[id]

		id => avail * cap
	end |> Dict

	return infeed
end

mutable struct RenewableEnergySource
	id
	technology
	node
	infeed

	function RenewableEnergySource(res_df::DataFrame, avail::Dict)

		ids, res_dict = get_res_dict(res_df)
		avail_dict = get_avail_dict(avail,res_dict)
		infeed = calc_infeed(res_dict[:Capacity], avail_dict)

		return new(ids, res_dict[:Technology], res_dict[:Node], infeed)
	end
end


mesh_file_names = ["../mesh/square_meshes/mesh_test_1.g",
                   "../mesh/square_meshes/mesh_test_0.5.g",
                   "../mesh/square_meshes/mesh_test_0.25.g",
                   "../mesh/square_meshes/mesh_test_0.125.g",
                   "../mesh/square_meshes/mesh_test_0.0625.g",
                   "../mesh/square_meshes/mesh_test_0.03125.g",
                   "../mesh/square_meshes/mesh_test_0.015625.g",
                   "../mesh/square_meshes/mesh_test_0.0078125.g"]

number_of_nodes = [4, 9, 25, 81, 289, 1089, 4225, 16641]
number_of_elements = [1, 2^2, 4^2, 8^2, 16^2, 32^2, 64^2, 128^2]
number_of_node_set_nodes = [1, 2, 4, 8, 16, 32, 64, 128] .+ 1

function test_square_mesh(n::Int)
    mesh_file_name = abspath(mesh_file_names[n])
    test_name = rpad("Testing square mesh: $(basename(mesh_file_name))", 96)
    @testset "$test_name" begin
        @suppress begin
            #
            # read method test
            #
            exo_id = Exodus.open_exodus_database(abspath(mesh_file_names[n]))
            init = Exodus.Initialization(exo_id)
            #
            # tests on the simple numbers
            #
            @test init.num_dim == 2
            @test init.num_nodes == number_of_nodes[n]
            @test init.num_elem == number_of_elements[n]
            @test init.num_elem_blk == 1
            @test init.num_node_sets == 4
            @test init.num_side_sets == 4
            #
            # test coordinates
            #
            coords = Exodus.read_coordinates(exo_id, init.num_dim, init.num_nodes)
            @test size(coords, 1) == number_of_nodes[n]
            @test size(coords, 2) == 2
            #
            # read block ids
            #
            block_ids = Exodus.read_block_ids(exo_id, init.num_elem_blk)
            @test size(block_ids, 1) == 1
            @test block_ids[1] == 1
            blocks = Exodus.read_blocks(exo_id, block_ids)
            block = blocks[1]
            @test block.block_id == 1
            @test block.num_elem == number_of_elements[n]
            @test block.num_nodes_per_elem == 4
            @test block.elem_type == "QUAD4"
            @test size(block.conn, 1) == 4 * number_of_elements[n]
            # @test size(block.conn, 1) == number_of_elements[n]
            # @test size(block.conn, 2) == 4
            #
            # test node set initialization
            #
            node_set_ids = Exodus.read_node_set_ids(exo_id, init.num_node_sets)
            @test node_set_ids == [1, 2, 3, 4]
            node_sets = Exodus.read_node_sets(exo_id, node_set_ids)
            @test size(node_sets, 1) == 4
            for (m, node_set) in enumerate(node_sets)
                # @show node_set
                @test node_set.node_set_id == node_set_ids[m]
                @test node_set.num_nodes == number_of_node_set_nodes[n]
                @test size(node_set.nodes, 1) == number_of_node_set_nodes[n]
            end
            #
            mesh = Exodus.Mesh(coords, blocks, node_sets)
            @test size(mesh.coords, 1) == number_of_nodes[n]
            @test size(mesh.coords, 2) == 2

            Exodus.close_exodus_database(exo_id)
        end
    end
end

for (n, mesh) in enumerate(mesh_file_names)
    test_square_mesh(n)
end


type SettlementPhysical <: SettlementType end
type SettlementCash <: SettlementType end

type SwaptionResults{S <: AbstractString}
  value::Float64
  additionalResults::Dict{S, Float64}
end

SwaptionResults() = SwaptionResults(0.0, Dict("spreadCorrection" => 0.0, "strike" => 0.0, "atmForward" => 0.0, "annuity" => 0.0, "swapLength" => 0.0, "stdDev" => 0.0, "vega" => 0.0))

function reset!(results::SwaptionResults)
  results.value = 0.0
  for k in keys(results.additionalResults)
    results.additionalResults[k] = 0.0
  end

  return results
end

type Swaption{E <: Exercise, S <: SettlementType, P <: PricingEngine} <: Option
  lazyMixin::LazyMixin
  swap::VanillaSwap
  exercise::E
  delivery::S
  results::SwaptionResults
  pricingEngine::P

  # function call{E, S}(::Type{Swaption}, lz::LazyMixin, swap::VanillaSwap, exercise::E, delivery::S, results::SwaptionResults)
  #   new{E, S, PricingEngine}(lz, swap, exercise, delivery, results)
  # end

  function call{E, S, P}(::Type{Swaption}, lz::LazyMixin, swap::VanillaSwap, exercise::E, delivery::S, results::SwaptionResults, pricingEngine::P = NullSwaptionEngine())
    new{E, S, P}(lz, swap, exercise, delivery, results, pricingEngine)
  end
end

Swaption{E <: Exercise}(swap::VanillaSwap, exercise::E) = Swaption(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults())
Swaption{E <: Exercise, P <: PricingEngine}(swap::VanillaSwap, exercise::E, pe::P) = Swaption(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults(), pe)

type NonstandardSwaption{E <: Exercise, S <: SettlementType, P <: PricingEngine} <: Option
  lazyMixin::LazyMixin
  swap::NonstandardSwap
  exercise::E
  delivery::S
  results::SwaptionResults
  pricingEngine::P

  function call{E, S, P}(::Type{NonstandardSwaption}, lz::LazyMixin, swap::NonstandardSwap, exercise::E, delivery::S, results::SwaptionResults, pricingEngine::P = NullSwaptionEngine())
    new{E, S, P}(lz, swap, exercise, delivery, results, pricingEngine)
  end
end

NonstandardSwaption{E <: Exercise}(swap::NonstandardSwap, exercise::E) = NonstandardSwaption(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults())
NonstandardSwaption{E <: Exercise, P <: PricingEngine}(swap::NonstandardSwap, exercise::E, pe::P) = NonstandardSwaption(LazyMixin(), swap, exercise, SettlementPhysical(), SwaptionResults(), pe)

function perform_calculations!(swaption::Union{Swaption, NonstandardSwaption})
  reset!(swaption.results)
  _calculate!(swaption.pricingEngine, swaption)

  return swaption
end

function npv(swaption::Union{Swaption, NonstandardSwaption})
  calculate!(swaption)

  return swaption.results.value
end

function calibration_basket(swaption::NonstandardSwaption, swaptionEngine::PricingEngine, swapIndex::SwapIndex, swaptionVol::SwaptionVolatilityStructure,
                            basketType::CalibrationBasketType)
  return calibration_basket(swaptionEngine, swaption, swapIndex, swaptionVol, basketType)
end

function clone(swaption::Swaption, pe::PricingEngine = swaption.pricingEngine)
  lazyMixin, res = pe == swaption.pricingEngine ? (swaption.lazyMixin, swaption.results) : (LazyMixin(), SwaptionResults())

  newSwaption = Swaption(lazyMixin, swaption.swap, swaption.exercise, swaption.delivery, res, pe)

  return newSwaption
end

get_pricing_engine_type{E, S, P}(::Swaption{E, S, P}) = P


using Test, CUDA, Knet, AutoGrad
using CUDA.CUDNN:
    cudnnDropoutForward,
    cudnnDropoutForward!,
    cudnnDropoutBackward,
    cudnnDropoutSeed,
    cudnnDropoutDescriptor,
        cudnnDropoutDescriptor_t,
        cudnnCreateDropoutDescriptor,
        cudnnSetDropoutDescriptor,
        cudnnGetDropoutDescriptor,
        cudnnRestoreDropoutDescriptor,
        cudnnDestroyDropoutDescriptor,
    cudnnDropoutGetStatesSize,
    cudnnDropoutGetReserveSpaceSize,
    handle

@testset "cudnn/dropout" begin

    N,P = 1000, 0.7
    x = Param(CUDA.rand(N))
    d = cudnnDropoutDescriptor(P)
    cudnnDropoutSeed[] = 1
    @test @gcheck cudnnDropoutForward(x; dropout = P)
    @test @gcheck cudnnDropoutForward(x, d)
    @test @gcheck cudnnDropoutForward!(similar(x), x; dropout = P)
    @test @gcheck cudnnDropoutForward!(similar(x), x, d)
    cudnnDropoutSeed[] = -1

end


"""
    StdDiagnostics

This module defines many standard diagnostic variables and groups that may
be used directly by experiments.
"""
module StdDiagnostics

using CLIMAParameters
using CLIMAParameters.Planet
using CLIMAParameters.Atmos
using CLIMAParameters.SubgridScale

using KernelAbstractions
using MPI
using OrderedCollections
using Printf
using StaticArrays

using ..Diagnostics # until old diagnostics groups are removed
using ..Atmos
using ..BalanceLaws
using ..ConfigTypes
using ..DGMethods
using ..DiagnosticsMachine
import ..DiagnosticsMachine:
    Settings,
    dv_name,
    dv_attrib,
    dv_args,
    dv_project,
    dv_scale,
    dv_PointwiseDiagnostic,
    dv_HorizontalAverage
using ..Mesh.Grids
using ..Mesh.Interpolation
using ..Mesh.Topologies
using ..MPIStateArrays
using Thermodynamics
using ..TurbulenceClosures
using ..VariableTemplates
using ..Writers

# Pre-defined diagnostic variables

# Atmos
include("atmos_les_diagnostic_vars.jl")
include("atmos_gcm_diagnostic_vars.jl")

# Pre-defined diagnostics groups

# Atmos
include("atmos_les_default.jl")
include("atmos_gcm_default.jl")

end # module StdDiagnostics


# Parameter iterators
postagv(m)=m[1];
deprelv(m)=m[2];
lcountv(m)=m[3];
rcountv(m)=m[4];
distancev(m)=m[5];
bufferm(m)=m[6];
parserv(m)=m[7];

# To initialize parser model from scratch
function initmodel1(o, s; ftype=Float32)
    model = Any[]

    # vector-level initialization
    dpostag, ddeprel, dcount = o[:embed]
    for (k,n,d) in ((:postag,17,dpostag),(:deprel,37,ddeprel),(:lcount,10,dcount),(:rcount,10,dcount),(:distance,10,dcount))
        push!(model, [ initr(d) for i=1:n ])
    end

    # buffer-lstm initialization
    r_b, wr_b = initbuf(o[:bufembed], o[:bufhidden]) # you may change lstm - gru with kwargs
    push!(model, [r_b, wr_b])

    # decision module initialization
    p = o[:arctype](s)
    f = features([p], o[:feats], model)
    mlpdims = (length(f)+r_b.hiddenSize, o[:hidden]..., p.nmove)
    info("mlpdims: $mlpdims")
    decider = Any[]
    for i=2:length(mlpdims)
        push!(decider, initx(mlpdims[i], mlpdims[i-1]))
        push!(decider, initx(mlpdims[i], 1))
    end
    push!(model, decider)
    optims = optimizers(model, o[:optimization])
    return model, optims
end

# To seperate feature vectors and all others
function splitmodel(pmodel)
    # feat-model
    featmodel = Array{Any}(5)
    for k in 1:5 # (:postag,:deprel,:lcount,:rcount,:distance)
        featmodel[k] = Any[]
        pmodel_k = pmodel[k]
        for i in 1:length(pmodel_k)
            push!(featmodel[k], pmodel_k[i])
        end
    end

    # buffer-lstm
    bufmodel = Any[]; buftemp = bufferm(pmodel);
    for i in 1:length(buftemp); push!(bufmodel, buftemp[i]);end;

    # mlp-model
    mlpmodel = Any[]; mlptemp = parserv(pmodel);
    for i in 1:length(mlptemp); push!(mlpmodel, mlptemp[i]); end;
    return (featmodel, bufmodel, mlpmodel)
end

function mlp(w,x; pdrop=(0,0))
    x = dropout(x,pdrop[1])
    for i=1:2:length(w)-2
        x = relu.(w[i]*x .+ w[i+1])
        x = dropout(x,pdrop[2])
    end
    return w[end-1]*x .+ w[end]
end

function scan_buffer(bufmodel, featmodel, parser)
    rbuf, wbuf, pvecs = bufmodel[1], bufmodel[2], postagv(featmodel)
    sentence = parser.sentence
    seqe = length(sentence)
    seqs = parser.wptr
    range = seqe:-1:seqs
    #instate = zeros(rbuf.dataType, rbuf.hiddenSize, 1)
    instate = zeros(Float32, 256, 1)
    ybuf = hout = cout = (gpu() >=0 ? KnetArray(instate) : instate)
    for i in range
        input = vcat(sentence.cavec[i], pvecs[sentence.postag[i]])
        input = (gpu() >=0 ? KnetArray(input) : input)
        ybuf, hout, cout, _ = rnnforw(rbuf, wbuf, input, hout, cout, hy=true, cy=true)
    end
    return ybuf
end

function scan_bufbatch(bufmodel, featmodel, parsers)
    # Todo: we may cache some parsers in future
    yalls = Any[]
    for p in parsers
        ybuf = scan_buffer(bufmodel, featmodel, p)
        push!(yalls, ybuf)
    end
    yalls = cat1d(yalls...)
    ncols = length(parsers)
    nrows = div(length(yalls), ncols)
    return reshape(yalls, nrows, ncols)
end

function oracleloss(allmodel, sentences, arctype, feats; losses=nothing, pdrop=(0.0, 0.0))
    parsers =  map(arctype, sentences)
    mcosts  =  Array{Cost}(parsers[1].nmove)
    parserdone = falses(length(parsers))
    featmodel, bufmodel, mlpmodel = splitmodel(allmodel)
    totalloss = 0.0

    while !all(parserdone)
        fmatrix = features(parsers, feats, featmodel)
        if gpu() >= 0
            fmatrix = KnetArray(fmatrix)
        end
        ybufs = scan_bufbatch(bufmodel, featmodel, parsers)
        input = vcat(fmatrix, ybufs)
        scores = mlp(mlpmodel, input, pdrop=pdrop)
        logprobs = logp(scores, 1)
        for (i, p) in enumerate(parsers)
            if parserdone[i]
                continue
            end
            movecosts(p, p.sentence.head, p.sentence.deprel, mcosts)
            goldmove = indmin(mcosts)

            if mcosts[goldmove] == typemax(Cost)
                parserdone[i] = true
                p.sentence.parse = p
            else
                totalloss -= logprobs[goldmove, i]
                move!(p, goldmove)
                if losses != nothing
                    loss1 = -getval(logprobs)[goldmove,i]
                    losses[1] += loss1
                    losses[2] += 1
                    if losses[2] < 1000
                        losses[3] = losses[1]/losses[2]
                    else
                        losses[3] = 0.999 * losses[3] + 0.001 * loss1
                    end
                end
            end
        end
    end
    return totalloss / length(sentences)
end

oraclegrad = grad(oracleloss)

# Implement test related

# endofparse to use in test
endofparse(p)=(p.sptr == 1 && p.wptr > p.nword)

# TODO: You need to modify score calculations
function oracletest(pmodel, corpus, arctype, feats, batchsize; pdrop=(0.0, 0.0))
    sentbatches = minibatch(corpus, batchsize)
    featmodel, bufmodel, mlpmodel = splitmodel(pmodel)

    for sentences in sentbatches
        parsers = map(arctype, sentences)
        mcosts = Array{Cost}(parsers[1].nmove) #Array(Cost, parsers[1].nmove)
        parserdone = falses(length(parsers))
        while !all(parserdone)
            fmatrix = features(parsers, feats, featmodel)
            if gpu()>=0
                fmatrix = KnetArray(fmatrix)
            end
            ybufs = scan_bufbatch(bufmodel, featmodel, parsers)
            input = vcat(fmatrix, ybufs)
            scores = mlp(mlpmodel, input, pdrop=pdrop)
            logprob = Array(logp(scores, 1))

            for (i,p) in enumerate(parsers)
                if parserdone[i]
                    continue
                end

                isorted = sortperm(logprob[:, i], rev=true) # ith column for ith instance
                for m in isorted # take the best&valid action
                    if moveok(p, m)
                        move!(p, m)
                        break
                    end
                end
                if endofparse(p)
                    parserdone[i] = true
                    p.sentence.parse = p
                end
            end
        end
    end
    return las(corpus)
end

# labeled-attachment score calculator
function las(corpus)
    nword = ncorr = 0
    for s in corpus
        p = s.parse
        nword += length(s)
        ncorr += sum((s.head .== p.head) .& (s.deprel .== p.deprel))
    end
    ncorr / nword
end

function empty_parses!(corpus)
    for s in corpus
        s.parse = nothing
    end
end


import Random: shuffle

@generated function myshuffle(xs)
    canShuffleEmpty = try
        shuffle([])
        true
    catch
        false
    end
    if canShuffleEmpty
        :(shuffle(xs))
    else
        :(length(xs)>0 ? shuffle(xs) : xs)
    end
end

@generated function myshuffle(rng, xs)
    canShuffleEmpty = try
        shuffle([])
        true
    catch
        false
    end
    if canShuffleEmpty
        :(shuffle(rng, xs))
    else
        :(length(xs)>0 ? shuffle(rng, xs) : xs)
    end
end

@doc """
    loop_update!(model, param::Parameter)
    loop_update!(model, T::Real,
                 Jz::AbstractArray,
                 Jxy::AbstractArray,
                 Gamma:AbstractArray)

Updates spin configuration by loop algorithm
under the temperature `T = param["T"]` and coupling constants `Jz, Jxy` and transverse field `Gamma`
"""
@inline function loop_update!(model::QuantumXXZ, param::Parameter)
    p = convert_parameter(model, param)
    return loop_update!(model, p...)
end

function loop_update!(model::QuantumXXZ, T::Real,
                      Jzs::AbstractArray, Jxys::AbstractArray, Gs::AbstractArray)
    rng = model.rng
    lo_types = [LET_FMLink, LET_AFLink, LET_Vertex, LET_Cross]
    nsites = numsites(model)
    S2 = model.S2
    nspins = nsites*S2
    nbonds = numbonds(model)
    nst = numsitetypes(model)
    nbt = numbondtypes(model)
    weights = zeros(4*nbt+nst)
    for i in 1:nbt
        nb = numbonds(model,i)*S2*S2
        z = nb*Jzs[i]
        x = nb*abs(Jxys[i])
        if z > x
            ## AntiFerroIsing like
            weights[(4i-3):(4i)] .= 0.5*[0.0, z-x, x, 0.0]
        elseif z < -x
            ## FerroIsing like
            weights[(4i-3):(4i)] .= 0.5*[-z-x, 0.0, 0.0, x]
        else
            ## XY like
            weights[(4i-3):(4i)] .= 0.25*[0.0, 0.0, x+z, x-z]
        end
    end
    for i in 1:nst
        ns = numsites(model,i)*S2
        weights[4nbt+i] = 0.5*ns*Gs[i]
    end
    accumulated_weights = cumsum(weights)
    op_dt = T/accumulated_weights[end]

    spins = model.spins[:]
    currents = collect(1:nspins)
    uf = UnionFind(nspins)

    ops = LocalLoopOperator[]
    iops = 1 # index of the next operator in the original string
    t = randexp(rng)*op_dt
    while t <= 1.0 || iops <= length(model.ops)
        if iops > length(model.ops) || t < model.ops[iops].time
            ## INSERT
            ot = searchsortedfirst(accumulated_weights, rand(rng)*accumulated_weights[end])
            if ot <= 4nbt
                ## Bond
                bt = ceil(Int, ot/4)
                ot = mod1(ot,4)
                lo_type = lo_types[ot]
                b = bonds(model,bt)[rand(rng, 1:numbonds(model,bt))]
                s1,s2 = source(b), target(b)
                ss1, ss2 = rand(rng, 1:S2, 2)
                if ifelse(spins[site2subspin(s1,ss1,S2)] == spins[site2subspin(s2,ss2,S2)],
                          lo_type == LET_FMLink || lo_type == LET_Cross,
                          lo_type == LET_AFLink || lo_type == LET_Vertex)
                    push!(ops, LocalLoopOperator(lo_type, t, nsites+b.id, (ss1, ss2)))
                    t += randexp(rng)*op_dt
                else
                    t += randexp(rng)*op_dt
                    continue
                end
            else
                ## Site
                st = ot-4nbt
                s = sites(model,st)[rand(rng, 1:numsites(model,st))]
                ss = rand(rng, 1:S2)
                push!(ops, LocalLoopOperator(LET_Cut, t, s, (ss,ss)))
                t += randexp(rng)*op_dt
            end
        else
            op = model.ops[iops]
            iops += 1
            if op.isdiagonal
                ## REMOVE
                continue
            elseif op.let_type == LET_Cut
                push!(ops, op)
            else
                push!(ops, op)
                op = ops[end]
                b = op.space-nsites
                ss1,ss2 = op.subspace
                bt = bondtype(model,b)
                otype = ifelse(rand(rng)*(weights[4bt-1]+weights[4bt])<weights[4bt-1], LET_Vertex, LET_Cross)
                op.let_type = otype
            end
        end

        op = ops[end]
        if op.let_type == LET_Cut
            s = op.space
            ss = op.subspace[1]
            subspin = site2subspin(s,ss,S2)
            op.bottom_id = currents[subspin]
            c = addnode!(uf)
            currents[subspin] = c
            op.top_id = c
            spins[subspin] *= ifelse(op.isdiagonal, 1, -1)
        else
            b = op.space - nsites
            ss1,ss2 = op.subspace
            s1 = source(model, b)
            s2 = target(model, b)
            subspin1 = site2subspin(s1,ss1,S2)
            subspin2 = site2subspin(s2,ss2,S2)
            if op.let_type == LET_FMLink || op.let_type == LET_AFLink
                c = unify!(uf, currents[subspin1], currents[subspin2])
                currents[subspin1] = currents[subspin2] = c
                op.bottom_id = op.top_id = c
            elseif op.let_type == LET_Cross
                op.bottom_id = currents[subspin1]
                op.top_id = currents[subspin2]
                spins[subspin1], spins[subspin2] = spins[subspin2], spins[subspin1]
                currents[subspin1], currents[subspin2] = currents[subspin2], currents[subspin1]
            else # if op.let_type == LET_Vertex
                unify!(uf, currents[subspin1], currents[subspin2])
                op.bottom_id = currents[subspin1]
                c = addnode!(uf)
                op.top_id = c
                currents[subspin1] = currents[subspin2] = c
                spins[subspin1] *= ifelse(op.isdiagonal, 1, -1)
                spins[subspin2] *= ifelse(op.isdiagonal, 1, -1)
            end
        end
    end # of while loop

    ## PBC for imaginary time axis
    subspin = 0
    for s in 1:nsites
        ups0 = zeros(Int,0)
        ups1 = zeros(Int,0)
        downs0 = zeros(Int,0)
        downs1 = zeros(Int,0)
        for ss in 1:S2
            subspin += 1
            push!(ifelse(spins[subspin]==1, ups0, downs0), subspin)
            push!(ifelse(model.spins[subspin]==1, ups1, downs1), subspin)
        end
        @assert length(ups0) == length(ups1)
        @assert length(downs0) == length(downs1)
        for (u, u2) in zip(ups0, myshuffle(rng,ups1))
            unify!(uf, u, currents[u2])
        end
        for (d, d2) in zip(downs0, myshuffle(rng,downs1))
            unify!(uf, d, currents[d2])
        end
    end

    model.ops = ops

    nc = clusterize!(uf)
    flips = rand(rng, [1,-1], nc)
    for s in 1:nspins
        model.spins[s] *= flips[clusterid(uf, s)]
    end
    for op in model.ops
        if op.let_type == LET_Cross || op.let_type == LET_Vertex || op.let_type == LET_Cut
            bid = clusterid(uf, op.bottom_id)
            tid = clusterid(uf, op.top_id)
            op.isdiagonal ⊻= (flips[bid] != flips[tid])
        end
    end

    return uf
end

@gen_convert_parameter(QuantumXXZ, ("T", 1, 1.0),
                                   ("Jz", numbondtypes, 1.0),
                                   ("Jxy", numbondtypes, 1.0),
                                   ("Gamma", numsitetypes, 0.0),
                                  )


# type definitions

export AbstractQC, AbstractQuState, AbstractQuCircuit

# abstracts
abstract AbstractQC{N}
abstract AbstractQuState{T,N}
abstract AbstractQuCircuit{N} <: AbstractQC{N}


# Makes VRT of all the geotiffs in a folder and in a region.  One vrt per UTM zone.

using VAWTools
import ArchGDAL
const AG = ArchGDAL

regions = [11]
dir = "../output/rasters-fit_sigma/production-with-dhdt/"

for region in [3,4,7,11,14][end:end]
    @show region
    fls = Dict() # UTM zones => list of files
    for f in readdir(dir)
        ff, ext = splitext(f)
        ext!=".tif" && continue
        reg, nr = parse.(Int, split(split(ff, '-')[end], '.') )
        reg!=region && continue
        # get all projections
        #p = readstring(`gdalsrsinfo -o proj4 $(joinpath(dir,f))`)
        # Faster with ArchGDAL?
        p = AG.registerdrivers() do
            AG.read(joinpath(dir,f)) do dataset
                proj4 = strip(AG.toPROJ4(AG.importWKT(AG.getproj(dataset))))
            end
        end
        zone = [[nr,parse(Int, split(s,'=')[2])] for s in split(p) if split(s, '=')[1]=="+zone"][1]
        push!(get!(fls, zone[2], []), joinpath(dir,f))
    end

    for (z,fs) in fls
        @show z
        @show vrt = joinpath(dir, "composite-$(VAWTools.int2str2(region))-zone_$z.vrt")
        run(`gdalbuildvrt $vrt $(fs...)`)
    end
end


function cross_map(r::MersenneTwister)
    map = YAML.load_file("maps/cross.yaml")
    exits = [(4, 1), (1, 4), (4, 7), (7, 4)]
    start = rand(r, exits)
    exit = rand(r, setdiff(exits, [start]))
    map["starts"] = [Dict("x"=>start[2], "y"=>start[1])]
    map["exits"] = [Dict("x"=>exit[2], "y"=>exit[1])]
    push!(map["objects"], Dict("type"=>"food", "color"=>"green",
                               "x"=>exit[2], "y"=>exit[1]))
    map
end

function cross_map(seed::Int64)
    r = MersenneTwister(seed)
    cross_map(r)
end

function cross_meta()
    Dict{String, Any}("max_step" => 100)
end

function cross_search_fitness(cont_f::Function; seed::Int64=0)
    e = Episode(Grid(cross_map(seed)); reward=food_reward, meta=cross_meta())
    run!(e, cont_f)
    steps = terminate(e) ? e.meta["max_step"] : e.meta["step"]
    1 - steps / e.meta["max_step"]
end

function cross_memorize_fitness(cont_f::Function; seed::Int64=0)
    map_function() = cross_map(seed)
    learn_fitness(cont_f, map_function, cross_meta)
end

function cross_map(r::MersenneTwister, strategy::Int64)
    map = YAML.load_file("maps/cross.yaml")
    exits = [(4, 1), (1, 4), (4, 7), (7, 4)]
    si = rand(r, 1:4)
    start = exits[si]
    exit = exits[mod((si-1)+strategy, 4)+1]
    map["starts"] = [Dict("x"=>start[2], "y"=>start[1])]
    map["exits"] = [Dict("x"=>exit[2], "y"=>exit[1])]
    push!(map["objects"], Dict("type"=>"food", "color"=>"green",
                               "x"=>exit[2], "y"=>exit[1]))
    map
end

function cross_strategy_fitness(cont_f::Function; seed::Int64=0)
    r = MersenneTwister(seed)
    strategy = rand(r, 1:3)
    map_function() = cross_map(r, strategy)
    learn_fitness(cont_f, map_function, cross_meta)
end


module FPRS

using FFTW
using LinearAlgebra
using Random         # for densities
using DelimitedFiles # readdlm for loadElser

include("utilities.jl")
include("densities.jl")
include("projections.jl")
include("loadElser.jl")
include("bookkeeping.jl")

export
# Utilities
randBall
randRotationMatrix
# Densities
genGaussWGrid
genGauss
genRandGaussWGrid
genRandGauss
# Projections
pM
pM0
pP
pTa
pTn
eM
eM0
eP
eTa
eTn
# loadElser
loadElser
loadElserPure
# bookkeeping
pri
pri1D
lpri

end # module


# module for handling the sniper map
module UrbanMaps

export
    UrbanMap,
    locals,
    inbounds

using Colladas # probably don't need this
using Maps

type UrbanMap <: Map
    buildings::Vector{Building}
    nBuildings::Int
    xSize::Int
    ySize::Int
    xl::Float64
    yl::Float64
    gridSize::Int
    # for rectangular maps only
    xCenters::Vector{Float64}
    yCenters::Vector{Float64}
    xLengths::Vector{Float64}
    yLengths::Vector{Float64}

    # TODO: need to deal with normalizing the building vertices for diff x,y sizes
    function UrbanMap(coll::Collada, xSize::Int64, ySize::Int64)
        self = new()

        nBuildings = coll.nObjects

        self.xSize    = xSize
        self.ySize    = ySize
        self.gridSize = xSize * ySize

        self.xl = float(xSize)
        self.yl = float(ySize)

        self.nBuildings = nBuildings

        # assumes a square grid
        self.buildings = [Polygon(xSize * c.points, c.vertices) for c in coll.normalizedObjects]

        return self
    end

    # Generates a map with n rectangular buildings
    # all input coordinates are normalized to unit square
    function UrbanMap(n::Int64, xCenters::Vector{Float64}, yCenters::Vector{Float64},
                      xLengths::Vector{Float64}, yLengths::Vector{Float64},
                      xSize::Int64, ySize::Int64)
        self = new()

        self.xSize      = xSize
        self.ySize      = ySize
        self.gridSize   = xSize * ySize
        self.nBuildings = n

        self.xl = float(xSize)
        self.yl = float(ySize)

        buildings = Array(Polygon, 0)

        # convert centers and lengths to vertices
        xm = [-0.5, 0.5, 0.5, -0.5]
        ym = [0.5, 0.5, -0.5, -0.5]
        verts = [1,2,3,4]
        for i = 1:n
            # make vertices
            xc = xCenters[i]
            yc = yCenters[i]
            xl = xLengths[i]
            yl = yLengths[i]
            points = zeros(4,3)
            for j = 1:4
                points[j,1] = xc + xm[j] * xl
                points[j,2] = yc + ym[j] * yl
            end
            p = Polygon(xSize * points, verts)
            push!(buildings, p)
        end
        self.buildings = buildings

        self.xCenters = xCenters
        self.yCenters = yCenters
        self.xLengths = xLengths
        self.yLengths = yLengths
        return self
    end

    function UrbanMap(t::Symbol, xSize::Int64, ySize::Int64)
        self = new()

        self.xSize = xSize
        self.ySize = ySize

        if t == :cirlce
            self.nBuildings = 1
            self.buildings = [Circular((0.5, 0.5), 0.1)]
        end

        return self
    end

end

function locals(map::UrbanMap)
    return (map.xCenters, map.yCenters, map.xLengths, map.yLengths)
end

# generates a map with rectangular buildings
function generateMaps()

end

function inbounds(map::UrbanMap, x::Vector{Float64})
    if 0.0 <= x[1] <= map.xl && 0.0 <= x[2] <= map.yl
        return true
    end
    return false
end

end # module


using Images, Base.Test, Color, FixedPointNumbers

macro chk(a, b)
    :(@test $a == $b && typeof($a) == typeof($b))
end

# MapNone
mapi = MapNone{Int}()
@chk map(mapi, 7) 7
@chk map(mapi, 0x07) 7
a7 = Int[7]
@chk map(mapi, [0x07]) a7
@test map(mapi, a7) === a7
mapi = MapNone{RGB24}()
g = Ufixed8(0.1)
@chk Images.map1(mapi, 0.1) g
@chk map(mapi, 0.1) RGB24(g,g,g)
@chk map(mapi, Gray(0.1)) RGB24(g,g,g)
@chk map(mapi, g) RGB24(g,g,g)
@chk map(mapi, true) RGB24(0xff,0xff,0xff)
@chk map(mapi, false) RGB24(0x00,0x00,0x00)
mapi = MapNone{RGB{Float32}}()
g = 0.1f0
@chk Images.map1(mapi, 0.1) g
@chk map(mapi, 0.1) RGB(g,g,g)

c = RGB(0.1,0.2,0.3)
mapi = MapNone{HSV{Float64}}()
@chk map(mapi, c) convert(HSV, c)

# issue #200
c = RGBA{Ufixed8}(1,0.5,0.25,0.8)
mapi = MapNone{Images.ColorTypes.BGRA{Ufixed8}}()
@chk map(mapi, c) convert(Images.ColorTypes.BGRA{Ufixed8}, c)

# BitShift
mapi = BitShift{Uint8,7}()
@chk map(mapi, 0xff) 0x01
@chk map(mapi, 0xf0ff) 0xff
@chk map(mapi, [0xff]) Uint8[0x01]
mapi = BitShift{Ufixed8,7}()
@chk map(mapi, 0xffuf8) 0x01uf8
@chk map(mapi, 0xf0ffuf16) 0xffuf8
mapi = BitShift{ARGB32,4}()
@chk map(mapi, 0xffuf8) ARGB32(0xff0f0f0f)
mapi = BitShift{RGB24,2}()
@chk map(mapi, Gray(0xffuf8)) RGB24(0x003f3f3f)
mapi = BitShift{ARGB32,2}()
@chk map(mapi, Gray(0xffuf8)) ARGB32(0xff3f3f3f)
@chk map(mapi, GrayAlpha{Ufixed8}(Gray(0xffuf8),0x3fuf8)) ARGB32(0x0f3f3f3f)
mapi = BitShift{RGB{Ufixed8},2}()
@chk map(mapi, Gray(0xffuf8)) RGB(0x3fuf8, 0x3fuf8, 0x3fuf8)
mapi = BitShift{ARGB{Ufixed8},2}()
@chk map(mapi, Gray(0xffuf8)) ARGB{Ufixed8}(0x3fuf8, 0x3fuf8, 0x3fuf8, 0xffuf8)
@chk map(mapi, GrayAlpha{Ufixed8}(Gray(0xffuf8),0x3fuf8)) ARGB{Ufixed8}(0x3fuf8, 0x3fuf8, 0x3fuf8, 0x0fuf8)
mapi = BitShift{RGBA{Ufixed8},2}()
@chk map(mapi, Gray(0xffuf8)) RGBA{Ufixed8}(0x3fuf8, 0x3fuf8, 0x3fuf8, 0xffuf8)
@chk map(mapi, GrayAlpha{Ufixed8}(Gray(0xffuf8),0x3fuf8)) RGBA{Ufixed8}(0x3fuf8, 0x3fuf8, 0x3fuf8, 0x0fuf8)
mapi = BitShift(ARGB{Ufixed8}, 8)
@chk map(mapi, RGB{Ufixed16}(1.0,0.8,0.6)) ARGB{Ufixed8}(1.0,0.8,0.6,1.0)
mapi = BitShift(RGBA{Ufixed8}, 8)
@chk map(mapi, RGB{Ufixed16}(1.0,0.8,0.6)) RGBA{Ufixed8}(1.0,0.8,0.6,1.0)
# Issue, #269, IJulia issue #294
bs = BitShift(Gray{Ufixed8}, 8)
v = Gray(ufixed16(0.8))
@chk map(bs, v) Gray{Ufixed8}(0.8)

# Clamp
mapi = ClampMin(Float32, 0.0)
@chk map(mapi,  1.2) 1.2f0
@chk map(mapi, -1.2) 0.0f0
mapi = ClampMin(RGB24, 0.0f0)
@chk map(mapi, RGB{Float32}(-5.3,0.4,0.8)) RGB24(0x000066cc)
mapi = ClampMax(Float32, 1.0)
@chk map(mapi,  1.2)  1.0f0
@chk map(mapi, -1.2) -1.2f0
mapi = ClampMax(RGB24, 1.0f0)
@chk map(mapi, RGB{Float32}(0.2,1.3,0.8)) RGB24(0x0033ffcc)
mapi = ClampMinMax(Float32, 0.0, 1.0)
@chk map(mapi,  1.2) 1.0f0
@chk map(mapi, -1.2) 0.0f0
mapi = ClampMinMax(ARGB32, 0.0f0, 1.0f0)
@chk map(mapi, RGB{Float32}(-0.2,1.3,0.8)) ARGB32(0xff00ffcc)
@chk map(mapi, ARGB{Float32}(-0.2,1.3,0.8,0.6)) ARGB32(0x9900ffcc)
mapi = Clamp(Float32)
@chk map(mapi,  1.2) 1.0f0
@chk map(mapi, -1.2) 0.0f0
mapi = Clamp(Ufixed12)
@chk map(mapi, Ufixed12(1.2)) one(Ufixed12)
mapi = Clamp(Gray{Ufixed12})
@chk map(mapi, Gray(Ufixed12(1.2))) Gray(one(Ufixed12))
mapi = ClampMinMax(RGB24, 0.0, 1.0)
@chk map(mapi, 1.2) RGB24(0x00ffffff)
@chk map(mapi, 0.5) RGB24(0x00808080)
@chk map(mapi, -.3) RGB24(0x00000000)
mapi = ClampMinMax(RGB{Ufixed8}, 0.0, 1.0)
@chk map(mapi, 1.2) RGB{Ufixed8}(1,1,1)
@chk map(mapi, 0.5) RGB{Ufixed8}(0.5,0.5,0.5)
@chk map(mapi, -.3) RGB{Ufixed8}(0,0,0)
mapi = Clamp(RGB{Ufixed8})
@chk map(mapi, RGB(1.2,0.5,-.3)) RGB{Ufixed8}(1,0.5,0)
mapi = Clamp(ARGB{Ufixed8})
@chk map(mapi, ARGB{Float64}(1.2,0.5,-.3,0.2)) ARGB{Ufixed8}(1.0,0.5,0.0,0.2)
@chk map(mapi, RGBA{Float64}(1.2,0.5,-.3,0.2)) ARGB{Ufixed8}(1.0,0.5,0.0,0.2)
@chk map(mapi, 0.2) ARGB{Ufixed8}(0.2,0.2,0.2,1.0)
@chk map(mapi, GrayAlpha{Float32}(Gray(0.2),1.2)) ARGB{Ufixed8}(0.2,0.2,0.2,1.0)
@chk map(mapi, GrayAlpha{Float32}(Gray(-.4),0.8)) ARGB{Ufixed8}(0.0,0.0,0.0,0.8)
mapi = Clamp(RGBA{Ufixed8})
@chk map(mapi, ARGB{Float64}(1.2,0.5,-.3,0.2)) RGBA{Ufixed8}(1.0,0.5,0.0,0.2)
@chk map(mapi, RGBA{Float64}(1.2,0.5,-.3,0.2)) RGBA{Ufixed8}(1.0,0.5,0.0,0.2)
@chk map(mapi, 0.2) RGBA{Ufixed8}(0.2,0.2,0.2,1.0)
@chk map(mapi, GrayAlpha{Float32}(Gray(0.2),1.2)) RGBA{Ufixed8}(0.2,0.2,0.2,1.0)
@chk map(mapi, GrayAlpha{Float32}(Gray(-.4),0.8)) RGBA{Ufixed8}(0.0,0.0,0.0,0.8)
# Issue #253
mapi = Clamp(BGRA{Ufixed8})
@chk map(mapi, RGBA{Float32}(1.2,0.5,-.3,0.2)) BGRA{Ufixed8}(1.0,0.5,0.0,0.2)

@chk clamp(RGB{Float32}(-0.2,0.5,1.8)) RGB{Float32}(0.0,0.5,1.0)
@chk clamp(ARGB{Float64}(1.2,0.5,-.3,0.2)) ARGB{Float64}(1.0,0.5,0.0,0.2)
@chk clamp(RGBA{Float64}(1.2,0.5,-.3,0.2)) RGBA{Float64}(1.0,0.5,0.0,0.2)

# ScaleMinMax
mapi = ScaleMinMax(Ufixed8, 100, 1000)
@chk map(mapi, 100) Ufixed8(0.0)
@chk map(mapi, 1000) Ufixed8(1.0)
@chk map(mapi, 10) Ufixed8(0.0)
@chk map(mapi, 2000) Ufixed8(1.0)
@chk map(mapi, 550) Ufixed8(0.5)
mapinew = ScaleMinMax(Ufixed8, [100,500,1000])
@test mapinew == mapi
mapinew = ScaleMinMax(Ufixed8, [0,500,2000], convert(Uint16, 100), convert(Uint16, 1000))
@test mapinew == mapi
mapi = ScaleMinMax(ARGB32, 100, 1000)
@chk map(mapi, 100) ARGB32(0x00,0x00,0x00,0xff)
@chk map(mapi, 550) ARGB32(0x80,0x80,0x80,0xff)
@chk map(mapi,2000) ARGB32(0xff,0xff,0xff,0xff)
mapi = ScaleMinMax(RGB{Float32}, 100, 1000)
@chk map(mapi,  50) RGB(0.0f0, 0.0f0, 0.0f0)
@chk map(mapi, 550) RGB{Float32}(0.5, 0.5, 0.5)
@chk map(mapi,2000) RGB(1.0f0, 1.0f0, 1.0f0)
A = Gray{Ufixed8}[Ufixed8(0.1), Ufixed8(0.9)]
@test mapinfo(RGB24, A) == MapNone{RGB24}()
mapi = ScaleMinMax(RGB24, A, zero(Gray{Ufixed8}), one(Gray{Ufixed8}))
@test map(mapi, A) == map(mapinfo(RGB24, A), A)
mapi = ScaleMinMax(Float32, [Gray(one(Ufixed8))], 0, 1) # issue #180
@chk map(mapi, Gray(Ufixed8(0.6))) 0.6f0
@test_throws ErrorException ScaleMinMax(Float32, 0, 0, 1.0) # issue #245
A = [Gray{Float64}(0.2)]
mapi = ScaleMinMax(RGB{Ufixed8}, A, 0.0, 0.2)
@test map(mapi, A) == [RGB{Ufixed8}(1,1,1)]

# ScaleSigned
mapi = ScaleSigned(Float32, 1/5)
@chk map(mapi, 7) 1.0f0
@chk map(mapi, 5) 1.0f0
@chk map(mapi, 3) convert(Float32, 3/5)
@chk map(mapi, -3) convert(Float32, -3/5)
@chk map(mapi, -6) -1.0f0
mapi = ScaleSigned(RGB24, 1.0f0/10)
@chk map(mapi, 12) RGB24(0x00ff00ff)
@chk map(mapi, -10.0) RGB24(0x0000ff00)
@chk map(mapi, 0) RGB24(0x00000000)

# ScaleAutoMinMax
mapi = ScaleAutoMinMax()
A = [100,550,1000]
@chk map(mapi, A) ufixed8([0.0,0.5,1.0])
mapi = ScaleAutoMinMax(RGB24)
@chk map(mapi, A) RGB24[0x00000000, 0x00808080, 0x00ffffff]

# scaling, ssd
img = Images.grayim(fill(typemax(Uint16), 3, 3))
mapi = Images.mapinfo(Ufixed8, img)
img8 = map(mapi, img)
@assert all(img8 .== typemax(Ufixed8))
A = 0
mnA, mxA = 1.0, -1.0
while mnA > 0 || mxA < 0
    A = randn(3,3)
    mnA, mxA = extrema(A)
end
offset = 30.0
img = convert(Images.Image, A .+ offset)
mapi = Images.ScaleMinMax(Ufixed8, offset, offset+mxA, 1/mxA)
imgs = map(mapi, img)
@assert minimum(imgs) == 0
@assert maximum(imgs) == 1
@assert eltype(imgs) == Ufixed8
imgs = Images.imadjustintensity(img, [])
@test_throws ErrorException Images.imadjustintensity(img, [1])
mnA = minimum(A)
@assert Images.ssd(imgs, (A.-mnA)/(mxA-mnA)) < eps()
A = reshape(1:9, 3, 3)
B = map(Images.ClampMin(Float32, 3), A)
@assert eltype(B) == Float32 && B == [3 4 7; 3 5 8; 3 6 9]
B = map(Images.ClampMax(Uint8, 7), A)
@assert eltype(B) == Uint8 && B == [1 4 7; 2 5 7; 3 6 7]

A = reinterpret(Ufixed8, [convert(UInt8,1):convert(UInt8,24);], (3, 2, 4))
img = reinterpret(RGB{Ufixed8}, A, (2,4))
@test separate(img) == permutedims(A, (2,3,1))

# sc
arr = zeros(4,4)
arr[2,2] = 0.5
@assert sc(arr)[2,2] == 0xffuf8
@assert sc(arr, 0.0, 0.75)[2,2] == 0xaauf8

# color conversion
gray = collect(linspace(0.0,1.0,5)) # a 1-dimensional image
gray8 = round(Uint8, 255*gray)
gray32 = Uint32[convert(UInt32, g)<<16 | convert(UInt32, g)<<8 | convert(UInt32, g) for g in gray8]
imgray = Images.Image(gray, Dict{ASCIIString,Any}([("colordim",0), ("colorspace","Gray")]))
buf = map(Images.mapinfo(Uint32, imgray), imgray) # Images.uint32color(imgray)
@test buf == gray32
rgb = RGB{Float64}[RGB(g, g, g) for g in gray]
buf = map(Images.mapinfo(Uint32, rgb), rgb) # Images.uint32color(rgb)
@test buf == gray32
r = red(rgb)
@test r == gray
img = Images.Image(reinterpret(RGB24, gray32)) # , ["colordim"=>0, "colorspace"=>"RGB24"])
buf = map(Images.mapinfo(Uint32, img), img) # Images.uint32color(img)
@test buf == gray32
rgb = repeat(gray, outer=[1,3])
img = Images.Image(rgb, Dict{ASCIIString,Any}([("colordim",2), ("colorspace","RGB"), ("spatialorder",["x"])]))
buf = map(Images.mapinfo(Uint32, img), img) # Images.uint32color(img)
@test buf == gray32
g = green(img)
@test g == gray
rgb = repeat(gray', outer=[3,1])
img = Images.Image(rgb, Dict{ASCIIString,Any}([("colordim",1), ("colorspace","RGB"), ("spatialorder",["x"])]))
buf = map(Images.mapinfo(Uint32, img), img) # Images.uint32color(img)
@test buf == gray32
b = blue(img)
@test b == gray

# map and indexed images
img = Images.ImageCmap([1 2 3; 3 2 1], [RGB{Ufixed16}(1.0,0.6,0.4), RGB{Ufixed16}(0.2, 0.4, 0.6), RGB{Ufixed16}(0.5,0.5,0.5)])
mapi = MapNone(RGB{Ufixed8})
imgd = map(mapi, img)
cmap = [RGB{Ufixed8}(1.0,0.6,0.4), RGB{Ufixed8}(0.2, 0.4, 0.6), RGB{Ufixed8}(0.5,0.5,0.5)]
@test imgd == reshape(cmap[[1,3,2,2,3,1]], (2,3))


module TrackError

using ..Util, ..Tagging, ..Condition, ..Traits
import ..Condition
export applytrackerr, condvar, tagerror

# Issue is that type of err value depends on
# gradient method used
# Solutions
# Specialize: condvar/applytrackerr based on type of Omega

"""
`applytrackerr(x, ω, initerr)`

Is `ω` in the domain of `x`?
Returns `(f(ω), v)` where `v = ω ∈ dom(x)`
"""
function applytrackerr(x, ω, initerr)
  ω_ = tagerror(ω, initerr) # zt: surprisingly expensive
  fx = x(ω_)
  (fx = fx, err = ω_.tags.err.val)
end

# tagnotrackerr(ω) = tag(ω, (donttrack = true,))
# applynotrackerr(x, ω) = apl(x, tagnotrackerr(ω))

# "Soft `condvar`: distance from `ω` to the domain of `x`"
# condvarₛ(x, ω, initerr = softtrue()) = applytrackerr(x, ω, initerr).err
# condvarₛ(x) = ω -> condvarₛ(x, ω)

"Is `ω` in the domain of `x`?"
condvar(x, ω, initerr = true) = applytrackerr(x, ω, initerr).err
condvar(x) = ω -> condvar(x, ω)

"Update `err` by conjoining cu`rrent error with `b`"
conjoinerror!(err::Box, b) = err.val &= b

dotrack(ω) = !haskey(ω.tags, :donttrack) && haskey(ω.tags, :err)

function Condition.condf(::trait(Err), ω, x, y)
  dotrack(ω) && conjoinerror!(ω.tags.err, y(ω))
  x(ω)
end

function Condition.cond!(::trait(Err), ω, bool)
  dotrack(ω) && conjoinerror!(ω.tags.err, bool)
  bool
end

"Tag `ω` with `err`"
tagerror(ω, initerr) = tag(ω, (err = Box(initerr),))

# tagerror(ω, initerr::AbstractBool) = tag(ω, (err = Ref(initerr),)) #FIXME UNCOMMENT THIS AND SPECIALISE BELOW TO OMEGA
# "Tag tω with error"
# function tagerror(tω::TaggedΩ, initerr::AbstractBool)
#   # If error already there, use that!
#   haskey(tω.tags, :err) ? tω : tag(tω, (err = Ref{Real}(initerr),))
# end

function combinetags(::Type{Val{:err}}, a, b)
  # @show a
  # @show b
  # @assert false
  @show a[]
  @show a[] & b[]
end

end


#
# Date created: 2022-03-07
# Author: aradclif
#
#
############################################################################################
# Experiments in efficiency improvements
using Base.Cartesian

A = rand(4,3,5);
Da = axes(A)
Db = Da
dims = (1,3)
Db = Base.setindex(Db, 1, dims[1])
Db = Base.setindex(Db, 1, dims[2])
B = similar(A, Db);
fill!(B, 0.0);
sum(B)
@macroexpand @nloops 3 i A begin
    (@nref 3 B d -> length(Db[d]) == 1 ? 1 : i_d) =
        +((@nref 3 B d -> length(Db[d]) == 1 ? 1 : i_d), (@nref 3 A i))
end
@macroexpand @nloops 3 i A begin
    (@nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d) =
        +((@nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d), (@nref 3 A i))
end
@macroexpand @nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d
@macroexpand @nexprs 3 d -> (j_d = length(axes(B, d)) == 1 ? 1 : i_d)
@macroexpand @nref 3 A i
Expr(:ref, :A, :(1), :i_2, :(3))
(LoopVectorization.StaticInt(1), LoopVectorization.StaticInt(3), LoopVectorization.StaticInt(1))

@generated function prepare(A, dims::NTuple{N, Int}) where {N}
    quote
        Da = axes(A)
        Db = Da
        @nexprs $N d -> (Db = Base.setindex(Db, Base.OneTo(1), dims[d]))
        Db
    end
end

@generated function vreduce!(f::Function, B::AbstractArray{T, N}, A::AbstractArray{T, N}) where {T, N}
    quote
        @nloops $N i A begin
            (@nref $N B d -> length(axes(B, d)) == 1 ? 1 : i_d) =
                f((@nref $N B d -> length(axes(B, d)) == 1 ? 1 : i_d), (@nref $N A i))
        end
        B
    end
end

# @generated function _vreduce!(f::F, B::AbstractArray{T, N}, A::AbstractArray{T, N}) where {F, T, N}
#     quote
#         ee = @macroexpand @nloops $N i A begin
#             (@nref $N B d -> length(axes(B, d)) == 1 ? 1 : i_d) =
#                 $(f.instance)((@nref $N B d -> length(axes(B, d)) == 1 ? 1 : i_d), (@nref $N A i))
#         end
#         ea = ee.args[2]
#         return quote
#             $ea
#             return B
#         end
#     end
#     # ee = quote
#     #     @macroexpand @nloops 3 i A begin
#     #     (@nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d) =
#     #         +((@nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d), (@nref 3 A i))
#     #     end
#     # end
#     # ee.args[2]
# end

function preparedims(B::AbstractArray{T, N}, A::AbstractArray{T, N}) where {T, N}
    Db = axes(B)
    Dᵥ = ntuple(d -> length(Db[d]) == 1 ? StaticInt(1) : Symbol(:i_, d), N)
end
# preparedims(B, A)

# function vvreduce!(f::F, B::AbstractArray{T, N}, A::AbstractArray{T, N}) where {F, T, N}
#     dims = preparedims(B, A)
#     _vvreduce!(f, B, A, dims)
# end
# @generated function _vvreduce!(f::F, B::AbstractArray{T, N}, A::AbstractArray{T, N}, dims::D) where {F, T, N, D}
#     e = Expr(:ref, :B, ntuple(d -> D.parameters[d] !== StaticInt{1} ? Symbol(:i_, d) : :(1), N)...)
#     quote
#         ex = @macroexpand @nloops $N i A begin
#             $e = $(F.instance)($e, (@nref $N A i))
#         end
#         loops = ex.args[2]
#         return quote
#             # @turbo $loops
#             $loops
#             return B
#         end
#         # @turbo $loops
#         # return B
#     end
# end
# _vvreduce!(+, B, A, preparedims(B, A))
function branches_quote(F, N::Int, D)
    e = Expr(:ref, :B, ntuple(d -> D.parameters[d] !== StaticInt{1} ? Symbol(:i_, d) : :(1), N)...)
    # or:
    # e = Expr(:ref, :B)
    # params = D.parameters
    # for d ∈ 1:N
    #     push!(e.args, params[d] !== StaticInt{1} ? Symbol(:i_, d) : :(1))
    # end
    ex = :(@nloops $N i A begin
               $e = $(F.instance)($e, (@nref $N A i))
           end)
    exm = macroexpand(Main, ex)
    loops = exm.args[2]
    return quote
        @turbo $loops
        return B
    end
end
@generated function vvvreduce!(f::F, B::AbstractArray{T, N}, A::AbstractArray{T, N}, dims::D) where {F, T, N, D}
    branches_quote(F, N, D)
end
function vvvreduce(f, A::AbstractArray{T}; dims::NTuple{N, Int}) where {N} where {T}
    if N == ndims(A)
        cat(_vvvreduce(f, A); dims=N)
    else
        Db = prepare(A, dims)
        # B = similar(A, Db)
        # fill!(B, 0.0)
        B = zeros(T, Db)
        vvvreduce!(f, B, A, preparedims(B, A))
    end
end
# vvvreduce(f, A::AbstractArray{T}; dims::Int) where {T} = vvvreduce(f, A; dims=(dims,))
@generated function _vvvreduce(f::F, A::AbstractArray{T, N}) where {F, T, N}
    quote
        s = zero($T)
        @turbo for i ∈ eachindex(A)
            s = $(F.instance)(s, A[i])
        end
        return s
    end
end
@timev vvvreduce(+, A)
A = rand(4,3,5);

# branches_quote(typeof(+), 3, typeof(preparedims(B, A)))
# dd = preparedims(B, A)
# @timev vvvreduce!(+, B, A, dd)

@benchmark vvvreduce(+, A, dims=dims)
@benchmark sum(A, dims=dims)
@timev vvvreduce(+, A, dims=(1,))
# This causes an error -- however, one can actually perform this faster by just
# reducing using a scalar, then doing cat(s, dims=N).
B = zeros(Float64, prepare(A, (1,)))
tp = typeof(+)
DD = typeof(preparedims(B, A))
@timev branches_quote(tp, 3, DD)
@timev branches_quote2(tp, 3, DD)
eval(ans)
@timev sum(A, dims=(5,));
@timev vvvreduce(+, A, dims=(5,));

@generated function indexexpr(D::NTuple{N, T}) where {N, T}
    ex = quote
        @ntuple $N d -> length(D[d]) == 1 ? Symbol(1) : Symbol("i_$(d)")
    end
    # return :($ex)
    ex
end
Isym = indexexpr(Db2)
@nref 3 B j -> ($(Isym[j]))

Db2 = prepare(A, dims)
B2 = similar(A, Db2);
fill!(B2, 0.0);
vreduce!(+, B2, A);
_vreduce!(+, B2, A)
vreduce6!(+, B2, A)
vreduce2!(+, B2, A);

s = 0.0
ex = quote @nloops 3 i A begin
    s += @nref 3 A i
end
end
ex2 = :(@nloops 3 i A begin
    s += @nref 3 A i
end)

@macroexpand @turbo for i ∈ eachindex(A)
    s = +(s, A[i])
end

LoopVectorization.@_turbo(ex)

ex = quote
    for i ∈ eachindex(A)
        s += A[i]
    end
end
ex2 = :(for i ∈ eachindex(A)
            s += A[i]
        end)

N = 3
static_dims = [1,2,3]
quote
    @turbo $loops
end
quote
    @turbo $loop3
end

loop3 = @macroexpand @nloops 3 i A begin
    (@nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d) =
        +((@nref 3 B d -> length(axes(B, d)) == 1 ? 1 : i_d), (@nref 3 A i))
end
ee = loop3.args[2]
quote
    @turbo $ee
end

function simdsum(A::Array{T, 3}) where {T}
    s = zero(T)
    @turbo for k ∈ axes(A, 3)
        nothing
        for j ∈ axes(A, 2)
            for i ∈ axes(A, 1)
                s = (+)(s, A[i, j, k])
            end
        end
    end
    s
end

simdsum(A)