module BrainFuck
using SparseArrays
export @bf_str, compile, interpret

"""

    bracketindices(tokens)

Returns a Dict of the indices of square brackets and their corresponding matching bracket.
"""
function bracketindices(tokens)
    stack = []
    bracketpairs = Dict{Int,Int}()
    for (index, token) in Iterators.enumerate(tokens)
        if token == '['
            push!(stack, index)
        elseif token == ']'
            if isempty(stack)
                throw(ArgumentError("Unmatched ']' at position $(index)."))
            end
            otherindex = pop!(stack)
            push!(bracketpairs, index => otherindex, otherindex => index)
        end
    end
    if !isempty(stack)
        throw(ArgumentError(
            "Unmatched '[' at position$(length(stack) > 1 ? "s" : "") $(join(stack, ", ", " and "))."
        ))
    end
    bracketpairs
end

"""

    clean(code::AbstractString)

Strips out all non-Brainfuck characters from the given string.
"""
clean(code::AbstractString) = replace(code, r"[^\-+,.<>[\]]" => "")

"""

    interpret(code; kwargs...)

Interpret the given string as Brainfuck `code`.

## Keyword arguments

- `memorysize::Integer` How many bytes to allocate for memory, default is `30000`.
- `input::Union{AbstractString,IO}` The string or buffer which the `,` command reads from.
    Defaults to `stdin`.
- `output::Union{IO,Type{String}}` The buffer which the `.` command write to. Returns the
    output as a string if `String` is passed as to `output`. Defaults to `stdout`.
"""
function interpret(code;
    memorysize::Integer = 30000,
    input::Union{AbstractString,IO} = stdin,
    output::Union{IO,Type{String}} = stdout
)
    tokens = clean(code)
    brackets = bracketindices(tokens)
    memory = spzeros(UInt8, memorysize)
    pointer = 1
    i = 1
    if input isa AbstractString
        input = IOBuffer(input)
    end
    returnstring = output isa Type{String}
    if returnstring
        output = IOBuffer()
    end
    while i <= length(tokens)
        token = tokens[i]
        if token == '-'
            memory[pointer] -= 0x01
        elseif token == '+'
            memory[pointer] += 0x01
        elseif token == '<'
            pointer = max(1, pointer - 1)
        elseif token == '>'
            pointer = min(memorysize, pointer + 1)
        elseif token == ','
            if !eof(input)
                skipchars(c -> c == 0x0d && peek(input) == 0x0a, input)
                memory[pointer] = read(input, UInt8)
            end
        elseif token == '.'
            print(output, Char(memory[pointer]))
        elseif token == '['
            if iszero(memory[pointer])
                i = brackets[i]
            end
        elseif token == ']'
            if !iszero(memory[pointer])
                i = brackets[i]
            end
        end
        i += 1
    end
    if returnstring
        String(take!(output))
    end
end

"""

    compile(code)

Compiles the BrainFuck code to a Julia function. The returned function has the same
keywords arguments as `interpret(code)`.
"""
compile(code) = eval(expr(code))

"""

    bf"+++++++++++++++++++++."

Defines a BrainFuck program.

## Example

```jldoctest
julia> helloworld = bf"\""
    ++++++++[>++++[>++>+++>+++>+<<<<-]>+>+>->>+[<]<-]>>.>
    ---.+++++++..+++.>>.<-.<.+++.------.--------.>>+.>++.
"\"";

julia> helloworld()
Hello World!

julia> helloworld(output = String)
"Hello World!\\n"

```
"""
macro bf_str(code)
    :(compile($code))
end

"""
Compiles the BrainFuck code to a Julia expr which will `eval` into a function.
"""
function expr(code)
    stack = [quote
        if input isa AbstractString
            input = IOBuffer(input)
        end
        returnstring = output isa Type{String}
        if returnstring
            output = IOBuffer()
        end
        pointer = 1
        memory = spzeros(UInt8, memorysize)
    end]
    tokens = Iterators.Stateful(clean(code))
    for token in tokens
        if occursin(token, "+-")
            i = token == '+' ? 0x01 : -0x01
            while occursin(peek(tokens), "+-")
                i += (popfirst!(tokens) == '+' ? 0x01 : -0x01)
            end
            push!(stack[end].args, :(memory[pointer] += $i))
        elseif occursin(token, "<>")
            i = token == '>' ? 1 : -1
            while occursin(peek(tokens), "<>")
                i += (popfirst!(tokens) == '>' ? 1 : -1)
            end
            push!(stack[end].args, :(pointer = clamp(pointer + $i, 1:memorysize)))
        elseif token == ','
            push!(stack[end].args, :(
                if !eof(input)
                    skipchars(c -> c == 0x0d && peek(input) == 0x0a, input)
                    memory[pointer] = read(input, UInt8)
                end
            ))
        elseif token == '.'
            push!(stack[end].args, :(print(output, Char(memory[pointer]))))
        elseif token == '['
            loop = Expr(:while, :(!iszero(memory[pointer])), Expr(:block))
            push!(stack[end].args, loop)
            push!(stack, loop.args[end])
        elseif token == ']'
            if length(stack) == 1
                throw(ArgumentError("Unmatched ']'."))
            end
            pop!(stack)
        end
    end
    if length(stack) != 1
        throw(ArgumentError("Unmatched '['."))
    end
    push!(stack[end].args, :(
        if returnstring
            String(take!(output))
        end
    ))
    Expr(
        :function,
        :((;
            memorysize::Integer = 30000,
            input::Union{AbstractString,IO} = stdin,
            output::Union{IO,Type{String}} = stdout
        )),
        only(stack)
    )
end

end # module


struct GymEnv{T,Ta,To,P} <: AbstractEnv
    pyenv::P
    observation_space::To
    action_space::Ta
    state::P
end
export GymEnv

mutable struct AtariEnv{IsGrayScale,TerminalOnLifeLoss,N,S<:AbstractRNG} <: AbstractEnv
    ale::Ptr{Nothing}
    name::String
    screens::Tuple{Array{UInt8,N},Array{UInt8,N}}  # for max-pooling
    actions::Vector{Int}
    action_space::Base.OneTo{Int}
    observation_space::Space{Array{ClosedInterval{UInt8},N}}
    noopmax::Int
    frame_skip::Int
    reward::Float32
    lives::Int
    rng::S
end
export AtariEnv

mutable struct OpenSpielEnv{S,G} <: AbstractEnv
    state::S
    game::G
end
export OpenSpielEnv

# ??? can we safely ignore the `game` field here
Base.hash(e::OpenSpielEnv, h::UInt) = hash(e.state, h)
Base.:(==)(e::OpenSpielEnv, ee::OpenSpielEnv) = e.state == ee.state

mutable struct SnakeGameEnv{A,N,G} <: AbstractEnv
    game::G
    latest_snakes_length::Vector{Int}
    latest_actions::Vector{CartesianIndex{2}}
    is_terminated::Bool
end
export SnakeGameEnv


# test the structs
@info "Testing the FT and NaN of the structs...";
@testset "FT and NaN --- Types" begin
    for FT in [Float32, Float64]
        node   = SPACSimple{FT}();
        cont1L = SPACContainer1L{FT}();
        cont2L = SPACContainer2L{FT}();

        for data in [node, cont1L, cont2L]
            @test FT_test(data, FT);
            @test NaN_test(data);
        end
    end
end

# test and benchmark the big_leaf_partition!
println();
@info "Testing the big_leaf_partition!...";
@testset "SoilPlantAirContinuum --- big_leaf_partition!" begin
    for FT in [Float32, Float64]
        node   = SPACSimple{FT}();
        zenith = FT(30);
        r_all  = FT(1000);

        big_leaf_partition!(node, zenith, r_all);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test the gain_risk_map
println();
@info "Testing the gain_risk_map Function...";
@testset "SoilPlantAirContinuum --- gain_risk_map" begin
    for FT in [Float32, Float64]
        node   = SPACSimple{FT}();
        photo  = C3CLM(FT);
        zenith = FT(30);
        r_all  = FT(1000);

        big_leaf_partition!(node, zenith, r_all);
        mat = gain_risk_map(node, photo);
        @test FT_test(mat, FT);
        @test NaN_test(mat);
    end
end

# test and benchmark the leaf_gas_exchange_nonopt!
println();
@info "Testing the leaf_gas_exchange_nonopt! Functions...";
@testset "SoilPlantAirContinuum --- leaf_gas_exchange_nonopt!" begin
    for FT in [Float32, Float64]
        node   = SPACSimple{FT}();
        photo  = C3CLM(FT);
        zenith = FT(30);
        r_all  = FT(1000);
        flow   = FT(4);
        f_sl   = FT(2.5);
        f_sh   = FT(1.5);

        big_leaf_partition!(node, zenith, r_all);
        leaf_gas_exchange_nonopt!(node, photo, flow);
        @test FT_test(node, FT);
        @test NaN_test(node);
        leaf_gas_exchange_nonopt!(node, photo, f_sl, f_sh);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test and benchmark the leaf_gas_exchange!
println();
@info "Testing the leaf_gas_exchange! Functions...";
@testset "SoilPlantAirContinuum --- leaf_gas_exchange!" begin
    for FT in [Float32, Float64]
        node   = SPACSimple{FT}();
        photo  = C3CLM(FT);
        zenith = FT(30);
        r_all  = FT(1000);
        flow   = FT(4);
        f_sl   = FT(2.5);
        f_sh   = FT(1.5);

        big_leaf_partition!(node, zenith, r_all);
        leaf_gas_exchange!(node, photo, flow);
        @test FT_test(node, FT);
        @test NaN_test(node);
        leaf_gas_exchange!(node, photo, f_sl, f_sh);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test and benchmark the leaf_temperature*
println();
@info "Testing the leaf_temperature* Functions...";
@testset "SoilPlantAirContinuum --- leaf_temperature*" begin
    for FT in [Float32, Float64]
        node = SPACSimple{FT}();
        rad  = FT(300);
        flow = FT(4);

        for result in [ leaf_temperature(node, rad, flow),
                        leaf_temperature_shaded(node, rad, flow),
                        leaf_temperature_sunlit(node, rad, flow) ]
            @test FT_test(result, FT);
            @test NaN_test(result);
        end
    end
end

# test and benchmark the optimize_flows!
println();
@info "Testing the optimize_flows! Functions...";
@testset "SoilPlantAirContinuum --- optimize_flows!" begin
    for FT in [Float32, Float64]
        node   = SPACSimple{FT}();
        photo  = C3CLM(FT);
        zenith = FT(30);
        r_all  = FT(1000);

        big_leaf_partition!(node, zenith, r_all);
        optimize_flows!(node, photo);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test and benchmark the atmosheric* functions
println();
@info "Testing the atmosheric* Functions...";
@testset "SoilPlantAirContinuum --- atmosheric*" begin
    for FT in [Float32, Float64]
        h = FT(1000);

        for result in [ atmospheric_pressure(h),
                        atmospheric_pressure_ratio(h),
                        ppm_to_Pa(h) ]
            @test FT_test(result, FT);
            @test NaN_test(result);
        end
    end
end

# test and benchmark the zenith_angle
println();
@info "Testing the zenith_angle Functions...";
@testset "SoilPlantAirContinuum --- zenith_angle" begin
    for FT in [Float32, Float64]
        latd = FT(10);
        decd = FT(10);
        lhad = FT(10);
        day  = FT(100)
        hour = FT(13)
        minu = FT(30)

        for result in [ zenith_angle(latd, decd, lhad),
                        zenith_angle(latd, day, hour),
                        zenith_angle(latd, day, hour, minu) ]
            @test FT_test(result, FT);
            @test NaN_test(result);
        end
    end
end

# test and benchmark the annual_profit
println();
@info "Testing the annual_profit Functions...";
@testset "SoilPlantAirContinuum --- annual_profit" begin
    arti = artifact"2020_leaf_invest_weather" *  "/gs_sample.csv";
    weat = DataFrame(CSV.File(arti));
    for FT in [Float32, Float64]
        node    = SPACSimple{FT}();
        photo   = C3CLM(FT);
        weatmat = Matrix{FT}(weat);

        gscp = annual_profit(node, photo, weatmat);
        @test FT_test(gscp, FT);
        @test NaN_test(gscp);
    end
end

# test and benchmark the annual_simulation!
println();
@info "Testing annual_simulation! Functions...";
@testset "SoilPlantAirContinuum --- annual_simulation!" begin
    arti = artifact"2020_leaf_invest_weather" *  "/gs_sample.csv";
    weat = DataFrame(CSV.File(arti));
    for FT in [Float32, Float64]
        node  = SPACSimple{FT}();
        photo = C3CLM(FT);
        df    = create_dataframe(FT, weat);

        annual_simulation!(node, photo, weat, df);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test and benchmark the leaf_allocation!
println();
@info "Testing the leaf_allocation! Functions...";
@testset "SoilPlantAirContinuum --- leaf_allocation!" begin
    for FT in [Float32, Float64]
        node  = SPACSimple{FT}();
        photo = C3CLM(FT);
        laba  = FT(1000);
        vmax  = FT(80);

        leaf_allocation!(node, laba);
        @test FT_test(node, FT);
        @test NaN_test(node);
        leaf_allocation!(node, photo, vmax);
        @test FT_test(node, FT);
        @test NaN_test(node);
        leaf_allocation!(node, photo, laba, vmax);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test and benchmark the optimize_leaf
println();
@info "Testing the optimize_leaf! Functions...";
@testset "SoilPlantAirContinuum --- optimize_leaf!" begin
    arti = artifact"2020_leaf_invest_weather" *  "/gs_sample.csv";
    weat = DataFrame(CSV.File(arti));
    for FT in [Float32, Float64]
        node    = SPACSimple{FT}();
        photo   = C3CLM(FT);
        weatmat = Matrix{FT}(weat);

        optimize_leaf!(node, photo, weatmat);
        @test FT_test(node, FT);
        @test NaN_test(node);
    end
end

# test the function to vary SPACSimple
println();
@info "Testing the vary_spac! Functions...";
@testset "SoilPlantAirContinuum --- vary_spac!" begin
    arti = artifact"2020_leaf_invest_weather" *  "/gs_sample.csv";
    weat = DataFrame(CSV.File(arti));
    facs = ["kl", "kw", "wb", "wc", "wk",
            "cc", "cv", "gm",
            "ga", "sd",
            "ta", "rh", "ca"];
    for FT in [Float32, Float64]
        node = SPACSimple{FT}();
        for _fac in facs
            vary_spac!(node, weat, _fac, FT(1.5));
            @test true;
        end
    end
end

# test the function for mSCOPE version radiation
println();
@info "Testing the mSCOPE verion setup...";
@testset "SoilPlantAirContinuum --- vary_spac!" begin
    for FT in [Float32, Float64]
        node = SPACMono{FT}();
        initialize_spac_canopy!(node);
        layer_fluxes!(node);
        @test NaN_test(node);
    end
end


### BasicMCJob

# BasicMCJob is used for sampling a single parameter via serial Monte Carlo
# It is the most elementary and typical Markov chain Monte Carlo (MCMC) job

mutable struct BasicMCJob <: MCJob
  model::GenericModel # Model of single parameter
  pindex::Integer # Index of single parameter in model.vertices
  parameter::Parameter # Points to model.vertices[pindex] for faster access
  sampler::MCSampler
  tuner::MCTuner
  range::BasicMCRange
  vstate::VariableStateVector # Vector of variable states ordered according to variables in model.vertices
  pstate::ParameterState # Points to vstate[pindex] for faster access
  sstate::MCSamplerState # Internal state of MCSampler
  outopts::Dict # Options related to output
  output::Union{VariableNState, VariableIOStream, Void} # Output of model's single parameter
  count::Integer # Current number of post-burnin iterations
  fmt_iter::Union{Function, Void}
  fmt_perc::Union{Function, Void}
  resetpstate::Bool # If resetpstate=true then pstate is reset by reset(job), else pstate is not modified by reset(job)
  verbose::Bool

  function BasicMCJob(
    model::GenericModel,
    pindex::Integer,
    sampler::MCSampler,
    tuner::MCTuner,
    range::BasicMCRange,
    vstate::VariableStateVector,
    outopts::Dict,
    resetpstate::Bool,
    verbose::Bool,
    check::Bool
  )
    instance = new()

    instance.model = model
    instance.pindex = pindex
    instance.vstate = vstate
    instance.resetpstate = resetpstate

    if check
      checkin(instance)
    end

    instance.outopts = isa(outopts, Dict{Symbol, Any}) ? outopts : convert(Dict{Symbol, Any}, outopts)
    augment_parameter_outopts!(instance.outopts)

    instance.parameter = instance.model.vertices[instance.pindex]
    instance.parameter.states = instance.vstate

    instance.sampler = sampler
    instance.tuner = tuner
    instance.range = range
    instance.verbose = verbose

    instance.pstate = instance.vstate[instance.pindex]
    if any(isnan, instance.pstate.value)
      if instance.parameter.pdf != nothing
        instance.pstate.value = rand(instance.parameter.pdf)
      elseif instance.parameter.prior != nothing
        instance.pstate.value = rand(instance.parameter.prior)
      else
        error("Not possible to initialize pstate with missing pstate.value and without parameter.pdf or parameter.prior")
      end
    end
    if value_support(instance.parameter) == Continuous &&
      instance.parameter.diffopts != nothing &&
      instance.parameter.diffopts.mode == :reverse
      set_tapes!(instance.parameter, instance.pstate)
    end
    initialize!(instance.pstate, instance.parameter, instance.sampler, instance.outopts)
    if value_support(instance.parameter) == Continuous
      instance.parameter.state = instance.pstate
    end

    instance.sstate = sampler_state(
      instance.parameter, instance.sampler, tuner, instance.pstate, instance.vstate, instance.outopts[:diagnostics]
    )

    if isa(instance.sampler, MuvAMWG)
      initialize_diagnosticvalues!(instance.pstate, instance.sstate)
    end

    instance.output = initialize_output(instance.pstate, range.npoststeps, instance.outopts)

    instance.count = 0

    instance.fmt_iter =
      if instance.tuner.verbose
        if isa(sampler, NUTS) && isa(instance.tuner, DualAveragingMCTuner)
          format_iteration(ndigits(instance.tuner.nadapt))
        else
          format_iteration(ndigits(range.burnin))
        end
      else
        nothing
      end

    instance.fmt_perc = instance.tuner.verbose ? format_percentage() : nothing

    instance
  end
end

BasicMCJob(
  model::GenericModel,
  sampler::MCSampler,
  range::BasicMCRange,
  v0::VariableStateVector;
  pindex::Integer=findfirst(v::Variable -> isa(v, Parameter), model.vertices),
  tuner::MCTuner=VanillaMCTuner(),
  outopts::Dict=Dict(:destination=>:nstate, :monitor=>[:value], :diagnostics=>Symbol[]),
  resetpstate::Bool=true,
  verbose::Bool=false,
  check::Bool=false
) =
  BasicMCJob(model, pindex, sampler, tuner, range, v0, outopts, resetpstate, verbose, check)

function BasicMCJob(
  model::GenericModel,
  sampler::MCSampler,
  range::BasicMCRange,
  v0::Dict{Symbol, S};
  pindex::Integer=findfirst(v::Variable -> isa(v, Parameter), model.vertices),
  tuner::MCTuner=VanillaMCTuner(),
  outopts::Dict=Dict(:destination=>:nstate, :monitor=>[:value], :diagnostics=>Symbol[]),
  resetpstate::Bool=true,
  verbose::Bool=false,
  check::Bool=false
) where S<:VariableState
  vstate = Array{S}(length(v0))
  for (k, v) in v0
    vstate[model.ofkey[k]] = v
  end
  BasicMCJob(model, pindex, sampler, tuner, range, vstate, outopts, resetpstate, verbose, check)
end

function BasicMCJob(
  model::GenericModel,
  sampler::MCSampler,
  range::BasicMCRange,
  v0::Vector;
  pindex::Integer=findfirst(v::Variable -> isa(v, Parameter), model.vertices),
  tuner::MCTuner=VanillaMCTuner(),
  outopts::Dict=Dict(:destination=>:nstate, :monitor=>[:value], :diagnostics=>Symbol[]),
  resetpstate::Bool=true,
  verbose::Bool=false,
  check::Bool=false
)
  vstate = default_state(model.vertices, v0, [outopts], [pindex])
  BasicMCJob(model, pindex, sampler, tuner, range, vstate, outopts, resetpstate, verbose, check)
end

function BasicMCJob(
  model::GenericModel,
  sampler::MCSampler,
  range::BasicMCRange,
  v0::Dict;
  pindex::Integer=findfirst(v::Variable -> isa(v, Parameter), model.vertices),
  tuner::MCTuner=VanillaMCTuner(),
  outopts::Dict=Dict(:destination=>:nstate, :monitor=>[:value], :diagnostics=>Symbol[]),
  resetpstate::Bool=true,
  verbose::Bool=false,
  check::Bool=false
)
  vstate = Array{Any}(length(v0))
  for (k, v) in v0
    vstate[model.ofkey[k]] = v
  end

  BasicMCJob(
    model,
    sampler,
    range,
    vstate,
    pindex=pindex,
    tuner=tuner,
    outopts=outopts,
    resetpstate=resetpstate,
    verbose=verbose,
    check=check
  )
end

function reset(job::BasicMCJob)
  reset!(job.sstate, job.pstate, job.parameter, job.sampler, job.tuner)

  if isa(job.output, VariableIOStream)
    reset(job.output)
    mark(job.output)
  end

  job.count = 0
end

function reset(job::BasicMCJob, x)
  reset!(job.pstate, x, job.parameter, job.sampler)
  reset(job)
end

function save(job::BasicMCJob)
  write(job.output, job.pstate)
  if job.outopts[:flush]
    flush(job.output)
  end
end

save(job::BasicMCJob, i::Integer) = copy!(job.output, job.pstate, i)

function run(job::BasicMCJob)
  fmt_iter::Union{Function, Void} = job.verbose ? format_iteration(ndigits(job.range.nsteps)) : nothing

  if isa(job.output, VariableIOStream)
    mark(job.output)
  end

  for i = 1:job.range.nsteps
    if job.verbose
      println("Iteration ", fmt_iter(i), " of ", job.range.nsteps)
    end

    iterate!(job, typeof(job.sampler), variate_form(job.parameter))

    if in(i, job.range.postrange)
      job.count += 1

      if job.output != nothing
        if isa(job.output, VariableNState)
          save(job, job.count)
        elseif isa(job.output, VariableIOStream)
          save(job)
        else
          error("To save output, :destination must be set to :nstate or :iostream, got $(job.outopts[:destination])")
        end
      end
    end
  end

  if isa(job.output, VariableIOStream)
    close(job.output)
  end
end

function checkin(job::BasicMCJob)
  pindex = find(v::Variable -> isa(v, Parameter), job.model.vertices)
  np = length(pindex)

  if np == 0
    error("The model does not have a parameter, but a BasicMCJob requires a parameter")
  elseif np == 1
    if pindex[1] != job.pindex
      error("Parameter located in model.vertices[$(pindex[1])], but pindex = $(job.pindex)")
    end
  elseif np >= 2
    warn("The model has $np parameters, but a BasicMCJob requires exactly one parameter")
    if in(job.pindex, pindex)
      error("Indices of parameters in model.vertices do not contain pindex")
    end
  end

  nv = num_vertices(job.model)
  nvstate = length(job.vstate)

  if nv != nvstate
    warn("Number of variables ($nv) not equal to number of variable states ($nvstate)")
  end

  pstate = job.vstate[job.pindex]

  if !isa(pstate, ParameterState)
    error("The parameter's state must be saved in a ParameterState subtype, got $(typeof(pstate)) state type")
  else
    check_support(job.model.vertices[job.pindex], pstate)
  end
end

output(job::BasicMCJob) = job.output

function show(io::IO, job::BasicMCJob)
  indentation = "  "

  println(io, "BasicMCJob:")
  print(io, indentation)
  show(io, job.parameter)
  print(io, "\n"*indentation)
  show(io, job.model)
  print(io, "\n"*indentation)
  show(io, job.sampler)
  print(io, "\n"*indentation)
  show(io, job.tuner)
  print(io, "\n"*indentation)
  show(io, job.range)
end


"""
`Model`

A struct representing a model.

**Fields**
- `variable_names`: Variable names (e.g. `["X", "Y", "Z"]`).
- `parameter_names`: Parameter names (e.g. `["v", "d1", "d2"]`).
- `problem`: `ODEProblem` or `SDEProblem` or `JumpProblem`.
- `solver_algorithm`: Solver used to solve the `problem`` (e.g. `SOSRI()`).
- `solver_parameters`: Parameters passed to the solver (e,g, `(saveat=0.1,)`)
- `input`: Input to the model representedas a tuple with two elements. The first
    element is a matrix representing the input events and the second element is
    a string that specifies, which parameter is manipulated by the events.
- `output`: Function that transforms `DifferentialEqualtions` solution
    into a matrix. Rows are time steps and columns are variables.

**Functions**
- `set_initial_conditions`
- `set_timespan`
- `create_model`
- `set_solver`
- `get_parameter_index`
- `get_parameter_value`
- `set_parameter_value`
- `set_input`
- `simulate_population`
"""
struct Model
    variable_names::Vector{String}
    parameter_names::Vector{String}
    problem::Union{ODEProblem, SDEProblem, JumpProblem}
    solver_algorithm::Any
    solver_parameters::NamedTuple
    input::Tuple{Matrix{Float64}, String}
    output::Function
end

"""
`_set_model_property(model::Model; kwargs...)`

Create a new copy of `model` with replaced fields as specified by the keyword
    arguments.
"""
function _set_model_property(model::Model; kwargs...)

    # Iterate all possible fields of Model struct
    fields = []
    for field in fieldnames(Model)

        # If the field is in kwargs, use it, otherwise use the original value
        if field in keys(kwargs)
            push!(fields, deepcopy(kwargs[field]))
        else
            push!(fields, deepcopy(getproperty(model, field)))
        end

    end

    # Create a model with new fields
    return Model(fields...)

end

function _get_problem_property(problem, property_name)

    # Property name must be a symbol
    property_name = Symbol(property_name)

    # For JumpProblem access problem.prob, otherwise problem directly
    if problem isa JumpProblem
        property = getproperty(problem.prob, property_name)
    else
        property = getproperty(problem, property_name)
    end

    # Return the found property
    return property

end

function _get_problem_property(model::Model, property_name)

    return _get_problem_property(model.problem, property_name)

end

function _set_problem_property(model::Model; kwargs...)

    # Remake problem with the new properties
    new_problem = remake(model.problem; kwargs...)

    # Replace the problem in the model
    new_model = _set_model_property(model, problem=new_problem)

    return new_model
end

"""
`set_initial_conditions(model::Model, u0)`

Set initial conditions for a model.
"""
function set_initial_conditions(model::Model, u0)
    return _set_problem_property(model; u0=u0)
end

"""
`set_timespan(model::Model, tspan)`

Set timespan for a model.
"""
function set_timespan(model::Model, tspan)
    return _set_problem_property(model; tspan=tspan)
end

function set_output(model::Model, output)
    return _set_model_property(model; output=output)
end

"""
`create_model`

Create a model based on the `DifferentialEqualtions` problem.

**Arguments**
- `variable_names`: Variable names (e.g. `["X", "Y", "Z"]`).
- `parameter_names`: Parameter names (e.g. `["v", "d1", "d2"]`).
- `problem`: `ODEProblem` or `SDEProblem` or `JumpProblem`.

**Returns**
- `model`: Model represented as a `Model` struct.
"""
function create_model(variable_names, parameter_names, problem)

    # Check that variable names match the length of the initial state
    u0 = _get_problem_property(problem, "u0")
    if length(variable_names) != length(u0)
        error("Incorrect number of variable names!")
    end

    # Check that parameter names match the length of the parameter vector
    p = _get_problem_property(problem, "p")
    if length(parameter_names) != length(p)
        error("Incorrect number of parameter names!")
    end

    # Set up a default algorithm based on the problem type
    if problem isa ODEProblem
        solver_algorithm = Tsit5()
    elseif problem isa SDEProblem
        solver_algorithm = SOSRI()
    elseif problem isa JumpProblem
        solver_algorithm = SSAStepper()
    else
        error("Unknown problem type!")
    end

    # Set up the other default parameters the do not depend on the problem type
    solver_parameters = (saveat=0.1,)
    input = (Matrix{Float64}(undef, 0, 0), "")
    output = sol -> Matrix(sol[:, :]')

    # Create the Model struct
    model = Model(
        variable_names,
        parameter_names,
        problem,
        solver_algorithm,
        solver_parameters,
        input,
        output
    )

    return model

end

"""
`set_solver(model::Model, algorithm=nothing; merge_kwargs=true, kwargs...)`

Set solver algorithm and its properties.

**Arguments**
- `model`: Model.

**Optional Arguments**
- `algorithm`: Solver algorithm (e.g. `Tsit5()` or `SOSRI()`).

**Keyword Arguments**
- `merge_kwargs`: If `false`, the original arguments are deleted.
- `kwargs...`: Solver parameters.

**Returns**
- `new_model`: Model with updated solver algorithm and parameters.
"""
function set_solver(model::Model, algorithm=nothing; merge_kwargs=true, kwargs...)

    if :callback in keys(kwargs)
        error("Callback is reserved for the input function!")
    end

    if isnothing(algorithm)
        # If the algorithm was not specified, use to original one
        solver_algorithm = model.solver_algorithm
    else
        # If the algorithm was specified, use the new one
        solver_algorithm = algorithm
    end

    if length(kwargs) > 0
        # If any kwargs were passed, add them to the solver_parameters

        if merge_kwargs
            # Merge kwargs with the existing parameters
            solver_parameters = merge(model.solver_parameters, NamedTuple(kwargs))
        else
            # Replace the old parameters with the new ones
            if :callback in keys(model.solver_parameters)
                # Preserve callback
                callback = model.solver_parameters[:callback]
                solver_parameters = merge(NamedTuple(kwargs), (callback=callback,))
            else
                solver_parameters = NamedTuple(kwargs)
            end
        end

    else
        # If no kwargs were passed, just copy the original ones
        solver_parameters = model.solver_parameters

    end

    # Build a new model
    new_model = _set_model_property(model; solver_algorithm=solver_algorithm,
        solver_parameters=solver_parameters)

    return new_model

end

"""
Set callback to the solver parameters.
"""
function _set_callback(model::Model, callback)

    # Merge callback with existing solver parameters
    solver_parameters = merge(model.solver_parameters, (callback=callback,))

    # Build a new model
    new_model = _set_model_property(model; solver_parameters=solver_parameters)

    return new_model

end

"""
`get_parameter_index(model::Model, parameter_name::String)`

Find index of the specified parameter.

**Arguments**
- `model`: Model.
- `parameter_name`: Name of the parameter.

**Returns**
- `index`: Index indicating the position of the parameter.
"""
function get_parameter_index(model::Model, parameter_name::String)

    parameter_names = model.parameter_names
    index = findfirst(parameter_names .== parameter_name)
    if isnothing(index)
        error("Parameter '$(parameter_name)' is not in the model!")
    end
    return index

end

"""
`get_parameter_value(model::Model, parameter_name::String)`

Get a parameter value.

**Arguments**
- `model`: Model.
- `parameter_name`: Name of the parameter.

**Returns**
- `parameter_value`: Value of the parameter.
"""
function get_parameter_value(model::Model, parameter_name::String)

    index = get_parameter_index(model, parameter_name)
    parameter_values = _get_problem_property(model, "p")
    return parameter_values[index]

end

"""
`set_parameter_value(model::Model, name::String, value::Number)`

Set a parameter value.

**Arguments**
- `model`: Model.
- `name`: Name of the parameter. Can also be in the form of "d1=d2=d3", which
    sets all parameters d1, d2, and d3 to the same value.
- `value`: New value for the parameter.

**Returns**
- `new_model`: Copy of `model` with updated parameter value.
"""
function set_parameter_value(model::Model, name::String, value::Number)

    # Get array of parameter values
    parameter_values = _get_problem_property(model, "p")

    if occursin("=", name)
        # If "=" occurs in `name`, split the string into individual parameters
        name_array = String.(split(name, "="))
        for name_split in name_array
            # Set each parameter to `value`
            index = get_parameter_index(model, name_split)
            parameter_values[index] = value
        end
    else
        # Set the desired parameter to `value`
        index = get_parameter_index(model, name)
        parameter_values[index] = value
    end

    # Rebuild the model with the new parameter value
    new_model = _set_problem_property(model, p=parameter_values)
    return new_model

end

"""
`model::Model, name_array::Array, value_array::Array`

Set parameter values.

**Arguments**
- `model`: Model.
- `name_array`: Parameter names. Can also be in the form of "d1=d2=d3", which
    sets all parameters d1, d2, and d3 to the same value.
- `value_array`: New values for the parameters.

**Returns**
- `new_model`: Copy of `model` with updated parameter values.
"""
function set_parameter_value(model::Model, name_array::Array, value_array::Array)

    # The input arrays must have the same length
    if length(name_array) != length(value_array)
        error("The array of parameter names and values must have the same length!")
    end

    # Set the parameters
    new_model = deepcopy(model)
    for (name, value) in zip(name_array, value_array)
        new_model = set_parameter_value(new_model, name, value)
    end

    return new_model

end

"""
`_create_callback(events::Matrix, i)`

Generate `PresetTimeCallback` from `events` modyfing `i`-th parameter.
"""
function _create_callback(events::Matrix, i)

    p = NaN  # holds the original value of the parameter

    # Initialization at the beginning of the integration
    initialize = function (c, u, t, integrator)

        # Protect the original parameters from overwriting
        integrator.p = copy(integrator.p)

        # Save the default parameter
        p = integrator.p[i]

        # Set the parameter to 0
        integrator.p[i] = 0.0

    end

    # Set the parameter to its default value at the beginning of the event
    tstops_on = events[:, 1]
    affect_on! = (integrator) -> integrator.p[i] = p
    callback_on = PresetTimeCallback(tstops_on, affect_on!,
        save_positions=(false, false), initialize=initialize)

    # Set the parameter to zero at the end of the event
    tstops_off = events[:, 2]
    affect_off! = (integrator) -> integrator.p[i] = 0.0
    callback_off = PresetTimeCallback(tstops_off, affect_off!,
        save_positions=(false, false))

    return CallbackSet(callback_on, callback_off)

end

"""
`_events_to_function(events::Matrix)`

Convert `events` to a function.
"""
function _events_to_function(events::Matrix)
    N = size(events, 1)
    function fun(t)
        for i = 1:N
            if events[i, 1] <= t < events[i, 2]
                return 1.0
            end
        end
        return 0.0
    end
    return fun
end

"""
`_create_discrete_callback(events::Matrix, i)`

Generate `DiscreteCallback` from `events` modyfing `i`-th parameter.
"""
function _create_discrete_callback(events::Matrix, i)

    p = NaN  # holds the original value of the parameter

    # Initialization at the beginning of the integration
    initialize = function (c, u, t, integrator)

        # Protect the original parameters from overwriting
        integrator.p = copy(integrator.p)

        # Save the default parameter
        p = integrator.p[i]

        # Set the parameter to 0
        integrator.p[i] = 0.0

    end

    # Convert the events to a function
    fun = _events_to_function(events)

    # Call the function at every step of the integration
    affect! = (integrator) -> integrator.p[i] = fun(integrator.t)*p
    condition = (u, t, integrator) -> true
    callback = DiscreteCallback(condition, affect!; initialize=initialize,
        save_positions=(false, false))

    return callback

end

"""
`set_input(model::Model, events::Matrix, parameter_name="I")`

Set input to the model.

**Arguments**
- `model`: Model.
- `events`: Matrix representing the input square function.
- `parameter_name`: Parameter that is being modified by the input.
"""
function set_input(model::Model, events::Matrix, parameter_name="I")

    parameter_index = get_parameter_index(model, parameter_name)

    # Create a callback
    if model.problem isa Union{ODEProblem, SDEProblem}
        # Use PresetTimeCallback for ODE and SDE models
        callback = _create_callback(events, parameter_index)
    elseif model.problem isa JumpProblem
        # Use DiscreteCallback for Jump models
        callback = _create_discrete_callback(events, parameter_index)
    end

    # Build a new model with the specified input
    new_model = _set_callback(model, callback)
    new_model = _set_model_property(new_model, input=(events, parameter_name))

    return new_model

end

"""
`simulate_population(model::Model, trajectories=1; save_trajectories=true,
    show_progress=false)`

Simulate a population.

**Arguments**
- `model`: Model.

**Optional Arguments**
- `trajectories`: Number of trajectories to simulate.

**Keyword Arguments**
- `save_trajectories`: If `false`, only mean is saved.
- `show_progress`: If `true`, show a progress bar in the terminal.
"""
function simulate_population(model::Model, trajectories=1;
    save_trajectories=true, show_progress=false)

    # Make sure that the original model is not modified
    model = deepcopy(model)

    # Prepare variables for PopulationSolution fields
    t = Vector{Float64}(undef, 0)  # time
    m = Matrix{Float64}(undef, 0, 0)  # mean
    U = Array{Float64, 3}(undef, 0, 0, 0)  # trajectories
    events = model.input[1]  # events

    # Initialize the porgress meter
    if show_progress
        progress_meter = ProgressMeter.Progress(trajectories; barlen=20)
    end

    # Simulate trajectories
    success_arr = fill(false, trajectories)  # success for individual trajectories
    lk = ReentrantLock()
    Threads.@threads for i = 1:trajectories

        model2 = deepcopy(model)

        # Solve the problem
        sol2 = solve(model2.problem, model2.solver_algorithm; model2.solver_parameters...)

        # Check if the simulation was successful
        if model2.problem isa Union{ODEProblem, SDEProblem}
            success_arr[i] = sol2.retcode == :Success
        elseif model2.problem isa JumpProblem
            success_arr[i] = sol2.retcode == :Default
        else
            error("Unknown model type!")
        end

        # Save the solution if the current simulation was successful
        if success_arr[i]

            lock(lk) do

                # Save time
                if isempty(t)
                    t = sol2.t
                end

                # Map solution to a matrix
                x = model2.output(sol2)

                # Add the trajectory contribution to the mean
                if isempty(m)
                    # Initialize the mean vector
                    m = (x ./ trajectories)
                else
                    # Add solution to the mean vector
                    m .+= (x ./ trajectories)
                end

                # Save the individual trajectory
                if save_trajectories
                    if isempty(U)
                        # Initialize the matrix for individual trajectories
                        nsamples = size(x, 1)
                        nvariables = size(x, 2)
                        U = fill(NaN, nsamples, nvariables, trajectories)
                    end
                    # Add the current trajectory to the trajectory list
                    U[:, :, i] = x
                end

            end

        else
            # Stop the simulation if the current simulation was unsuccessful
            break
        end

        if show_progress
            lock(lk) do
                ProgressMeter.next!(progress_meter)
            end
        end

    end

    # Build the output struct
    success = all(success_arr)
    solution = PopulationSolution(t, m, U, events, success)

    return solution

end

"""
`scan`

Scan parameters of a model.

**Arguments**
- `model`: Model.
- `parameter_names`: Parameter names to scan.
- `parameter_values`: Parameter values for each parameter to scan over.
- `summary_function`: A function that takes in a model, evaluates it, and
    returns some parameters. If the function is called without any arguments,
    it returns the names of the parameters as a vector of strings. For example
    ```
    julia> summary_function(model)
    [1.1, 2.2, 3.3]
    julia> summary_function()
    ["Period", "Amplitude", "Phase"]
    ```

**Keyword Arguments**
- `show_progress`: If `true`, show a progress bar in the terminal.

**Returns**
- `df`: Dataframe with results.
"""
function scan(model::Model, parameter_names, parameter_values,
    summary_function; show_progress=false)

    # Protect the original model from overwriting
    model = deepcopy(model)

    # Find all possible parameter combinations
    parameter_value_combinations = find_all_combinations(parameter_values)
    n = size(parameter_value_combinations, 1)

    # Initialize summary matrix
    summary_names = summary_function()
    scan_summary = fill(NaN, n, length(summary_names))

    # Initialize progress bar
    if show_progress
        progressmeter = ProgressMeter.Progress(n; barlen=20)
    end

    # Iterate all parameter value combinations
    lk = ReentrantLock()
    Threads.@threads for i = 1:n

        # Set model parameters
        pmodel = set_parameter_value(model, parameter_names, parameter_value_combinations[i, :])

        # Evaluate the sumary function
        pmodel_summary = summary_function(pmodel)

        # Save the model summary to the overall matrix
        scan_summary[i, :] = pmodel_summary

        # Update progress bar
        if show_progress
            lock(lk) do
                ProgressMeter.next!(progressmeter)
            end
        end

    end

    # Build the output dataframe
    matrix = hcat(parameter_value_combinations, scan_summary)
    names = vcat(parameter_names, summary_names)
    df = DataFrame(matrix, names)

    return df

end

"""
`estimate_binary_arnold_tongue`

Use binary search to estimate arnold tongue for a model.

**Arguments**
- `model`: Model.
- `periods`: A two-element array like `[minimal_period, maximal_period]`.
- `amplitudes`: Input amplitudes to scan.

**Keyword Arguments**
- `trajectories`: Number of trajectories in a population.
- `i_solution`: Number of the state to use to estimate the period.
- `n_lags`: Number of lags used to estimate the autocorrelation functions.
- `input_parameter_name`: Name of the parameter that control the input amplitude.
- `min_time`: Minimal time to consider (to remove the effect of initial conditions).
- `max_time`: Maximal time of the simulation.
- `show_progress`: If true, show a progress bar in the terminal.
- `period_error`: Maximal allowed output period deviation from the input period.
- `boundary_error`: Maximal allowed error to estimate the boundary between the
    entrained and unentrained regions.
- `maximal_error_attempts`: How many attempts should be given to an errored simulation.
"""
function estimate_binary_arnold_tongue(model::Model, periods, amplitudes;
    trajectories=1, i_solution=1, n_lags=1000, input_parameter_name="I",
    min_time=nothing, max_time=nothing, show_progress=false, period_error=0.01,
    boundary_error=0.01, maximal_error_attempts=0)

    # Protect the original model from overwriting
    model = deepcopy(model)

    # Set the maximal time of the simulation
    if !isnothing(max_time)
        model = set_timespan(model, (0.0, max_time))
    end

    # Get minimal and maximal amplitude and the number of periods to scan
    minimal_amplitude = amplitudes[1]
    maximal_amplitude = amplitudes[end]
    n_periods = length(periods)

    # Each element is the minimal amplitude of entrainment for the given period
    entrainment_amplitudes = fill(NaN, n_periods)

    # Initialize a progress bar
    if show_progress
        progressmeter = ProgressMeter.Progress(n_periods; barlen=20)
    end

    # Iterate all period-amplitude combinations
    lk = ReentrantLock()
    Threads.@threads for i = 1:n_periods  # Threads.@threads

        # Set the forcing (input) period
        forcing_period = periods[i]

        # Create events for the current forcing period
        if model.problem isa JumpProblem
            end_time = model.problem.prob.tspan[2]
        else
            end_time = model.problem.tspan[2]
        end
        nevents = end_time / forcing_period + 1
        events = create_events([(:DD, forcing_period*0.5),
            (:LD, nevents, forcing_period*0.5, forcing_period*0.5)])

        # Set offset (solution before this time is discarted)
        if isnothing(min_time)
            offset = end_time*0.5
        else
            offset = min_time
        end

        # pmodel = deepcopy(model)
        pmodel = set_input(model, events, input_parameter_name)

        """
        This function will be optimized using the binary boundary search. It
        simulates the model for the given forcing amplitude and returns `true`
        if the model is entrained and `false` if the model is unentrained.
        """
        function fun(forcing_amplitude)

            # Set the forcing amplitude and events to the model
            pmodel = set_parameter_value(pmodel, input_parameter_name, forcing_amplitude)

            # Simulate the population
            solution = PopulationSolution()
            success = false
            counter = 0
            while !success
                try
                    # lock(lk) do
                    solution = simulate_population(pmodel, trajectories;
                        save_trajectories=false)
                    # end
                    success = true
                catch err
                    counter += 1
                    if counter >= maximal_error_attempts
                        @warn "[LAST] Simulation threw an error: $err"
                        solution = PopulationSolution()
                        success = true
                    else
                        @warn "[$counter] Simulation threw an error: $err"
                    end
                end
            end

            # If the simulation was successful, estimate the period
            if solution.success

                # Discard the simulation start
                solution = select_time(solution, mintime=offset)
                t = solution.time
                m = solution.mean[:, i_solution]

                R_period, _ = estimate_period(t, m; n_lags=n_lags)

                return abs(R_period - forcing_period) <= period_error

            else

                return false

            end

        end

        # Perform binary search
        entrainment_amplitudes[i] = binary_boundary_search(fun,
            [minimal_amplitude, maximal_amplitude], boundary_error)

        # Update the progress bar
        if show_progress
            ProgressMeter.next!(progressmeter)
        end

    end

    # Create the output dataframe
    arnold_tongue = DataFrame(periods=periods, amplitudes=entrainment_amplitudes)

    return arnold_tongue

end


function action_info(p::POMCPPlanner, b; tree_in_info=false)
    local a::actiontype(p.problem)
    info = Dict{Symbol, Any}()
    try
        tree = POMCPTree(p.problem, b, p.solver.tree_queries)
        a = search(p, b, tree, info)
        p._tree = tree
        if p.solver.tree_in_info || tree_in_info
            info[:tree] = tree
        end
    catch ex
        # Note: this might not be type stable, but it shouldn't matter too much here
        # a = convert(actiontype(p.problem), default_action(p.solver.default_action, p.problem, b, ex))
        a = convert(actiontype(p.problem), default_action(RandomPolicy(p.problem, rng=MersenneTwister(1)), p.problem, b, ex))
        info[:exception] = ex
    end
    return a, info
end

action(p::POMCPPlanner, b) = first(action_info(p, b))

function search(p::POMCPPlanner, b, t::POMCPTree, info::Dict)
    all_terminal = true
    nquery = 0
    start_us = CPUtime_us()
    info[:tree_depth] = 0
    dep_count=0
    for i in 1:p.solver.tree_queries
        nquery += 1
        if CPUtime_us() - start_us >= 1e6*p.solver.max_time
            break
        end
        s = rand(p.rng, b)
        if !POMDPs.isterminal(p.problem, s)
            reward, depth = simulate(p, s, POMCPObsNode(t, 1), p.solver.max_depth)
            info[:tree_depth] = (info[:tree_depth] * dep_count + depth) / (dep_count + 1)
            dep_count+=1
            # info[:tree_depth] = max(info[:tree_depth], depth)
            all_terminal = false
        end
    end
    info[:search_time_us] = CPUtime_us() - start_us
    info[:tree_queries] = nquery

    if all_terminal
        throw(AllSamplesTerminal(b))
    end

    h = 1
    best_node = first(t.children[h])
    best_v = t.v[best_node]
    @assert !isnan(best_v)
    for node in t.children[h][2:end]
        if t.v[node] >= best_v
            best_v = t.v[node]
            best_node = node
        end
    end

    return t.a_labels[best_node]
end

solve(solver::POMCPSolver, pomdp::POMDP) = POMCPPlanner(solver, pomdp)

function simulate(p::POMCPPlanner, s, hnode::POMCPObsNode, steps::Int)
    dep = 1.0
    if steps == 0 || isterminal(p.problem, s)
        return 0.0, dep
    end
    t = hnode.tree
    h = hnode.node

    ltn = log(t.total_n[h])
    best_nodes = empty!(p._best_node_mem)
    best_criterion_val = -Inf
    for node in t.children[h]
        n = t.n[node]
        if n == 0 && ltn <= 0.0
            criterion_value = t.v[node]
        elseif n == 0 && t.v[node] == -Inf
            criterion_value = Inf
        else
            criterion_value = t.v[node] + p.solver.c*sqrt(ltn/n)
        end
        if criterion_value > best_criterion_val
            best_criterion_val = criterion_value
            empty!(best_nodes)
            push!(best_nodes, node)
        elseif criterion_value == best_criterion_val
            push!(best_nodes, node)
        end
    end
    ha = rand(p.rng, best_nodes)
    a = t.a_labels[ha]

    sp, o, r = @gen(:sp, :o, :r)(p.problem, s, a, p.rng)

    hao = get(t.o_lookup, (ha, o), 0)
    if hao == 0
        dep = dep + 1.0
        hao = insert_obs_node!(t, p.problem, ha, sp, o)
        v = estimate_value(p.solved_estimator,
                           p.problem,
                           sp,
                           POMCPObsNode(t, hao),
                           steps-1)
        R = r + discount(p.problem)*v
    else
        reward, depth = simulate(p, sp, POMCPObsNode(t, hao), steps-1)
        dep = depth + 1
        R = r + discount(p.problem) * reward
    end

    t.total_n[h] += 1
    t.n[ha] += 1
    t.v[ha] += (R-t.v[ha])/t.n[ha]

    return R, dep
end


module DataModel

abstract AbstractDataModel

export AbstractDataModel, getgrad, getllik, getfullgrad, get_MAP, getfullhess
function getgrad(dm::AbstractDataModel;nobs=0)
    # should return a function grad(x) that calculates the gradient at x
    error("getgrad not implemented for AbstractDataModel")
end

function getllik(dm::AbstractDataModel;nobs=0)
    # should return a function llik(x) that calculates the loglikelihood at x.
    error("getllik not implemented for AbstractDataModel")
end

function getfullgrad(dm::AbstractDataModel;nobs=0)
    # should return a function grad(x) that calculates the gradient at x
    error("getfullgrad not implemented for AbstractDataModel")
end

function getfullhess(dm::AbstractDataModel;nobs=10)
  # should return a function hess(x) that calculates the hessian at x
  error("getfullgrad not implemented for AbstractDataModel")
end

function getN(dm::AbstractDataModel)
  # should return number of data items
  error("getN not implemented for AbstractDataModel")
end

end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form 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 copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 COPYRIGHT HOLDER OR CONTRIBUTORS 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.
# *********************************************************************************
"""
    Financial

Financial data struct with inner constructor:
```julia
function Financial(;
    om_cost_escalation_pct::Float64 = 0.025,
    elec_cost_escalation_pct::Float64 = 0.023,
    offtaker_tax_pct::Float64 = 0.26,
    offtaker_discount_pct = 0.083,
    third_party_ownership::Bool = false,
    owner_tax_pct::Float64 = 0.26,
    owner_discount_pct::Float64 = 0.083,
    analysis_years::Int = 25,
    VoLL::Union{Array{R,1}, R} where R<:Real = 1.00,
    microgrid_premium_pct::Float64 = 0.3,
    macrs_five_year::Array{Float64,1} = [0.2, 0.32, 0.192, 0.1152, 0.1152, 0.0576],  # IRS pub 946
    macrs_seven_year::Array{Float64,1} = [0.1429, 0.2449, 0.1749, 0.1249, 0.0893, 0.0892, 0.0893, 0.0446]
)
```

!!! note
    When `third_party_ownership` is `false` the offtaker's discount and tax percentages are used throughout the model:
    ```julia
        if !third_party_ownership
            owner_tax_pct = offtaker_tax_pct
            owner_discount_pct = offtaker_discount_pct
        end
    ```
"""
struct Financial
    om_cost_escalation_pct::Float64
    elec_cost_escalation_pct::Float64
    offtaker_tax_pct::Float64
    offtaker_discount_pct
    third_party_ownership::Bool
    owner_tax_pct::Float64
    owner_discount_pct::Float64
    analysis_years::Int
    VoLL::Union{Array{R,1}, R} where R<:Real
    microgrid_premium_pct::Float64
    macrs_five_year::Array{Float64,1}
    macrs_seven_year::Array{Float64,1}

    function Financial(;
        om_cost_escalation_pct::Float64 = 0.025,
        elec_cost_escalation_pct::Float64 = 0.023,
        offtaker_tax_pct::Float64 = 0.26,
        offtaker_discount_pct = 0.083,
        third_party_ownership::Bool = false,
        owner_tax_pct::Float64 = 0.26,
        owner_discount_pct::Float64 = 0.083,
        analysis_years::Int = 25,
        VoLL::Union{Array{R,1}, R} where R<:Real = 1.00,
        microgrid_premium_pct::Float64 = 0.3,
        macrs_five_year::Array{Float64,1} = [0.2, 0.32, 0.192, 0.1152, 0.1152, 0.0576],  # IRS pub 946
        macrs_seven_year::Array{Float64,1} = [0.1429, 0.2449, 0.1749, 0.1249, 0.0893, 0.0892, 0.0893, 0.0446]
    )
        if !third_party_ownership
            owner_tax_pct = offtaker_tax_pct
            owner_discount_pct = offtaker_discount_pct
        end

        return new(
            om_cost_escalation_pct,
            elec_cost_escalation_pct,
            offtaker_tax_pct,
            offtaker_discount_pct,
            third_party_ownership,
            owner_tax_pct,
            owner_discount_pct,
            analysis_years,
            VoLL,
            microgrid_premium_pct,
            macrs_five_year,
            macrs_seven_year
        )
    end
end


module Tracking

    using CUDA
    using DocStringExtensions
    using GNSSSignals
    using LoopVectorization
    using StaticArrays
    using StructArrays
    using TrackingLoopFilters

    using Acquisition: AcquisitionResults
    using Unitful: upreferred, Hz, dBHz, ms
    import Base.zero, Base.length, Base.resize!, LinearAlgebra.dot

    export get_early, get_prompt, get_late
    export get_early_index,get_prompt_index, get_late_index
    export get_correlator
    export get_accumulator, get_accumulators
    export get_num_accumulators
    export get_carrier_doppler, get_carrier_phase
    export get_code_doppler, get_code_phase
    export get_correlator_sample_shifts
    export get_early_late_sample_spacing
    export get_early_late_index_shift
    export get_secondary_code_or_bit_found
    export get_correlator_carrier_phase, get_correlator_carrier_frequency
    export get_state
    export get_system
    export get_cn0
    export track
    export get_bits
    export get_num_bits

    export TrackingState
    export NumAnts
    export NumAccumulators
    export MomentsCN0Estimator
    export AbstractCN0Estimator
    export EarlyPromptLateCorrelator
    #export VeryEarlyPromptLateCorrelator
    export SecondaryCodeOrBitDetector
    export GainControlledSignal

    struct NumAnts{x}
    end

    NumAnts(x) = NumAnts{x}()

    struct NumAccumulators{x}
    end

    NumAccumulators(x) = NumAccumulators{x}()

    include("code_replica.jl")
    include("carrier_replica.jl")
    include("downconvert.jl")
    include("cn0_estimation.jl")
    include("discriminators.jl")
    include("bit_buffer.jl")
    include("correlator.jl")
    include("secondary_code_or_bit_detector.jl")
    include("tracking_state.jl")
    include("tracking_results.jl")
    include("tracking_loop.jl")
    include("gpsl1.jl")
    include("gpsl5.jl")
    include("galileo_e1b.jl")
    include("boc.jl")
    include("downconvert_and_correlate.jl")
end


function read_model_configuration(pm::PredictionModel{BMTMKLModel}, data::Dict{String, Any})
	if haskey(data, "model_configuration")
		mc_data = data["model_configuration"]
		# general setting
		if haskey(mc_data, "alpha") && haskey(mc_data, "beta") && haskey(mc_data, "mu") && haskey(mc_data, "sigma2")
			return ModelConfiguration(pm, mc_data["alpha"], mc_data["beta"], mc_data["mu"], mc_data["sigma2"])
		else
			mc = ModelConfiguration(pm)
			# refined settings if present
			if haskey(mc_data, "bias_precision_alpha") && haskey(mc_data, "bias_precision_beta")
				mc.parameters["α_ɣ"] = mc_data["bias_precision_alpha"]
				mc.parameters["β_ɣ"] = mc_data["bias_precision_beta"]
			end
			if haskey(mc_data, "weights_precision_alpha") && haskey(mc_data, "weights_precision_beta")
				mc.parameters["α_λ"] = mc_data["weights_precision_alpha"]
				mc.parameters["β_λ"] = mc_data["weights_precision_beta"]
			end
			if haskey(mc_data, "outcome_precision_alpha") && haskey(mc_data, "outcome_precision_beta")
				mc.parameters["α_ε"] = mc_data["outcome_precision_alpha"]
				mc.parameters["β_ε"] = mc_data["outcome_precision_beta"]
			end
			if haskey(mc_data, "intermed_results_precision_alpha") && haskey(mc_data, "intermed_results_precision_beta")
				mc.parameters["α_ν"] = mc_data["intermed_results_precision_alpha"]
				mc.parameters["β_ν"] = mc_data["intermed_results_precision_beta"]
			end
			if haskey(mc_data, "kernel_weights_precision_alpha") && haskey(mc_data, "kernel_weights_precision_beta")
				mc.parameters["α_⍵"] = mc_data["kernel_weights_precision_alpha"]
				mc.parameters["β_⍵"] = mc_data["kernel_weights_precision_beta"]
			end
			if haskey(mc_data, "bias_mean") && haskey(mc_data, "bias_sigma2")
				mc.parameters["μ_b"] = mc_data["bias_mean"]
				mc.parameters["σ_b"] = mc_data["bias_sigma2"]
			end
			if haskey(mc_data, "kernel_weights_mean") && haskey(mc_data, "kernel_weights_sigma2")
				mc.parameters["μ_e"] = mc_data["kernel_weights_mean"]
				mc.parameters["σ_e"] = mc_data["kernel_weights_sigma2"]
			end
			if haskey(mc_data, "weights_mean") && haskey(mc_data, "weights_sigma2")
				mc.parameters["μ_a"] = mc_data["weights_mean"]
				mc.parameters["Σ_a"] = mc_data["weights_sigma2"]
			end
			if haskey(mc_data, "intermed_results_mean") && haskey(mc_data, "intermed_results_sigma2")
				mc.parameters["μ_g"] = mc_data["intermed_results_mean"]
				mc.parameters["Σ_g"] = mc_data["intermed_results_sigma2"]
			end
		end
	end
	# standard settings if no config present
	ModelConfiguration(pm)
end

function init!(m::PredictionModel{BMTMKLModel}, ic::InferenceConfiguration)
	(base_kernels, pathway_specific_kernels) = compute_all_kernels(m.data, collect(values(m.data.cell_lines)), subsume_pathways = ic.subsume_pathways)
	m.precomputations["base_kernels"] = base_kernels
	m.precomputations["pathway_specific_kernels"] = pathway_specific_kernels
end

function update_model_config!(pm::PredictionModel{BMTMKLModel}, mc::ModelConfiguration)
	mc.parameters["T"] = length(pm.data.drugs)
	mc.parameters["N"] = OrderedDict{Drug, Int64}()
	for drug in collect(keys(pm.data.results))
		# N[t] is the number of cell lines for which we have a result and that are in the training set
		mc.parameters["N"][drug] = length(filter((cl) -> !cl.in_test_set && haskey(pm.data.results[drug].outcome_values, cl), collect(values(pm.data.cell_lines))))
	end
	# recompute overall number of kernels (+4 accomodates for combined kernels)
	mc.parameters["K"] = length(pm.precomputations["base_kernels"]) + sum(map(length, collect(values(pm.precomputations["pathway_specific_kernels"])))) + 4
end

function post_init!(m::PredictionModel{BMTMKLModel})
    # collect outcomes statistics, only on results used for training
    for (t, drug) in enumerate(keys(m.data.results))
        vals = Float64[]
        [if !cl.in_test_set push!(vals, m.data.results[drug].outcome_values[cl]) end for cl in keys(m.data.results[drug].outcome_values)]
        m.data.results[drug].outcome_mean = mean(vals)
        m.data.results[drug].outcome_std = stdm(vals, m.data.results[drug].outcome_mean)
    end
    all_drugs = collect(keys(m.data.results)) # the tasks
    cell_lines = collect(values(m.data.cell_lines)) # all cell lines in the experiment, excluding held-out data

    # the kernels between cell lines
    kernels = DataStructures.OrderedDict{Drug, Vector{Matrix{Float64}}}()
    # the target values
    targets = DataStructures.OrderedDict{Drug, Vector{Float64}}()

    # cross kernels between training data and test data
    cross_kernels = DataStructures.OrderedDict{Drug, Vector{Matrix{Float64}}}()
    # target values for test data
    test_targets = DataStructures.OrderedDict{Drug, Vector{Float64}}()

    for (t, drug) in enumerate(all_drugs)
        #init
        kernels[drug] = Vector{Matrix{Float64}}()
        cross_kernels[drug] = Vector{Matrix{Float64}}()

        # determine cell lines available for this drug
        result_cell_lines = filter((cl) ->  !cl.in_test_set && haskey(m.data.results[drug].outcome_values, cl), cell_lines)
        test_result_cell_lines = filter((cl) -> cl.in_test_set && haskey(m.data.results[drug].outcome_values, cl), cell_lines)

        # get their indicex in original cell line array
        idx_in_cell_lines = findall((in)(result_cell_lines), cell_lines)
        test_idx_in_cell_lines = findall((in)(test_result_cell_lines), cell_lines)

        #construct the kernels
        for v in m.data.views
            push!(kernels[drug], m.precomputations["base_kernels"][v][idx_in_cell_lines, idx_in_cell_lines])
            push!(cross_kernels[drug], m.precomputations["base_kernels"][v][test_idx_in_cell_lines, idx_in_cell_lines])
            if length(m.data.pathway_information) != 0
                for pw_kernel in m.precomputations["pathway_specific_kernels"][v]
                    push!(kernels[drug], pw_kernel[idx_in_cell_lines, idx_in_cell_lines])
                    push!(cross_kernels[drug], pw_kernel[test_idx_in_cell_lines, idx_in_cell_lines])
                end
            end
        end
        # compute additional kernels
        gene_expression = m.precomputations["base_kernels"][GeneExpression][idx_in_cell_lines, idx_in_cell_lines]
        methylation = m.precomputations["base_kernels"][Methylation][idx_in_cell_lines, idx_in_cell_lines]
        cnv = m.precomputations["base_kernels"][CNV][idx_in_cell_lines, idx_in_cell_lines]

        push!(kernels[drug], gene_expression .* methylation)
        push!(kernels[drug], gene_expression .* cnv)
        push!(kernels[drug], cnv .* methylation)
        push!(kernels[drug], gene_expression .* methylation .* cnv)

        cross_gene_expression = m.precomputations["base_kernels"][GeneExpression][test_idx_in_cell_lines, idx_in_cell_lines]
        cross_methylation = m.precomputations["base_kernels"][Methylation][test_idx_in_cell_lines, idx_in_cell_lines]
        cross_cnv = m.precomputations["base_kernels"][CNV][test_idx_in_cell_lines, idx_in_cell_lines]

        push!(cross_kernels[drug], cross_gene_expression .* cross_methylation)
        push!(cross_kernels[drug], cross_gene_expression .* cross_cnv)
        push!(cross_kernels[drug], cross_cnv .* cross_methylation)
        push!(cross_kernels[drug], cross_gene_expression .* cross_methylation .* cross_cnv)

        train_outcome = map(cl -> m.data.results[drug].outcome_values[cl], result_cell_lines)
        test_outcome = map(cl -> m.data.results[drug].outcome_values[cl], test_result_cell_lines)

        #normalize the outcome data for each drug across the cell lines
        targets[drug] = (train_outcome .- m.data.results[drug].outcome_mean) ./ m.data.results[drug].outcome_std

        # test outcomes, not normalized
        test_targets[drug] = test_outcome
    end

	m.precomputations["kernels"] = kernels
	m.precomputations["cross_kernels"] = cross_kernels
	m.precomputations["targets"] = targets
	m.precomputations["test_targets"] = test_targets
	log_message("done post init")
	nothing
 end

function compute_all_kernels(experiment::Experiment, cell_lines::Vector{CellLine}, cell_lines_test::Union{Vector{CellLine}, Nothing}=nothing; subsume_pathways::Bool=true)
	log_message("in compute kernel method")
	base_kernels = OrderedDict{Type{<:ViewType}, Matrix{Float64}}()
    pathway_specific_kernels = OrderedDict{Type{<:ViewType}, Vector{Matrix{Float64}}}()
    # compute a base kernel for each view, containing all cell lines that are present in all views
    num_views = length(experiment.views)
    for v in experiment.views
        # get views of this type for all cell lines
        # v = experiment.views[v_id]
        data_views = [map_data_views(cell_lines, v)]
        if cell_lines_test != nothing
            push!(data_views, map_data_views(cell_lines_test, v))
        end
        # compute the base kernel for this view
        kct = @elapsed base_kernels[v] = compute_kernel(data_views...)
        if cell_lines_test == nothing
            log_message("computing base kernel for $v $kct seconds)")
        end
        if length(experiment.pathway_information) != 0
            if cell_lines_test == nothing
                log_message("computing pathway specific kernels...")
            end
            pathway_specific_kernels[v] = Vector{Matrix{Float64}}()

            # # pathway specific kernels method 1: compute a kernel for each pathway
            if !subsume_pathways
                num_pws = length(experiment.pathway_information)
                tt = 0
                for pathway in collect(values(experiment.pathway_information))
                    tt += @elapsed push!(pathway_specific_kernels[v], compute_kernel(data_views..., pathway=pathway))
                end

                # method 1 end
            else
                # pathway specific kernels method 2:
                # - for each cell line:
                #   - for each pathway:
                #       - compute average (expression data) or max value (otherwise) over pathway genes
                # - compute cell line similarities as usual
                tt = @elapsed begin
                    cell_line_gene_set_views = map((x) -> Vector{Vector{Float64}}(undef, length(x)), data_views)
                    for (cl_set_idx, cl_set) in enumerate(data_views)
                        for (cl_idx, dv) in enumerate(cl_set)
                            cell_line_gene_set_views[cl_set_idx][cl_idx] = zeros(length(experiment.pathway_information))
                            for (pw_idx, pathway) in enumerate(values(experiment.pathway_information))

                                (m_values, _) = prepare_kernel(dv, dv, pathway)
                                if v in [RNASeq, GeneExpression]
                                    cell_line_gene_set_views[cl_set_idx][cl_idx][pw_idx] = mean(m_values)
                                else
                                    cell_line_gene_set_views[cl_set_idx][cl_idx][pw_idx] = maximum(m_values)
                                end
                            end
                        end
                    end
                    len_views = length(cell_line_gene_set_views)
                    pw_kernel = zeros(length(cell_line_gene_set_views[1]), length(cell_line_gene_set_views[len_views]))
                    for i in 1:size(pw_kernel, 1), j in 1:size(pw_kernel, 2)
                        pw_kernel[i,j] = kernel_function(cell_line_gene_set_views[1][i], cell_line_gene_set_views[len_views][j])
                    end
                    push!(pathway_specific_kernels[v], pw_kernel)
                end
            # # method 2 end
            end
            if cell_lines_test == nothing
                log_message("computing pathway specific kernels for $v took $tt seconds)")
            end
        else
            if cell_lines_test == nothing
                log_message("Note: no pathway information provided")
            end
        end
    end
    if cell_lines_test == nothing
        log_message("done computing kernels")
    end
    (base_kernels, pathway_specific_kernels)
end

function map_data_views(cell_lines::Vector{CellLine}, view_type::Type{<:ViewType})
    ret = Vector{DataView{<:KeyType, <:ViewType}}()
    for cl in cell_lines
        if haskey(cl.views, view_type)
            push!(ret, cl.views[view_type])
        else
            if view_type == RPPA
                push!(ret, DataView{Protein, NAViewType}(cl.id))
            else
                push!(ret, DataView{Gene, NAViewType}(cl.id))
            end
        end
    end
    ret
end

get_complete_cell_lines(experiment::Experiment) = filter(cl -> length(cl.views) == length(experiment.views), collect(values(experiment.cell_lines)))

function compute_kernel(dvs::Vector{DataView}; pathway::Union{Pathway, Nothing} = nothing)
    N = length(dvs)
    k = zeros(N,N)
    prep_time = 0
    comp_time = 0
    kernel_compute_time = @elapsed begin
        Threads.@threads for i in 1:N
            for  j in i:N
                kpt, kct, k[i,j] = compute_kernel_value(dvs[i], dvs[j], pathway)
                prep_time += kpt
                comp_time += kct
                if isnan(k[i,j])
                    k[i,j] = 1e-7
                end
                if i≠j k[j,i] = k[i,j] end
            end
        end
        if !isposdef(k)
            h_k = Array(Hermitian(k))
            # @warn "kernel is not positive definite, using only upper triangular form" isposdef(h_k) sum_diff=sum(abs.(h_k - k))
            k = h_k
        end
    end
    log_message("computed kernel ($(N*(N-1)) computations), preparation time=$prep_time computation time=$comp_time")
    k
end

function compute_kernel(dvs_train::Vector{DataView}, dvs_test::Vector{DataView}; pathway::Union{Pathway, Nothing}=nothing)
    N = length(dvs_train)
    M = length(dvs_test)
    k = zeros(N,M)
    Threads.@threads for i in 1:N
        for j in 1:M
            kpt, kct, k[i,j] = compute_kernel_value(dvs_train[i], dvs_test[j], pathway)
            if isnan(k[i,j])
                k[i,j] = 1e-7
            end
        end
    end
    k
end

function compute_kernel_value(a::DataView, b::DataView, pathway::Union{Pathway, Nothing})
    kpc = @elapsed kernel_params = prepare_kernel(a,b,pathway)
    kcc = @elapsed kfv = kernel_function(kernel_params...)
    kpc, kcc, kfv
end

# general view type vs. NAViewType
function prepare_kernel(a::DataView{K, <:ValueViewType}, b::DataView{K, NAViewType}, pathway::Union{Pathway, Nothing}) where {K<:KeyType}
    keys_to_use = get_common_keys(a, b, pathway)
    (map(m -> a.measurements[m].normalized_value, keys_to_use), zeros(length(keys_to_use)))
end

# NAViewType vs. general view type
function prepare_kernel(a::DataView{K, NAViewType}, b::DataView{K, <:ValueViewType}, pathway::Union{Pathway, Nothing}) where {K<:KeyType}
    (x_prime, x) = prepare_kernel(b,a, pathway)
    (x, x_prime)
end

# RNASeqCall vs. NAViewType
function prepare_kernel(a::DataView{K, RNASeqCall}, b::DataView{K, NAViewType}, pathway::Union{Pathway{K}, Nothing}) where {K<:KeyType}
    keys_to_use = get_common_keys(a, b, pathway)
    (map(m -> a.measurements[m].normalized_value, keys_to_use), [falses(length(keys_to_use))...])
end

# general view type vs. general view type
function prepare_kernel(a::DataView{K, V}, b::DataView{K, V}, pathway::Union{Pathway, Nothing}) where {K<:KeyType, V <:ValueViewType}
    common_keys = get_common_keys(a, b, pathway)
    # hack: when no keys are shared, we have small similarity
    if length(common_keys) == 0
        x = [-2.]
        x_prime = [2.]
        return (x, x_prime  )
    end
    x = map(k -> a.measurements[k].normalized_value, common_keys)
    x_prime = map(k -> b.measurements[k].normalized_value, common_keys)
    (x, x_prime)
end

function prepare_kernel(a::DataView{K, ExomeSeq}, b::DataView{K, ExomeSeq}, pathway::Union{Pathway{K}, Nothing})  where {K<:KeyType}
    common_keys = get_common_keys(a, b, pathway)
    x = Vector{Float64}()
    x_prime = Vector{Float64}()
    factor_exact_matches = 1.
    for k in common_keys
        # check for match in protein change (this happens rarely and we want to emphasize this)
        # if there is match, make sure to include only the ExomeSeq measurement containing it
        measurements_a = a.measurements[k]
        measurements_b = b.measurements[k]

        pcs_a = map(es -> es.protein_change, measurements_a)
        pcs_b = map(es -> es.protein_change, measurements_b)

        protein_change_intersection = intersect(pcs_a, pcs_b)
        if length(protein_change_intersection) > 0
            for pc in protein_change_intersection
                push!(x, get_measurement_value(filter(m -> m.protein_change == pc, measurements_a))...)
                push!(x_prime, get_measurement_value(filter(m -> m.protein_change == pc, measurements_a))...)
                if pc != "None" factor_exact_matches += 1 end
            end
        else
            push!(x, mean(get_measurement_value(a, k)))
            push!(x_prime, mean(get_measurement_value(b, k)))
        end

    end
    if length(common_keys) == 0
        x = [-2.]
        x_prime = [2.]
    end
    (x, x_prime, Float64(length(x) * factor_exact_matches))
end

function prepare_kernel(a::DataView{K, ExomeSeq}, b::DataView{K, NAViewType}, pathway::Union{Pathway{K}, Nothing}) where {K<:KeyType}
    x = map(k -> mean(get_measurement_value(a, k)), [a.used_keys...])
    x_prime = zeros(length(x))
    (x, x_prime)
end

prepare_kernel(a::DataView{K,V}, b::DataView{K,V}, pathway::Union{Pathway, Nothing}) where {K <: KeyType, V <: NAViewType} = ([randn()],[randn()])

function get_common_keys(a::DataView, b::DataView, pathway::Union{Pathway, Nothing})
    # @info "getting common keys" pathway
    if pathway != nothing
        return extract_pathway_keys(a, b, pathway)
    end
    if (b.view_type == NAViewType) || (b.used_keys == a.used_keys)
        return collect(a.used_keys)
    elseif haskey(a.common_keys, b.cell_line_id)
        return a.common_keys[b.cell_line_id]
    elseif haskey(b.common_keys, a.cell_line_id)
        return b.common_keys[a.cell_line_id]
    end
    common_keys = collect(intersect(a.used_keys, b.used_keys))
    a.common_keys[b.cell_line_id] = common_keys
    b.common_keys[a.cell_line_id] = common_keys
    common_keys
end

function extract_pathway_keys(a::DataView, b::DataView, pathway::Pathway)
    # if b.view_type == NAViewType
    #     return intersect(pathway.genes, a.used_keys)
    # end
    # if haskey(a.common_keys, b.cell_line_id)
    #     # @info "found precomputed key intersection"
    #     return intersect(pathway.genes, a.common_keys[b.cell_line_id])
    # end
    # if haskey(b.common_keys, a.cell_line_id)
    #     # @info "found precomputed key intersection"
    #     return intersect(pathway.genes, b.common_keys[a.cell_line_id])
    # end
    common_keys = get_common_keys(a,b,nothing)
    intersect(pathway.genes, common_keys)
end

function extract_pathway_keys(a::DataView{<:KeyType, RPPA}, b::DataView{<:KeyType, RPPA}, pathway::Pathway)
    common_keys = get_common_keys(a,b,nothing)

    keys_to_use = Vector{KeyType}()
    pathway_hgnc_ids = map(k -> k.hgnc_id, pathway.genes)
    hgnc_ids_in_common_keys = map(g -> g.hgnc_id, common_keys)

    hgnc_intersection = intersect(pathway_hgnc_ids, hgnc_ids_in_common_keys)
    key_idx = findall((in)(hgnc_intersection), hgnc_ids_in_common_keys)
    common_keys[key_idx]
end

#actual kernel functions
# Gaussian
kernel_function(x::Vector{Float64}, x_prime::Vector{Float64}, l::Float64 = Float64(length(x))) = exp(-dot(x-x_prime, x-x_prime)/l)
# Jaccard
kernel_function(x::Vector{Bool}, x_prime::Vector{Bool}) = x'*x_prime/(x'*x + x_prime'*x_prime - x'*x_prime)


const stl_machine = (
    function () # Grammar
    # Basic regexp
    object      = re"[A-Za-z][0-9A-Za-z_\-]*"
    space       = re"[\t\n\r ]"
    oct         = re"0o[0-7]+"
    dec         = re"[-+]?[0-9]+"
    hex         = re"0x[0-9A-Fa-f]+"
    prefloat    = re"[-+]?([0-9]+\.[0-9]*|[0-9]*\.[0-9]+)"

    # Complex regexp
    float   = prefloat | re.cat(prefloat | re"[-+]?[0-9]+", re"[eE][-+]?[0-9]+")
    number  = oct | dec | hex | float
    coord   = number × number × number
    vertex  = re"vertex" × coord
    loop    = re"outer" × re"loop" * re.rep(space | vertex) * re"endloop"
    normal  = re"normal" × coord
    facet   = re"facet" × normal × loop × re"endfacet"
    solid   = re"solid" × object * re.rep(space | facet) * re"endsolid" × object
    list    = re.rep(space | solid)

    # RegExp/States Actions
    coord.actions[:enter]   = [:init_coord]
    facet.actions[:exit]    = [:add_facet]
    loop.actions[:enter]    = [:init_loop]
    normal.actions[:exit]   = [:new_facet]
    number.actions[:enter]  = [:mark]
    number.actions[:exit]   = [:store_coord]
    object.actions[:enter]  = [:mark]
    object.actions[:exit]   = [:store_name]
    solid.actions[:exit]    = [:add_solid]
    vertex.actions[:exit]   = [:add_vertex]

    return Automa.compile(list)
end
)()

const stl_actions = Dict(
    :add_facet      => :(push!(facets, facet)),
    :add_solid      => :(), # TODO: handle several objects
    :add_vertex     => :(add_vertex!(facet, coord)),
    :init_coord     => :(coord = Vector{Float64}()),
    :init_loop      => :(),
    :mark           => :(mark = p),
    :new_facet      => :(facet = Facet(coord)),
    :store_coord    => :(push!(coord, round(parse(Float64, data[mark:p - 1]); digits=3))),
    :store_name     => :(name = data[mark:p-1]),
)

const stl_context = Automa.CodeGenContext()

@eval function parse_stl(data::String)
    # Variables to store data
    name = ""
    coord = Vector{Float64}()
    facet = Facet()
    facets = Vector{Facet}()

    # Marks
    mark = 0

    # generate code to initialize variables used by FSM
    $(Automa.generate_init_code(stl_context, stl_machine))

    # set end and EOF positions of data buffer
    p_end = p_eof = sizeof(data)

    # generate code to execute FSM
    $(Automa.generate_exec_code(stl_context, stl_machine, stl_actions))

    design = Design(name, facets)

    # check if FSM properly finished and returm the state of the FSM
    return design, cs == 0 ? :ok : cs < 0 ? :error : :incomplete
end


export LightSource, DistantLight, PointLight, DiskLight

################################
# Visibility Tester is a ray from point to light source

"""
    VisibilityTester(p1::Point3, p2::Point3, eps1::Float64, eps2::Float64)

Returns a ray to test intersection on the line segment from `p1` to `p2` with the given ray
epsilon values.

"""
function VisibilityTester(p1::Point3, p2::Point3, eps1::Float64, eps2::Float64)
    distance = norm(p2-p1)
    return Ray(p1, (p2-p1)/distance, eps1, distance*(1.0 - eps2), 1)
end

"""
    VisibilityTester(origin::Point3, direction::Vector3, eps::Float64)

Returns a ray to test intersection on the line segment from `origin` infinitely far in the
given direction, with a given ray epsilon.

"""
function VisibilityTester(origin::Point3, direction::Vector3, eps::Float64)
    return Ray(origin, direction, eps, Inf, 1)
end

################################
# Point light

"""
    PointLight{S<:Spectrum} <: LightSource

Point light source. It's located at a given point and radiates isotropically. The intensity
value given is interpreted as the intensity in Watts/m^2 at a distance of 1 meter.

# Fields

- `position::Point3`, the location of the light source.
- `intensity::S`, the intensity at a distance of 1 meter in Watts/m^2.
- `light_to_world::Transformation`, transformation to shift the location of the light
   source.
- `world_to_light::Transformation`, inverse transformation of the above.

"""
struct PointLight{S<:Spectrum} <: LightSource
    position::Point3
    intensity::S
    light_to_world::Transformation
    world_to_light::Transformation
end

"""
    PointLight(L::Spectrum, l2w::Transformation)

Create a new `PointLight` with the given spectrum, located at the origin but shifted by
the given `Transformation`.

"""
function PointLight(L::Spectrum, l2w::Transformation)
    PointLight(l2w(Point3(0)), L, l2w, inv(l2w))
end

"""
    sample_L(light::PointLight, p::Point3)

Get an intensity sample from the given `DistantLight` and a ray from the light towards the
given point. The sample is always the same since the light source is a uniform parallel
light.

Returns:

- Intensity `Spectrum`
- Direction `Vector3` of the light.
- A weight coefficient which is always 1.0
- A ray from the light source to the given point.

"""
function sample_L(light::PointLight, p::Point3)
    distance = norm(light.position - p)
    L = light.intensity / distance^2
    pdf = 1.0
    vis = VisibilityTester(p, light.position, 2e-5, 0.0)
    return L, pdf, vis
end

"""
    generate_ray(light::PointLight, bounds::BoundingSphere, S::Sample)

Returns a ray from the light source at the given `BoundingSphere`, parametrized by the
given random sample.

"""
function generate_ray(light::PointLight, bounds::BoundingSphere, S::Sample)
    ray = ray_from_point(bounds.radius, S.imgX, S.imgY)
    return light.light_to_world(ray)
end

################################
# Distanct light source

"""
    DistantLight{S<:Spectrum} <: LightSource

A point light source infinitely far away, i.e. parallel rays with a given intensity spectrum
interpreted as Watts / m^2.

By default the light source is "located" infinitely far away along the positive Z axis and
the light rays travel in the -Z direction.

# Fields

- `intensity::S`, the intensity of the light in Watts/m^2.
- `light_to_world::Transformation`, transformation to define the direction of the light
   source.
- `world_to_light::Transformation`, inverse transformation of the above.

"""
struct DistantLight{S<:Spectrum} <: LightSource
    intensity::S
    light_to_world::Transformation
    world_to_light::Transformation
end

"""
    function DistantLight(L::Spectrum, l2w::Transformation)

Create a new `DistantLight` given the intensity spectrum (in units of Watts / m^2) and
a `Transformation` used to rotate the direction of the light.

"""
function DistantLight(L::Spectrum, l2w::Transformation)
    DistantLight(L, l2w, inv(l2w))
end

"""
    sample_L(light::DistantLight, p::Point3)

Get an intensity sample from the given `DistantLight` and a ray from the light towards the
given point. The sample is always the same since the light source is a uniform parallel
light.

Returns:

- Intensity `Spectrum`
- A weight coefficient which is always 1.0
- A ray from the light source to the given point.

"""
function sample_L(light::DistantLight, p::Point3)
    return light.intensity, 1.0, VisibilityTester(p, light.light_to_world(NEG_Z_AXIS), 1e-9)
end

"""
    generate_ray(light::DistantLight, bounds::BoundingSphere, S::Sample)

Returns a ray from the light source at the given `BoundingSphere`, parametrized by the
given random sample.

"""
function generate_ray(light::DistantLight, bounds::BoundingSphere, S::Sample)
    ray = ray_parallel(bounds.radius, S.imgX, S.imgY)
    shift = translation(bounds.center)
    return shift(light.light_to_world(ray))
end

###### ########################
# TODO: Area light
struct AreaLight <: LightSource end

################################
# Disk Light (TODO)

"""
    DiskLight{S<:Spectrum} <: LightSource

A light source which is a uniformly illuminated disk far away. Expressed in terms of its
angular radius. By default it's in the positive Z direction, and can be rotated with
a transformation.

# Fields:

- `radius::Float64`, angular radius (in radians) of the source
- `intensity::S`,
- `light_to_world::Transformation`, transformation to adjust the direction of the source.
   Chosen to rotate the +Z axis into the desired direction of the light.
- `world_to_light::Transformation`, inverse transformation of the above

"""
struct DiskLight{S<:Spectrum} <: LightSource
    radius::Float64
    intensity::S
    light_to_world::Transformation
    world_to_light::Transformation
end

"""
    DiskLight(direction::Vector3, radius::Float64, intensity::Spectrum)

Create new `DiskLight` in the desired direction, angular radius and intensity.

"""
function DiskLight(direction::Vector3, radius::Float64, intensity::Spectrum)
    rot_z = rotate_z_to(direction)
    Ω = 2π * (1 - cos(radius))
    return DiskLight(radius, Ω.*intensity, rot_z, inv(rot_z))
end

@inline solid_angle(D::DiskLight) = 2π * (1 - cos(D.radius))

sample_L(light::DiskLight, p::Point3) = sample_L(light, p, rand(), rand())

function sample_L(light::DiskLight, p::Point3, u1::Real, u2::Real)
    direction = light.light_to_world(disk_light_sample(light.radius, u1, u2))
    return light.intensity, 1.0, VisibilityTester(p, direction, 2e-9)
end

"""
    disk_light_sample(radius::Real, u1::Real, u2::Real)

Create a vector which points from origin to a point on the surface of a disk with given
angular radius centered on the +Z axis. It is parametrized by two numbers, `u1` and `u2`,
which range from zero to one. Using uniform random numbers will result in the surface of
the disk to be uniformly sampled.

"""
function disk_light_sample(radius::Real, u1::Real, u2::Real)
    cos_th = cos(radius)
    z = cos_th + (1.0 - cos_th) * u1
    phi = 2π * u2
    r = sqrt(1.0 - z^2)
    return Vector3(r*cos(phi), r*sin(phi), z)
end

"""
    generate_ray(light::DiskLight, bounds::BoundingSphere, S::Sample)

Returns a ray from the light source at the given `BoundingSphere`, parametrized by the
given random sample.

"""
function generate_ray(light::DiskLight, bounds::BoundingSphere, S::Sample)
    # Pick a direction, i.e. a point on the (infinitely distant) disk
    # and make the transformation to rotate Z into that direction
    V = disk_light_sample(light.radius, S.lensU, S.lensV)
    rot = light.light_to_world * rotate_z_to(V)

    # Make a ray from infinitely far along the Z axis and rotate it into the direction
    # given by T. Also shift it relative to the origin of the bounding sphere.

    ray = ray_parallel(bounds.radius, S.imgX, S.imgY)
    T = translation(bounds.center) * rot
    return T(ray)
end


import Mustache
import Distributions

type Report
   num::Integer
   name::String
   times::Vector{Float64}
   analysis::SampleAnalysis
   outliers::Outliers
end

function render_iter(s::String, iter)
    [{s => x} for x in iter]
end

function report_to_dict(report::Report)
   kde_times = Distributions.kde(report.times)
   {"name" => report.name,
    "number" => report.num,
    "times" => render_iter("x", report.times),
    "kdetimes" => render_iter("x", kde_times.x),
    "kdepdf" => render_iter("x", kde_times.density),
    "mean_point" => report.analysis.mean.point,
    "mean_confidence_level" => report.analysis.mean.confidence_level,
    "mean_lbound" => report.analysis.mean.lbound,
    "mean_ubound" => report.analysis.mean.ubound,
    "outlier_effect" => report.analysis.outlier_variance.effect,
    "outlier_fraction" => report.analysis.outlier_variance.frac,
    "std_lbound" => report.analysis.std.lbound,
    "std_point" => report.analysis.std.point,
    "std_ubound" => report.analysis.std.ubound}
end

function read_css(template_dir::String)
    readall(abspath(joinpath(template_dir, "criterion.css")))
end

function report(rfile::String, reports::Vector{Report})
    report_dicts = {report_to_dict(r) for r in reports}
    tmpl = Mustache.template_from_file("../templates/report.tpl")
    template_dir = abspath("../templates/")
    fh = open(rfile,"w")
    Mustache.render(fh, tmpl, {"name" => "aaa",
                    "include" => [{"dir" => string("file://", template_dir),
				   "css" => read_css(template_dir)}],
		    "report" => report_dicts})
    close(fh)
end

function report(rfile::String, r::Report)
    report(rfile, [r])
end

function report(rs::Vector{Report})
    report("criterion.html", rs)
end

function report(r::Report)
    report("criterion.html", [r])
end

# JUnit Output
function junit(reports::Vector{Report})

    attr_esc = (esc) -> begin
        if     esc == '\'' "'"
        elseif esc == '"'  """
        elseif esc == '<'  "<"
        elseif esc == '>'  ">"
        elseif esc == '&'  "&"
        else string(esc)
        end
    end

    esc_html = (str) -> join([attr_esc(c) for c in str], "")

    report_out = String[]
    ntests = length(reports)
    header = string("<?xml versions\"1.0\" encoding=\"UTF-8\"?>\n",
                    "<testsuite name=\"Criterion benchmarks\" tests=\"$ntests\">\n")
    push!(report_out, header)
    for report in reports
       name = esc_html(report.name)
       time = report.analysis.mean.point
       push!(report_out, "<testcase name=\"$name\" time=\"$time\"/>\n")
   end
   push!(report_out, "</testsuite>\n")
   join(report_out, "")
end

function junit(report::Report)
    junit([report])
end