#### thermodynamics

export new_thermo_state, recover_thermo_state

"""
    new_thermo_state(atmos::AtmosModel, state::Vars, aux::Vars)

Create a new thermodynamic state, based on the `state`, and _not_
the `aux` state.

!!! note
    This method calls the iterative saturation adjustment
    procedure for EquilMoist models.
"""
function new_thermo_state end

# First dispatch on compressibility
new_thermo_state(atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state(compressibility_model(atmos), atmos, state, aux)

new_thermo_state(::Compressible, atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state(
        atmos,
        energy_model(atmos),
        moisture_model(atmos),
        state,
        aux,
    )
new_thermo_state(::Anelastic1D, atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state_anelastic(atmos, state, aux)

"""
    recover_thermo_state(atmos::AtmosModel, state::Vars, aux::Vars)

An atmospheric thermodynamic state.

!!! warn
    For now, we are directly calling new_thermo_state to avoid
    inconsistent aux states in kernels where the aux states are
    out of sync with the boundary state.

# TODO: Define/call `recover_thermo_state` when it's safely implemented
  (see https://github.com/CliMA/ClimateMachine.jl/issues/1648)
"""
function recover_thermo_state end

# First dispatch on compressibility
recover_thermo_state(atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state(compressibility_model(atmos), atmos, state, aux)

recover_thermo_state(
    ::Compressible,
    atmos::AtmosModel,
    state::Vars,
    aux::Vars,
) = new_thermo_state(
    atmos,
    energy_model(atmos),
    moisture_model(atmos),
    state,
    aux,
)

recover_thermo_state(::Anelastic1D, atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state_anelastic(atmos, state, aux)

function new_thermo_state(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::DryModel,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    return PhaseDry(param_set, e_int, state.ρ)
end

function new_thermo_state(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
)
    e_int = internal_energy(atmos, state, aux)
    param_set = parameter_set(atmos)
    return PhaseEquil(
        param_set,
        e_int,
        state.ρ,
        state.moisture.ρq_tot / state.ρ,
        moist.maxiter,
        moist.tolerance,
    )
end

function new_thermo_state(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::NonEquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    q = PhasePartition(
        state.moisture.ρq_tot / state.ρ,
        state.moisture.ρq_liq / state.ρ,
        state.moisture.ρq_ice / state.ρ,
    )

    return PhaseNonEquil{eltype(state), typeof(param_set)}(
        param_set,
        e_int,
        state.ρ,
        q,
    )
end

function new_thermo_state(
    atmos::AtmosModel,
    energy::θModel,
    moist::DryModel,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    θ_liq_ice = state.energy.ρθ_liq_ice / state.ρ
    return PhaseDry_ρθ(param_set, state.ρ, θ_liq_ice)
end

function new_thermo_state(
    atmos::AtmosModel,
    energy::θModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    θ_liq_ice = state.energy.ρθ_liq_ice / state.ρ
    return PhaseEquil_ρθq(
        param_set,
        state.ρ,
        θ_liq_ice,
        state.moisture.ρq_tot / state.ρ,
        moist.maxiter,
        moist.tolerance,
    )
end

function new_thermo_state(
    atmos::AtmosModel,
    energy::θModel,
    moist::NonEquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    θ_liq_ice = state.energy.ρθ_liq_ice / state.ρ
    q = PhasePartition(
        state.moisture.ρq_tot / state.ρ,
        state.moisture.ρq_liq / state.ρ,
        state.moisture.ρq_ice / state.ρ,
    )

    return PhaseNonEquil{eltype(state), typeof(param_set)}(
        param_set,
        state.ρ,
        θ_liq_ice,
        q,
    )
end


abstract type InverterComponent <: DynamicComponent end

"""
    mutable struct DynamicInverter{
        C <: Converter,
        O <: OuterControl,
        IC <: InnerControl,
        DC <: DCSource,
        P <: FrequencyEstimator,
        F <: Filter,
    } <: DynamicInjection
        name::String
        ω_ref::Float64
        converter::C
        outer_control::O
        inner_control::IC
        dc_source::DC
        freq_estimator::P
        filter::F
        base_power::Float64
        n_states::Int
        states::Vector{Symbol}
        ext::Dict{String, Any}
        internal::InfrastructureSystemsInternal
    end

A dynamic inverter is composed by 6 components, namely a Converter, an Outer Control, an Inner Control,
a DC Source, a Frequency Estimator and a Filter. It requires a Static Injection device that is attached to it.

# Arguments
- `name::String`: Name of inverter.
- `ω_ref::Float64`: Frequency reference set-point in pu.
- `converter <: Converter`: Converter model for the PWM transformation.
- `outer_control <: OuterControl`: Outer-control controller model.
- `inner_control <: InnerControl`: Inner-control controller model.
- `dc_source <: DCSource`: DC Source model.
- `freq_estimator <: FrequencyEstimator`: Frequency Estimator (typically a PLL) model.
- `filter <: Filter`: Filter model.
- `base_power::Float64`: Base power
- `n_states::Int`: Number of states (will depend on the components).
- `states::Vector{Symbol}`: Vector of states (will depend on the components).
- `ext::Dict{String, Any}`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct DynamicInverter{
    C <: Converter,
    O <: OuterControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
} <: DynamicInjection
    name::String
    ω_ref::Float64
    converter::C
    outer_control::O
    inner_control::IC
    dc_source::DC
    freq_estimator::P
    filter::F
    base_power::Float64
    n_states::Int
    states::Vector{Symbol}
    ext::Dict{String, Any}
    internal::InfrastructureSystemsInternal
end

function DynamicInverter(
    name::String,
    ω_ref::Float64,
    converter::C,
    outer_control::O,
    inner_control::IC,
    dc_source::DC,
    freq_estimator::P,
    filter::F,
    base_power::Float64 = 100.0,
    ext::Dict{String, Any} = Dict{String, Any}(),
) where {
    C <: Converter,
    O <: OuterControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
}
    n_states = _calc_n_states(
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
    )
    states = _calc_states(
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
    )

    return DynamicInverter{C, O, IC, DC, P, F}(
        name,
        ω_ref,
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
        base_power,
        n_states,
        states,
        ext,
        InfrastructureSystemsInternal(),
    )
end

function DynamicInverter(;
    name::String,
    ω_ref::Float64,
    converter::C,
    outer_control::O,
    inner_control::IC,
    dc_source::DC,
    freq_estimator::P,
    filter::F,
    base_power::Float64 = 100.0,
    n_states = _calc_n_states(
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
    ),
    states = _calc_states(
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
    ),
    ext::Dict{String, Any} = Dict{String, Any}(),
    internal = IS.InfrastructureSystemsInternal(),
) where {
    C <: Converter,
    O <: OuterControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
}
    return DynamicInverter(
        name,
        ω_ref,
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
        base_power,
        n_states,
        states,
        ext,
        internal,
    )
end

function DynamicInverter(
    name::String,
    ω_ref::Float64,
    converter::C,
    active_power_control::AC,
    reactive_power_control::RC,
    inner_control::IC,
    dc_source::DC,
    freq_estimator::P,
    ac_filter::F,
) where {
    C <: Converter,
    AC <: ActivePowerControl,
    RC <: ReactivePowerControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
}
    outer_control = OuterControl(active_power_control, reactive_power_control)
    return DynamicInverter(
        name,
        ω_ref,
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        ac_filter,
    )
end

get_name(device::DynamicInverter) = device.name
get_ω_ref(device::DynamicInverter) = device.ω_ref
get_ext(device::DynamicInverter) = device.ext
get_states(device::DynamicInverter) = device.states
get_n_states(device::DynamicInverter) = device.n_states
get_converter(device::DynamicInverter) = device.converter
get_outer_control(device::DynamicInverter) = device.outer_control
get_inner_control(device::DynamicInverter) = device.inner_control
get_dc_source(device::DynamicInverter) = device.dc_source
get_freq_estimator(device::DynamicInverter) = device.freq_estimator
get_filter(device::DynamicInverter) = device.filter
get_base_power(device::DynamicInverter) = device.base_power
get_internal(device::DynamicInverter) = device.internal
get_P_ref(value::DynamicInverter) = get_P_ref(get_active_power(get_outer_control(value)))
get_V_ref(value::DynamicInverter) = get_V_ref(get_reactive_power(get_outer_control(value)))

set_base_power!(value::DynamicInverter, val) = value.base_power = val

function _calc_n_states(
    converter,
    outer_control,
    inner_control,
    dc_source,
    freq_estimator,
    filter,
)
    return converter.n_states +
           outer_control.n_states +
           inner_control.n_states +
           dc_source.n_states +
           freq_estimator.n_states +
           filter.n_states
end

function _calc_states(
    converter,
    outer_control,
    inner_control,
    dc_source,
    freq_estimator,
    filter,
)
    return vcat(
        converter.states,
        outer_control.states,
        inner_control.states,
        dc_source.states,
        freq_estimator.states,
        filter.states,
    )
end

function get_REControlB_states(Q_Flag::Int)
    if Q_Flag == 0
        return [:Vt_filt, :I_icv]
    elseif Q_Flag == 1
        return [:Vt_filt, :ξ_icv]
    else
        error("Unsupported value of Q_Flag")
    end
end

function get_activeRETypeAB_states(Freq_Flag::Int)
    if Freq_Flag == 1
        return [:p_flt, :ξ_P, :p_ext, :p_ord], 4
    elseif Freq_Flag == 0
        return [:p_ord], 1
    else
        error("Unsupported value of Freq_Flag")
    end
end

function get_reactiveRETypeAB_states(Ref_Flag::Int, PF_Flag::Int, V_Flag::Int)
    if (Ref_Flag == 0) && ((PF_Flag == 1) && (V_Flag == 1))
        return [:pr_flt, :ξ_Q], 2
    elseif (Ref_Flag == 0) && ((PF_Flag == 1) && (V_Flag == 0))
        return [:pr_flt], 1
    elseif (Ref_Flag == 0) && ((PF_Flag == 0) && (V_Flag == 1))
        return [:q_flt, :ξq_oc, :q_LL, :ξ_Q], 4
    elseif (Ref_Flag == 0) && ((PF_Flag == 0) && (V_Flag == 0))
        return [:q_flt, :ξq_oc, :q_LL], 3
    elseif (Ref_Flag == 1) && ((PF_Flag == 1) && (V_Flag == 1))
        return [:pr_flt, :ξ_Q], 2
    elseif (Ref_Flag == 1) && ((PF_Flag == 1) && (V_Flag == 0))
        return [:pr_flt], 1
    elseif (Ref_Flag == 1) && ((PF_Flag == 0) && (V_Flag == 1))
        return [:V_cflt, :ξq_oc, :q_LL, :ξ_Q], 4
    elseif (Ref_Flag == 1) && ((PF_Flag == 0) && (V_Flag == 0))
        return [:V_cflt, :ξq_oc, :q_LL], 3
    else
        error("Unsupported value of Ref_Flag, PF_Flag or V_Flag")
    end
end


module InstrumentOperator

using OffsetArrays
using Polynomials
using NCDatasets
using Interpolations
using Parameters
using DocStringExtensions
using Distributions
using ImageFiltering   # for convolutions
using Unitful
using OrderedCollections
using UnitfulEquivalences
using YAML

@derived_dimension PerTime inv(Unitful.𝐓)

include("types.jl")                   # All types used in this module
include("ils_tables_io.jl")           # IO functions
include("prepare_ils.jl")             # function for ILS preparation
include("instrument_convolutions.jl") # conv using ImageFiltering, can be changed later
include("noise_functions.jl")
include("io_L1_files.jl")

export FixedKernelInstrument, VariableKernelInstrument, FTSInstrument, create_instrument_kernel, conv_spectra

end


module Signals

using ..Network
export take_due_signals, connect, Signal, add_acceptor, add_donor, can_add_acceptor, can_add_donor,
    name_back_spike, BackSpike, name_delta_v, DeltaV, name_dc_instruction, DCInstruction,
    name_new_dc, NewDC, ForwardSignal, BackwardSignal

abstract type ForwardSignal <: Signal end
abstract type BackwardSignal <: Signal end
function amplify end

# struct MarkramDelta{T <: AbstractFloat} <: ForwardSignal
#     delta::T
# end

# struct ExctDeltaCond{T <: AbstractFloat} <: ForwardSignal
#     delta_cond::T
# end

# struct InhbtDeltaCond{T <: AbstractFloat} <: ForwardSignal
#     delta_cond::T
# end

# struct DeltaCond{T <: AbstractFloat} <: ForwardSignal
#     delta_cond::T
# end

const name_delta_v = "DeltaV"
struct DeltaV{T <: AbstractFloat} <: ForwardSignal
    delta_v::T
end

amplify(s::DeltaV{T}, w::T) where {T <: AbstractFloat} = DeltaV(s.delta_v * w)

## state_dict is for recording states, not for network-serialization. temporarily commented.
#state_dict(s::DeltaV{T}) where {T <: AbstractFloat} = Dict(["delta_v" => s.delta_v])

const name_dc_instruction = "DCInstruction"
struct DCInstruction{T <: AbstractFloat} <: ForwardSignal
    previous::T
    new::T
end

amplify(s::DCInstruction{T}, w::T) where {T <: AbstractFloat} = DCInstruction(s.previous * w, s.new * w)

const name_new_dc = "NewDC"
struct NewDC{T <: AbstractFloat} <: ForwardSignal
    current::T
end

const name_back_spike = "BackSpike"
struct BackSpike <: BackwardSignal
end

## state_dict is for recording states, not for network-serialization. temporarily commented.
#state_dict(s::DCInstruction{T}) where {T <: AbstractFloat} = Dict(["precious" => s.previous, "new" => s.new])

function can_add_acceptor end # (agent, signal_name) -> bool
function can_add_donor end
function add_acceptor end # (agent, signal_name, acptr_address) -> ()
function add_donor end # (agent, signal_name, dnr_address) -> ()

function connect(network::Dict{String, Population{T}},
                 signal_name::String,
                 donor_address::Address,
                 acceptor_address::Address) where {T <: Unsigned}
    dnr = get_agent(network, donor_address)
    acptr = get_agent(network, acceptor_address)
    if can_add_acceptor(dnr, signal_name) && can_add_donor(acptr, signal_name)
        add_acceptor(dnr, signal_name, acceptor_address)
        add_donor(acptr, signal_name, donor_address)
    else
        error(
            "connect failed. Donor: $donor_address; acceptor: $acceptor_address; signal_name: $signal_name"
        )
    end
end

end # Module end


__precompile__()

module PolicyGradient

using mjWrap
using LearningStrategies
import LearningStrategies: setup!, update!
using Baseline
using Policy
using Tools
using ExpFunctions

using ReverseDiff

using IterativeSolvers

export PolicyGradModel, NPGStrategy, setup!, update!, PolicyRollout

struct PolicyGradModel
   theta::Vector     # policy parameters
   mjsys::mjWrap.mjSet
   policy::Policy.AbstractPolicy

   # scratch stuff
   samples::mjWrap.TrajSamples
   meansamples::mjWrap.TrajSamples

   trace::Dict{Symbol,Vector{Float64}}

   niter::Int

   function PolicyGradModel(theta,
                            mjsys, policy,
                            T, numT, meanT, niter)
      ns = mjsys.ns
      samples     = mjw.allocateTrajSamples(mjsys, T, numT, ns)
      meansamples = mjw.allocateTrajSamples(mjsys, T, meanT, ns)
      return new(theta, mjsys, policy,
                 samples, meansamples,
                 Dict(:stocR => Vector{Float64}(),
                      :meanR => Vector{Float64}(),
                      :evalscore => Vector{Float64}() ), niter)
                 #zeros(niter), zeros(niter), zeros(niter), niter)
   end
end

struct NPGStrategy{dtype<:AbstractFloat} <: LearningStrategy
   vpg::Vector{dtype}
   npg::Vector{dtype}
   grad::Vector{dtype}

   fullFIM::Bool
   gradll::Matrix{dtype}
   FIM::Matrix{dtype}

   adv::Vector{dtype}
   ret::Vector{dtype}

   baseline::Baseline.AbstractBaseline

   norm_step_size::Float64
   gamma::Float64
   gae::Float64

   cg_iter::Int
   cg_reg::Float64
   cg_tol::Float64

   obsmat::Matrix{dtype}
   ctrlmat::Matrix{dtype}

   function NPGStrategy{dtype}(model,
                               baseline,
                               fullFIM,
                               norm_step_size,
                               gamma, gae,
                               cg_iter, cg_reg, cg_tol) where dtype<:AbstractFloat
      #dtype = eltype(model.policy.theta) #Float64
      T, numT = model.samples.T, model.samples.numT

      nparam  = length(model.theta)
      vpg     = zeros(dtype, nparam)
      npg     = zeros(dtype, nparam)
      grad    = zeros(dtype, nparam)

      if fullFIM
         gradll  = zeros(dtype, nparam, T*numT)
         FIM     = zeros(dtype, nparam, nparam)
      else
         gradll  = zeros(dtype, 1, 1)
         FIM     = zeros(dtype, 1, 1)
      end

      adv     = zeros(dtype, T*numT)
      ret     = zeros(dtype, T*numT)

      tmpobs  = rand(dtype, model.mjsys.ns, T*numT)
      tmpctrl = rand(dtype, model.mjsys.nu, T*numT)

      new(vpg, npg, grad,
          fullFIM, gradll, FIM, adv, ret,
          baseline,
          norm_step_size,
          gamma, gae,
          cg_iter, cg_reg, cg_tol,
          tmpobs, tmpctrl)
   end
end
Base.show(io::IO, s::NPGStrategy) = print(io, "NPG: $(s.fullFIM?"FIM":"HVP"), δ=$(s.norm_step_size)\n")

function grad_builder(f::Function, input)
   ctp = ReverseDiff.GradientTape(f, input)
   if length(ctp.tape) <= 10000
      info("Compiling tape")
      ctp = ReverseDiff.compile(ctp)
   end
   return (out,in)->ReverseDiff.gradient!(out, ctp, in) # taped
end

# feature of learning strats; need to have iter and null to get iteration num
function update!(model::PolicyGradModel, s::NPGStrategy, iter, null)

   tic()

   samples  = model.samples
   mjsys    = model.mjsys
   pol      = model.policy

   b        = s.baseline

   nsamples = samples.T*samples.numT

   # Parallize here......................................................
   rsum = sum(samples.reward, 1) # rollouts done in different strat

   if eltype(pol.theta) == Float64
      ctrlmat = reshape(samples.ctrl, mjsys.nu, nsamples)
      obsmat  = reshape(samples.obs, mjsys.ns, nsamples)
   else
      copy!(s.ctrlmat, convert(Matrix{Float32}, reshape(samples.ctrl, mjsys.nu, nsamples)))
      copy!(s.obsmat, convert(Matrix{Float32}, reshape(samples.obs, mjsys.ns, nsamples)))
      ctrlmat = s.ctrlmat
      obsmat  = s.obsmat
   end

   println("returns")
   @time Tools.compute_returns!(s.ret, samples.reward, s.gamma)

   println("predict")
   @time Baseline.predict!(b, reshape(samples.obs, mjsys.ns, nsamples))

   println("advantages")
   @time Tools.computeGAEadvantages!(s.adv, b.bline,
                                     samples.reward,
                                     reshape(samples.obs, mjsys.ns, nsamples), s.gae, s.gamma)

   println("fit")
   @time Baseline.fit!(b, s.ret) # update fit

   if s.fullFIM
      DT = eltype(s.vpg)

      info("FULL FIM")
      println("gradll")
      @time Policy.gradLL!(s.gradll, pol, obsmat, ctrlmat)

      println("vpg")
      # vpg[:] = gradll*adv / nsamples
      @time BLAS.gemv!('N', 1/DT(nsamples), s.gradll, s.adv, DT(0.0), s.vpg)

      println("fim")
      #FIM[:] = gradll*gradll' / nsamples
      @time begin
         #const α = 0.9
         #FIM = copy(s.FIM)
         A_mul_Bt!(s.FIM, s.gradll, s.gradll)
         scale!(s.FIM, 1/DT(nsamples))
         #@. s.FIM = FIM + α*(s.FIM - FIM) # keep older FIM data around; works for higher α
      end

      println("cg")
      @time s.npg[:] = Tools.cpu_cg_solve(s.FIM, s.vpg, s.cg_iter, s.cg_reg, s.cg_tol)
      #@time IterativeSolvers.cg!(s.npg, s.FIM, s.vpg; tol = s.cg_tol, maxiter = s.cg_iter)
   else
      info("HVP algorithm")
      @time my∇LLxA = grad_builder( (theta, a)->Policy.loglikelihood(theta, pol, obsmat, ctrlmat)*a,
                                   (pol.theta, s.adv) )
      println("vpg")
      @time begin
         my∇LLxA((s.vpg, s.adv), (pol.theta, s.adv))
         scale!(s.vpg, 1/nsamples)
      end

      # TODO don't need in-place Loglikelihood for output; need inplace during comp
      minus = similar(s.adv')
      alpha = similar(s.adv')

      function hvp!(z, vec)
         const ϵ = sqrt(eps(eltype(vec))) * norm(vec)
         const eps_vec = ϵ * vec
         Policy.loglikelihood!(alpha, pol.theta+eps_vec, pol, obsmat, ctrlmat)
         Policy.loglikelihood!(minus, pol.theta-eps_vec, pol, obsmat, ctrlmat)
         alpha .-= minus
         scale!(alpha, 1.0/2ϵ)
         s.adv .= alpha'
         my∇LLxA((z, s.adv), (pol.theta, s.adv))
         #scale!(z, 1/nsamples)
      end
      println("cg")
      @time s.npg[:] = Tools.hvp_cg_solve(hvp!, s.vpg, s.cg_iter, s.cg_reg, s.cg_tol)
      scale!(s.npg, nsamples)
   end

   alpha = sqrt(s.norm_step_size / dot(s.npg, s.vpg))
   println("grad")
   s.grad .= alpha.*s.npg

   println("update")
   Policy.updatetheta(pol, pol.theta + s.grad) # update parameters
   model.theta[:] = pol.theta

   ########################################################## mean rollout
   println("no noise")
   meanreward = sum(model.meansamples.reward, 1)

   ############################################################### cleanup
   iterscore = mean(rsum)
   meanscore = mean(meanreward)
   if iter > 1
      @printf "\n  Stoc: %f, Mean: %f" iterscore meanscore
   else
      @printf "\n  Mean of rewards: %f" iterscore
   end
   @printf "\n         NPG Norm: %f\n\n" norm(s.npg) #normgrad[iter]
   @printf "  Reward of Mean Rollout: %f\n" meanscore
   itertime = toc()
   @printf "  Total time: %f\n" itertime
   push!(model.trace[:stocR], iterscore)
   push!(model.trace[:meanR], meanscore)
end

mutable struct PolicyRollout <: LearningStrategy
   f::FunctionSet
   policyaction::Function
   policymean::Function
   function PolicyRollout(f)
      new(f, ()->nothing, ()->nothing)
   end
end
Base.show(io::IO, s::PolicyRollout) = print(io, "π rollout\n")

function setup!(s::PolicyRollout, model::PolicyGradModel)
   s.policyaction = (c,o)->Policy.getaction!(c, model.policy, o)
   s.policymean   = (c,o)->Policy.getmean!(c, model.policy, o)
end

function stochasticrollout(mjsys::mjWrap.mjSet,
                           samples::mjWrap.TrajSamples,
                           ctrlfunc!::Function, f::FunctionSet, iter::Int)
   # stochastic rollout
   f.setmodel!(mjsys, iter) # change model for this iteration
   f.initstate!(mjsys, samples)
   f.setcontrol!(mjsys, samples.ctrl) # sets policy noise in ctrl
   @time mjWrap.rollout(mjsys, samples,
                        ctrlfunc!, f.observe!, f.reward)
end

function update!(model::PolicyGradModel, s::PolicyRollout, iter, null)
   println("policy rollout, stoc & mean")

   mjsys       = model.mjsys
   samples     = model.samples
   meansamples = model.meansamples

   tic()
   stochasticrollout(mjsys, samples, s.policyaction, s.f, iter)

   # policy's mean rollout
   s.f.initstate!(mjsys, meansamples)
   mjWrap.rollout(mjsys, meansamples,
                  s.policymean,
                  s.f.observe!, s.f.reward)
   toc()
end

end


module AlphaGo
using CuArrays
using Flux
using Flux: crossentropy, back!, mse, @treelike
using StatsBase: Weights
using Printf: @sprintf

struct IllegalMove <:Exception end
abstract type Position end
abstract type GameEnv end

struct PositionWithContext
  position::Position
  next_move
  result::Int
end

export MCTSPlayer, NeuralNet, pick_move, play_move!,
      initialize_game!, set_result!,
      should_resign, is_done, result, train!, selfplay,
      train, play, load_model, GoEnv, GomokuEnv, Position, all_legal_moves,
      ChessEnv, result

include("game/go/go.jl")
include("game/gomoku/gomoku.jl")
include("game/env.jl")
include("mcts.jl")
include("mcts_play.jl")
include("features.jl")
include("neural_net.jl")
include("selfplay.jl")
include("train.jl")
include("play.jl")

end #module


# mbe_demo_pitl.jl
#
# This script runs the control law on the target processor.

# Include our custom models.
include("models.jl")

# We work inside a module because it's more flexible than the main workspace.
module MBEDemoPITL

#############
# Resources #
#############

# Use/import the modules we'll need.
using Eratosthenes # The simulation engine (setup, simulate, mc)
using ..MBEModels  # The MBEModels module included above (and "above" this module)
import YAML        # For loading the YAML file directly
import HDF5        # For loading the log file
using Plots        # For plotting our results

# Create a directory for the log file.
if !isdir("out"); mkdir("out"); end

########################
# Software-in-the-Loop #
########################

# Specify the name of the YAML file that defines our scenario.
yaml_file = joinpath(@__DIR__, "underactuated.yaml")

# Load the YAML file and change the model name to use the PITL version.
# Also, change the log name.
yaml = YAML.load_file(yaml_file)
yaml["vehicles"][1]["computers"][1]["software"][1]["model"] = "ReducedEffortUnderactuatedControllerPITL"
yaml["sim"]["log_file"] = "out/underactuated-pitl.h5"

# The parameters are now stored under constants/parameters, so move them. Also,
# set up the UDP parameters here.
yaml["vehicles"][1]["computers"][1]["software"][1]["constants"] = Dict(
    "parameters"=>yaml["vehicles"][1]["computers"][1]["software"][1]["constants"],
    "conn"=>Dict(
        "target_ip"=>ip"192.168.1.195",
        "send_port"=>2000,
        "receive_port"=>2001
    )
);

# Set up the senario and run it.
scenario = setup(yaml)
scenario = simulate(scenario)

#########
# Plots #
#########

# Read some data from the logs.
t, q_BI, ω_BI_B = HDF5.h5open(scenario.sim.log_file, "r") do logs
    (dropdims(read(logs, "/smallsat/body/state/time"), dims=1),
     read(logs, "/smallsat/body/state/data/q_BI"),
     read(logs, "/smallsat/body/state/data/ω_BI_B"))
end

# Choose a friendly plotting package.
plotlyjs()

# Define each plot that we'll need.
display(plot(t, transpose(q_BI),
             label  = ["q1" "q2" "q3" "q4"],
             xlabel = "Time (s)",
             ylabel = "Attitude Quaternion"))
display(plot(t, 180/π * transpose(ω_BI_B),
             label  = ["ω1" "ω2" "ω3"],
             xlabel = "Time (s)",
             ylabel = "Rotation Rate (deg/s)"))

# Compare to the baseline.
(q_BI_baseline, ω_BI_B_baseline) = HDF5.h5open("out/underactuated.h5", "r") do logs
    (read(logs, "/smallsat/body/state/data/q_BI"),
     read(logs, "/smallsat/body/state/data/ω_BI_B"))
end
display(plot(t, transpose(q_BI) - transpose(q_BI_baseline),
    label  = ["q1" "q2" "q3" "q4"],
    xlabel = "Time (s)",
    ylabel = "Quaternion Differences",
    title  = "Differences in Attitude Quaternions"))
display(plot(t, 180/π * transpose((ω_BI_B) - transpose(ω_BI_B_baseline)),
    label  = ["ω1" "ω2" "ω3"],
    xlabel = "Time (s)",
    ylabel = "Rate (deg/s)",
    title  = "Differences in Rotation Rate"))

end # module


# This file is a part of SimilaritySearch.jl

"""
    execute_callback(opt::NeighborhoodSize, index)

SearchGraph's callback for adjusting neighborhood strategy
"""
function execute_callback(opt::NeighborhoodSize, index)
    N = index.neighborhood
    N.ksearch = ceil(Int, N.minsize + log(N.logbase, length(index)))
end

"""
    find_neighborhood(index::SearchGraph{T}, item, pools; hints=index.hints)

Searches for `item` neighborhood in the index, i.e., if `item` were in the index whose items should be its neighbors (intenal function).
`res` is always reused since `reduce` creates a new KnnResult from it (a copy if `reduce` in its simpler terms)
"""
function find_neighborhood(index::SearchGraph, item, pools::SearchGraphPools; hints=index.hints)
    n = length(index)
    if n > 0
        res = getknnresult(index.neighborhood.ksearch, pools)
        search(index.search_algo, index, item, res, hints, pools)
        reduce_neighborhood(index.neighborhood.reduce, index, item, res, pools)
    else
        Int32[]
    end
end

"""
    push_neighborhood!(index::SearchGraph, item, neighbors, callbacks; push_item=true)

Inserts the object `item` into the index, i.e., creates an edge from items listed in L and the
vertex created for ìtem` (internal function)
"""
function push_neighborhood!(index::SearchGraph, item, neighbors, callbacks; push_item=true)
    push_item && push!(index.db, item)
    push!(index.links, neighbors)
    push!(index.locks, Threads.SpinLock())
    n = length(index)
    n == 1 && return
    ## vstate = getvisitedvertices(index)
    @inbounds for id in neighbors
        push!(index.links[id], n)  # sat push?
    end

    callbacks !== nothing && execute_callbacks(callbacks, index)

    if index.verbose && length(index) % 100_000 == 0
        println(stderr, "added n=$(length(index)), neighborhood=$(length(neighbors)), $(string(index.search_algo)), $(Dates.now())")
    end
end

const GlobalSatKnnResult = [KnnResult(1)]
function __init__neighborhood()
    for _ in 2:Threads.nthreads()
        push!(GlobalSatKnnResult, KnnResult(1))
    end
end

"""
    SatNeighborhood()

New items are connected with a small set of items computed with a SAT like scheme (**cite**).
It starts with `k` near items that are reduced to a small neighborhood due to the SAT partitioning stage.
"""
struct SatNeighborhood <: NeighborhoodReduction end
Base.copy(::SatNeighborhood) = SatNeighborhood()

"""
    DistalSatNeighborhood()

New items are connected with a small set of items computed with a Distal SAT like scheme (**cite**).
It starts with `k` near items that are reduced to a small neighborhood due to the SAT partitioning stage but in reverse order of distance.
"""
struct DistalSatNeighborhood <: NeighborhoodReduction end
Base.copy(::DistalSatNeighborhood) = DistalSatNeighborhood()

"""
    struct IdentityNeighborhood

It does not modifies the given neighborhood
"""
struct IdentityNeighborhood <: NeighborhoodReduction end
Base.copy(::IdentityNeighborhood) = IdentityNeighborhood()

## functions

function sat_should_push(sat_neighborhood::T, index, item, id, dist, near::KnnResult) where T
    @inbounds obj = index[id]
    dist = dist < 0.0 ? evaluate(index.dist, item, obj) : dist
    push!(near, 0, dist)

    @inbounds for linkID in sat_neighborhood
        d = evaluate(index.dist, index[linkID], obj)
        push!(near, linkID, d)
    end

    argmin(near) == 0
end

function reduce_neighborhood(::IdentityNeighborhood, index::SearchGraph, item, res, pools::SearchGraphPools)
    copy(res.id)
end

"""
    reduce(sat::DistalSatNeighborhood, index::SearchGraph, item, res, pools)

Reduces `res` using the DistSAT strategy.
"""
@inline function reduce_neighborhood(::DistalSatNeighborhood, index::SearchGraph, item, res, pools::SearchGraphPools)
    N = Vector{Int32}(undef, 2)
    resize!(N, 0)
    @inbounds for i in length(res):-1:1  # DistSat => works a little better but also produces a bit larger neighborhoods
        id, dist = getpair(res, i)
        sat_should_push(N, index, item, id, dist, getsatknnresult(pools)) && push!(N, id)
    end

    N
end

@inline function reduce_neighborhood(::SatNeighborhood, index::SearchGraph, item, res, pools::SearchGraphPools)
    N = Vector{Int32}(undef, 2)
    resize!(N, 0)
    @inbounds for (id, dist) in res  # DistSat => works a little better but also produces a bit larger neighborhoods
        sat_should_push(N, index, item, id, dist, getsatknnresult(pools)) && push!(N, id)
    end

    N
end


#(When threading works, not yet):
#       To get the benefit of threading start julia with multiple threads:
#       JULIA_NUM_THREADS=x julia (where you choose x)

#A multiprocess write method for saving large image sequences in NRRD format
#This function is useful when retrieving a stack from img takes a lot of time due to lazy computation
function multiproc_write(fname::AbstractString, img::AbstractArray{T}, out_type=T) where {T}
    mmapa = prep_nrrd_write(fname, img, out_type)
    multiproc_write!(mmapa, img)
end

function prep_nrrd_write(fname::AbstractString, img::AbstractArray{T,N}, out_type=T) where {T,N}
    bname, ext = splitext(fname)
    if !isempty(ext) && !in(ext, [".nhdr", ".nrrd"])
        error("Must write in NRRD format (.nhdr or .nrrd extension)")
    end
    s = open(bname*".raw", "w+")
    mmapa = Mmap.mmap(s, Array{out_type, N}, size(img); shared=true)
    header = NRRD.headerinfo(out_type, axes(img))
    header["datafile"] = bname*".raw"
    open(fname, "w") do io
        write(io, magic(format"NRRD"))
        NRRD.write_header(io, "0004", header)
    end
    return mmapa
end

function multiproc_write!(mmapA::AbstractArray{T,3}, img::AbstractArray{T2,3}) where {T, T2}
    @assert size(mmapA) == size(img)
    mmapA[:,:,:] = img[:,:,:]
end

#Multithreading would be nice, but currently this doesn't seem to work well, perhaps because IO operations and external C library calls are restricted to one thread (see julia parallel docs as off march 2018)
function multiproc_write!(mmapA::AbstractArray{T,4}, img::AbstractArray{T2,4}) where {T, T2}
    @assert size(mmapA) == size(img)
    np = nworkers()
    if np == 0
        error("No workers available.  Add them with addprocs() (also make sure code and the input image are loaded on the worker processes")
    end
    print("Defining image and output array on other processes...\n")
    if np > 1
        @eval @everywhere img = $(img)
    end
    #nthreads = Base.Threads.nthreads()
    #print("Writing image with $(nthreads) threads\n")
    print("Writing image with $(np) procs\n")
    cur_i = 1
    while cur_i<=size(img,4)
        i_start = cur_i
        cur_i = min(size(img,4), i_start + np-1)
        idxs = [i_start:cur_i...]
        rrefs = []
        ws = workers()
#    Threads.@threads for i=1:size(img,4)
        for i=1:length(idxs)
            push!(rrefs, remotecall(getindex, ws[i], img, :,:,:,idxs[i]))
        end
        for i=1:length(idxs)
            mmapA[:,:,:,idxs[i]] = fetch(rrefs[i])
        end
        print("Wrote $cur_i stacks out of $(size(img,4)) so far.\n")
        cur_i+=1
    end
    print("Finished\n")
end


# literal zero for initialisation
zero(::Type{Quire8}) = Quire8(0x0000_0000)
zero(::Type{Quire16}) = Quire16(0x0000_0000_0000_0000_0000_0000_0000_0000)
zero(::Type{Quire32}) = Base.zext_int(Quire32,0x0)

# literal one
one(::Type{Quire8}) = Quire8(0x00001000)
one(::Type{Quire16}) = Quire16(0x01000000000000000000000000000000)
one(::Type{Quire32}) = fma(zero(Quire32),one(Posit32),one(Posit32))

# Fused multiply-add q+b*c
fma(q::Quire8,b::Posit8,c::Posit8) = ccall((:q8_fdp_add, SoftPositPath), Quire8, (Quire8,Posit8,Posit8),q,b,c)
fma(q::Quire16,b::Posit16,c::Posit16) = ccall((:q16_fdp_add, SoftPositPath), Quire16, (Quire16,Posit16,Posit16),q,b,c)
fma(q::Quire32,b::Posit32,c::Posit32) = ccall((:q32_fdp_add, SoftPositPath), Quire32, (Quire32,Posit32,Posit32),q,b,c)

# Fused multiply-sub q-b*c
fms(q::Quire8,a::Posit8,b::Posit8) = ccall((:q8_fdp_sub, SoftPositPath), Quire8, (Quire8,Posit8,Posit8),q,a,b)
fms(q::Quire16,a::Posit16,b::Posit16) = ccall((:q16_fdp_sub, SoftPositPath), Quire16, (Quire16,Posit16,Posit16),q,a,b)
fms(q::Quire32,a::Posit32,b::Posit32) = ccall((:q32_fdp_sub, SoftPositPath), Quire32, (Quire32,Posit32,Posit32),q,a,b)


using Test
using NQCDynamics
using Distributions: Normal

@testset "ElectronicDistribution" begin
    diabatic = NonadiabaticDistributions.Diabatic()

    @testset "SingleState" begin
        @test NonadiabaticDistributions.SingleState(1) === NonadiabaticDistributions.SingleState(1, diabatic)
    end

    @testset "ElectronicPopulation" begin
        dist = NonadiabaticDistributions.ElectronicPopulation([1, 0])
        dist2 = NonadiabaticDistributions.ElectronicPopulation([1, 0], diabatic)
        @test dist.statetype === dist2.statetype
        @test dist.populations == dist2.populations
    end
end

@testset "CombinedDistribution" begin
    electronic = NonadiabaticDistributions.SingleState(1)
    nuclear = DynamicalDistribution(Normal(), Normal(), (3, 2))
    combined = electronic * nuclear
    @test combined isa NonadiabaticDistributions.NonadiabaticDistribution
end


using JLD
using Plots
using JSON

# local
# @everywhere include("/home/genis/cluster_home/model_fitting_julia/gradient_utils.jl")
# @everywhere include("/home/genis/cluster_home/model_fitting_julia/functions_julia.jl")
# path="/home/genis/cluster_home/model_fitting_julia/"

#portatil
#
path="/home/emma/fit_behaviour_julia/"
# include(path*"/scripts/functions/functions_julia.jl")
# include(path*"/scripts/functions/utils.jl")
# include(path*"/scripts/functions/analysis.jl")
# include(path*"/scripts/functions/parameter_functions.jl")
# include(path*"/scripts/functions/simulations.jl")

SUBJECTS=[1,2,4,5]

#model_fit="urgency_x0_biasd_st3_dxdt42"
model_fit="StLinearDw_dx"

##### PARAM simulations ######3

tau=1.0
Tframe=0.5
Nframes=10
T=Tframe*Nframes

dt=tau/100.0
NT=T/dt
#############################

#plotly(ticks=:native)
#fig,axes=subplots(4,4,figsize=(20/2.54,20/2.54))

#colors_subject=["#cf4396","#59b94c","#a258c6","#a8b546","#5a6dc5",
#"#d7a044","#ab8ad0","#5a8437","#d43f4b","#60c395","#d47baf",
#"#358461","#c8683a","#4faad1","#8c7537","#b75d6c"]

for isubject in 1:length(SUBJECTS)

    subject=SUBJECTS[isubject]

    path_save=path*"figures/"
    filename_save=path_save*"check_PR_fit_PR_sim"*model_fit*"_subject"*string(subject)*".jld"
    data=JLD.load(filename_save)
    PR_sim=data["PR_sim"]
    PR_fit=data["PR_fit"]
    #axes[isubject].plot([0,1],[0,1],"k--")
    #axes[isubject].plot(PR_sim,PR_fit[1:length(PR_sim)],".",color=colors_subject[isubject])
    #axes[isubject].set_title(string(subject) )
    #if isubject==4
        #axes[isubject].set_xlabel("PR simulations")
    #end
    #if isubject==1 || isubject==4
        #axes[isubject].set_ylabel("PR Brunton")
    #end
end

#fig.tight_layout()
#fig.show()
#path_save=path*"figures/"
#fig.savefig(path_save*"PR_sim_PR_fit_model"*model_fit*".png")


"""
    AbstractEvolution

Abstract base type for strategies for how to evolve a population into a new population.
"""
abstract type AbstractEvolution end

"""
    evolve(e::AbstractEvolution, population)

Return a new population evolved by `e`.

New population may or may not contain same individuals as before.
"""
evolve(e::AbstractEvolution, population) = _evolve(e, population)

"""
    NoOpEvolution <: AbstractEvolution
    NoOpEvolution()

Does not evolve the given population.
"""
struct NoOpEvolution <: AbstractEvolution end
_evolve(::NoOpEvolution, pop) = pop

"""
    AfterEvolution <: AbstractEvolution
    AfterEvolution(evo, fun)

Return `fun(newpop)` where `newpop = evolve(e.evo, pop)` where `pop` is original population to evolve.
"""
struct AfterEvolution{F, E} <: AbstractEvolution
    evo::E
    fun::F
end

function _evolve(e::AfterEvolution, pop)
    newpop = evolve(e.evo, pop)
    return e.fun(newpop)
end

"""
    EliteSelection <: AbstractEvolution
    EliteSelection(nselect::Integer, evo=NoOpEvolution())

Selects the `nselect` highest fitness candidates to be passed on to `evo`.
"""
struct EliteSelection{N, E} <: AbstractEvolution
    nselect::N
    evo::E
end
EliteSelection(n::Integer) = EliteSelection(n, NoOpEvolution())
_evolve(e::EliteSelection, pop::AbstractArray{<:AbstractCandidate}) = evolve(e.evo, partialsort(pop, 1:e.nselect, by=fitness, rev=true))

"""
    SusSelection <: AbstractEvolution
    SusSelection(nselect, evo, rng=rng_default)

Selects candidates for further evolution using stochastic universal sampling.
"""
struct SusSelection{N, R, E} <: AbstractEvolution
    nselect::N
    evo::E
    rng::R
end
SusSelection(nselect, evo) = SusSelection(nselect, evo, rng_default)

function _evolve(e::SusSelection, pop::AbstractArray{<:AbstractCandidate})
    csfitness = cumsum(fitness.(pop))

    gridspace = csfitness[end] / e.nselect
    start = rand(e.rng)*gridspace

    selected = similar(pop, e.nselect)
    for i in eachindex(selected)
        candind = findfirst(x -> x >= start + (i-1)*gridspace, csfitness)
        selected[i] = pop[candind]
    end
    return evolve(e.evo, selected)
end

"""
    TournamentSelection <: AbstractEvolution
    TournamentSelection(nselect, k, p::Real, evo, rng=rng_default)

Selects candidates for further evolution using tournament selection.

Holds `nselect` tournaments with one winner each where each tournament has `k`
random candidates from the given population.

Winner of a tournament is selected as the candidate with highest fitness with a
probability `p`, second highest fitness with a probability `p(p-1)`, third highest
fitness with a probability of `p((p-1)^2)` and so on.
"""
struct TournamentSelection{N,P,E,R} <: AbstractEvolution
    nselect::N
    k::N
    p::P
    evo::E
    rng::R
    function TournamentSelection(nselect::N, k::N, p::P, evo::E, rng::R=rng_default) where {N, P, E, R}
        @assert 0 <= p <= 1 "0 <= p <= 1 not fulfilled! p = $p"
        return new{N, Vector{P}, E, R}(nselect, k, p .* ((1 - p) .^ collect(0:k-1)), evo, rng)
    end
end

function _evolve(e::TournamentSelection, pop::AbstractArray{<:AbstractCandidate})
    # Step 1: Create a random tournament order with as little repetition as possible so that we can select
    # e.nselect candidates out of e.nselect tournaments with e.k random candidates in each tournament
    n = e.nselect * e.k
    nrem = mod(n, length(pop))
    nrep = max(0, n ÷ length(pop))
    torder = mapfoldl(i -> shuffle(e.rng, pop), vcat, 1:nrep, init = shuffle(e.rng, pop)[1:nrem])

    selected = similar(pop, e.nselect)
    for (i, cands) in enumerate(Iterators.partition(torder, e.k))
        score = rand(e.rng, e.k) .* e.p
        selected[i] = partialsort(cands, argmax(score), by=fitness, rev=true)
    end
    return evolve(e.evo, selected)
end

"""
    CombinedEvolution <: AbstractEvolution
    CombinedEvolution(evos::AbstractArray)
    CombinedEvolution(evos...)

Combines the evolved populations from several `AbstractEvolution`s into one population.
"""
struct CombinedEvolution{E<:AbstractArray} <: AbstractEvolution
    evos::E
end
CombinedEvolution(evos...) = CombinedEvolution(collect(evos))
_evolve(e::CombinedEvolution, pop) = mapfoldl(evo -> evolve(evo, pop), vcat, e.evos)

"""
    EvolutionChain <: AbstractEvolution
    EvolutionChain(evos::AbstractArray)
    EvolutionChain(evos...)

Chains multiple `AbstractEvolution`s in a sequence so that output from the first is input to the next and so on.
"""
struct EvolutionChain{E<:AbstractArray} <: AbstractEvolution
    evos::E
end
EvolutionChain(evos...) = EvolutionChain(collect(evos))
_evolve(e::EvolutionChain, pop) = foldr(evolve, reverse(e.evos); init=pop)

"""
    PairCandidates <: AbstractEvolution
    PairCandidates(evo::AbstractEvolution)

Creates pairs of candidates in a population and calls `evolve(evo, pairs)` where `pairs` is the array of pairs.
"""
struct PairCandidates{E<:AbstractEvolution} <: AbstractEvolution
    evo::E
end
function _evolve(e::PairCandidates, pop)
    padpop = iseven(length(pop)) ? pop : vcat(pop, pop[1])
    pairs = collect(zip(padpop[1:2:end], padpop[2:2:end]))
    return foldl(evolve(e.evo, pairs), init=[]) do popout, (c1,c2)::Tuple
        length(popout) < length(pop) - 1 ? vcat(popout, c1, c2) : vcat(popout, c1)
    end
end

"""
    ShuffleCandidates
    ShuffleCandidates()
    ShuffleCandidates(rng)

Shuffles the population using `rng`.

Useful with [`PairCandidates`](@ref) in case the prior selection does not shuffle the population.
"""
struct ShuffleCandidates{R} <: AbstractEvolution
    rng::R
end
ShuffleCandidates() = ShuffleCandidates(rng_default)
_evolve(e::ShuffleCandidates, pop) = shuffle(e.rng, pop)

"""
    EvolveCandidates <: AbstractEvolution
    EvolveCandidates(fun)

Applies `fun(c)` for each candidate `c` in a given population.

Useful with [`evolvemodel`](@ref).
"""
struct EvolveCandidates{F} <: AbstractEvolution
    fun::F
end
_evolve(e::EvolveCandidates, pop) = map(e.fun, pop)


############################################################################################
# Types and structs - Resampling Methods

############################################################################################
"""
$(TYPEDEF)

Container that determines if jittering has to be applied

# Fields
$(TYPEDFIELDS)
"""
struct JitterTune{B<:UpdateBool}
    "Bolean if adaptive jittering is applied. If false, always max steps are jittered."
    adaption::B
    "Number of jitttering steps performed while jittering"
    Nsteps::Iterator
    "Jittering threshold for correlation of parameter."
    threshold::Float64
    "Minimum amount of jittering steps."
    min::Int64
    "Maximum amount of jittering steps."
    max::Int64
    function JitterTune(
        adaption::B,
        threshold::Float64,
        min::Integer=1,
        max::Integer=30
    ) where {B<:UpdateBool}
        ArgCheck.@argcheck 0 < min <= max
        ArgCheck.@argcheck 0.0 < threshold <= 1.0 "threshold needs to be positive"
        return new{B}(adaption, Iterator(0), threshold, min, max)
    end
end

############################################################################################
function update!(tune::JitterTune)
    init!(tune.Nsteps, 0)
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Check if we can stop jittering step based on correlation ρ.

# Examples
```julia
```

"""
function jitter!(tune::JitterTune, adaption::UpdateTrue, ρ::T) where {T<:AbstractFloat}
    ## Add step to Nsteps
    update!(tune.Nsteps)
    ## Check if jittering has to be continued
    #!NOTE: Need to start with >= tune.max and < tune.min conditions, and then with ρ
    if tune.Nsteps.current >= tune.max
        return false
    elseif tune.Nsteps.current < tune.min
        return true
    #!NOTE: if not above or below boundary, check if correlation threshold is fullfiled
    # If correlation smaller than threshold, stop jittering
    elseif ρ < tune.threshold
        return false
    else
        return true
    end
end

function jitter!(tune::JitterTune, adaption::UpdateFalse, ρ::T) where {T<:AbstractFloat}
    ## Add step to Nsteps
    update!(tune.Nsteps)
    ## Continue jittering until max steps are reached
    if tune.Nsteps.current >= tune.max
        return false
    else
        return true
    end
end

function jitter!(tune::JitterTune, ρ::T) where {T<:AbstractFloat}
    return jitter!(tune, tune.adaption, ρ)
end

############################################################################################
#export
export JitterTune, jitter!, update!


export @alchemy

macro alchemy(block)
    mol = repr(block)
    mol = replace(mol, ";"=>"\n")
    mol = strip(filter(c->!occursin(c,"{}():"), mol))
    mol = String(mol)
    mol = Molecules.parse_string(mol)
    quote
        plot($mol)
    end
end

macro alchemy(mol::Symbol)
    return :(Alchemy.plot($mol)) |> esc
end

macro alchemy(fpath::String)
    mol = Molecules.parse_file(fpath)
    return :(Alchemy.plot($mol)) |> esc
end

function plot(atoms::Vector{T}) where T <: Molecules.Atom
    S = Scene()
    plot!(S, atoms)
    return S
end

function plot!(S, atoms::Vector{T}) where T <: Molecules.Atom
    for i = eachindex(atoms)
        center = Point3f0(atoms[i].xyz)
        symbol = Molecules.symbol(atoms[i])
        r = get_atom_radius(symbol)
        c = get_atom_color(symbol)

        sphere = Sphere(center, r)

        mesh!(S, sphere, show_axis=false, color=c)
    end

    # Add chemical bonds for close atoms
    for i in eachindex(atoms)
        A1 = atoms[i]
        r1 = A1.xyz
        for j in (i+1):length(atoms)
            A2 = atoms[j]
            r2 = A2.xyz
            d = √sum((r1 .- r2).^2)
            if d < 1.55
                xyz = [Point3f0(A1.xyz)]
                uvw = [Point3f0(A2.xyz - A1.xyz)]
                arrows!(S, xyz, uvw, linewidth=0.1, arrowsize=0.0, linecolor=stegeman)
            end
        end
    end

    return S
end


############################################################################################
"""
$(TYPEDEF)
Super type for various particle resampling techniques.
"""
abstract type ResamplingMethod end

############################################################################################
"""
$(TYPEDEF)

Stores information about resampling steps.

# Fields
$(TYPEDFIELDS)
"""
struct ResampleTune{R<:ResamplingMethod}
    "Method for resampling."
    method::R
    "Stores if last iteration was resampled."
    update::Updater
    function ResampleTune(method::R, update::Updater) where {R<:ResamplingMethod}
        return new{R}(method, update)
    end
end
function update!(tune::ResampleTune)
    update!(tune.update)
    return nothing
end
function init!(tune::ResampleTune, bool::Bool)
    init!(tune.update, bool)
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Resample particles, dispatched on ResamplingMethod subtypes.

# Examples
```julia
```

"""
function resample! end

function resample!(
    _rng::Random.AbstractRNG,
    type::S,
    container::AbstractMatrix{<:Integer},
    iter::Integer,
    weights::Vector{<:Real},
    n::Integer=length(weights),
) where {S<:ResamplingMethod}
    return resample!(_rng, type, view(container, :, iter), iter, weights, n)
end

############################################################################################
"""
$(SIGNATURES)
More stable, faster version of rand(Categorical) if weights sum up to 1.

# Examples
```julia
```

"""
function randcat(_rng::Random.AbstractRNG, p::AbstractVector{<:Real})
    T = eltype(p)
    r = rand(_rng, T)
    cp = p[1]
    s = 1
    n = length(p)
    while cp <= r && s < n
        @inbounds cp += p[s += 1]
    end
    return s
end

############################################################################################
# Export
export
    ResamplingMethod,
    ResampleTune,
    resample!,
    randcat


using Pkg
Pkg.activate(".")
using ConceptnetNumberbatch
using Serialization

local cptnet

# Load a serialized version of ConceptNetEnglish
fid = open("./_conceptnet_/numberbatch-en-17.06.txt.bin")
cptnet = deserialize(fid)
close(fid)

## Get an embedding matrix
phrase = "this is a phrase that containz some iwords"
ConceptnetNumberbatch.word_embeddings(cptnet, phrase, keep_size=false, search_mismatches=false)
@time embs=ConceptnetNumberbatch.word_embeddings(cptnet, phrase, keep_size=false, search_mismatches=false)
@time embs=ConceptnetNumberbatch.word_embeddings(cptnet, phrase, keep_size=false, search_mismatches=true)
println("Loaded $(size(embs, 2)) embedding vectors (out of $(length(split(phrase)))), $(size(embs,1)) elements each.")


"""
Randomly selects the genetic operator from the vector
according to its weight and applies it.
"""
struct FixedGeneticOperatorsMixture <: GeneticOperatorsMixture
    operators::Vector{GeneticOperator} # available operations
    weights::Weights{Float64, Float64, Vector{Float64}}  # fixed weights

    function FixedGeneticOperatorsMixture(
        operators::AbstractVector{<:GeneticOperator},
        rates::AbstractVector{Float64} = fill(1.0/length(operators), length(operators)) # defaults to uniform distribution of rates
    )
        length(operators) == length(rates) ||
            throw(DimensionMismatch("Number of mutators does not match the number of their rates"))
        new(GeneticOperator[op for op in operators], weights(rates))
    end
end

"""
Default implementation of `apply!()` for operators mixture.
"""
function apply!(opmix::GeneticOperatorsMixture,
                v::AbstractVector{<:Real}, target_index::Int)
    op, tag = next(opmix)
    apply!(op, v, target_index)
end

"""
    next(fa::FixedGeneticOperatorsMixture)

Gets the random genetic operator from the mixture.

The probability to select an operator is proportional to its weight.

Returns a tuple of the genetic operator and its index in the mix.
"""
function next(opmix::FixedGeneticOperatorsMixture)
    i = sample(1:length(opmix.operators), opmix.weights)
    return opmix.operators[i], i
end

"""
Frequency-adapting genetic operators mixture.
"""
struct FAGeneticOperatorsMixture <: GeneticOperatorsMixture
    operators::Vector{GeneticOperator} # available operators
    fa::FrequencyAdapter               # adapter of operator frequencies

    FAGeneticOperatorsMixture(
        operators::AbstractVector{<:GeneticOperator};
        c = 1E-2, eta = 1E-2, pmin = 0.01, pmax = 1.0
    ) = new(GeneticOperator[op for op in operators],
            FrequencyAdapter(length(operators);
                             c = c, eta = eta, pmin = pmin, pmax = pmax))
end

frequencies(opmix::FAGeneticOperatorsMixture) = frequencies(opmix.fa)

function next(opmix::FAGeneticOperatorsMixture)
    i = next(opmix.fa)
    return opmix.operators[i], i
end

function adjust!(opmix::FAGeneticOperatorsMixture, op_index::Int, candi_index::Int,
                 new_fitness::F, old_fitness::F, is_improved::Bool) where F
    # KLUDGE we don't know the fitness scheme, but if there is no improvement in fitness,
    # new_fitness==old_fitness, so it's ok to take the abs()
    update!(opmix.fa, op_index, is_improved ? abs(new_fitness-old_fitness) : 0.0)

    # also adjust the actual operator
    adjust!(opmix.operators[op_index], 0, candi_index, new_fitness, old_fitness, is_improved)
end

# FIXME use logging
function trace_state(io::IO, fa::FAGeneticOperatorsMixture, mode::Symbol)
    if mode == :verbose
        println(io, "FrequencyAdapting mixture rates:")
        for i in eachindex(fa.operators)
            @printf(io, " %s=%.2f", typeof(fa.operators[i]), fa.fa.p[i])
        end
    end
    println(io)
end


using Autosampler, ImagineFormat, CSV, DataFrames
using Test, Random

if !isfile("recording.ai")
    include("create_ai.jl")
end

@testset "Autosampler.jl" begin
    rows = ["chemical1_10uM,RA1", "chemical1_01uM,RA2"]
    rowid = [split(row, ',')[1] for row in rows]
    # The sequence of random numbers might change across Julia versions,
    # so determine what the sequence is
    Random.seed!(1234321)
    p = vcat(randperm(2), randperm(2), randperm(2), randperm(2))
    for stimfile in ("stimuli_tabs.csv", "stimuli_comma_quotes.csv")
        Random.seed!(1234321)
        mktemp() do pathcsv, iocsv
            # Generate sequence file
            stim_randomize(iocsv, stimfile, 4)
            str = readlines(pathcsv)
            @test str == rows[p]
            # Test insertion into .imagine file
            pathi, io = mktemp()
            close(io)
            update_imagine("recording.imagine", pathcsv; updatedfile=pathi, um_per_pixel=0.33)
            header = ImagineFormat.parse_header(pathi)
            @test split(header["stimulus sequence"], '\$') == rowid[p]
            @test header["stimulus scan hi"] == 1000:1000:8000
            @test length(header["stimulus frame hi"]) == 8
            rm(pathi)
        end
    end
end


abstract type AbstractParam <: AbstractDelayed end

"""
$(TYPEDSIGNATURES)

A node of the space used to assign parameter labels. Only Param objects are extracted into
the vals dictionary for output. Therefore, every parameter of the user's function must have
a corresponding Param object in the space.
"""
struct Param <: AbstractParam
    label::Symbol
    obj::AbstractDelayed
end


## Effect differences ##

const EFFECT_FUNCS = Dict{Symbol,Function}()

"Convert effect formula to a state difference (additions, deletions, etc.)"
function effect_diff(domain::Domain, state::State, effect::Term)
    effect_fn! = get(EFFECT_FUNCS, effect.name, add_effect!)
    return effect_fn!(GenericDiff(), domain, state, effect)
end

function add_effect!(diff::GenericDiff,
                     domain::Domain, state::State, effect::Term)
    if effect isa Compound # Evaluated all nested functions
        args = Term[a isa Compound ? Const(evaluate(domain, state, a)) : a
                    for a in effect.args]
        effect = Compound(effect.name, args)
    end
    push!(diff.add, effect)
    return diff
end

function del_effect!(diff::GenericDiff,
                     domain::Domain, state::State, effect::Term)
    effect = effect.args[1]
    if effect isa Compound # Evaluated all nested functions
        args = Term[a isa Compound ? Const(evaluate(domain, state, a)) : a
                    for a in effect.args]
        effect = Compound(effect.name, args)
    end
    push!(diff.del, effect)
    return diff
end
EFFECT_FUNCS[:not] = del_effect!

function and_effect!(diff::GenericDiff,
                     domain::Domain, state::State, effect::Term)
    for eff in get_args(effect)
        combine!(diff, effect_diff(domain, state, eff))
    end
    return diff
end
EFFECT_FUNCS[:and] = and_effect!

function when_effect!(diff::GenericDiff,
                      domain::Domain, state::State, effect::Term)
    cond, eff = effect.args[1], effect.args[2]
    if !satisfy(domain, state, cond) return diff end
    return combine!(diff, effect_diff(domain, state, eff))
end
EFFECT_FUNCS[:when] = when_effect!

function forall_effect!(diff::GenericDiff,
                        domain::Domain, state::State, effect::Term)
    cond, eff = effect.args[1], effect.args[2]
    for s in satisfiers(domain, state, cond)
        combine!(diff, effect_diff(domain, state, substitute(eff, s)))
    end
    return diff
end
EFFECT_FUNCS[:forall] = forall_effect!

function assign_effect!(diff::GenericDiff,
                        domain::Domain, state::State, effect::Term)
    term, val = effect.args[1], effect.args[2]
    diff.ops[term] = val
    return diff
end
EFFECT_FUNCS[:assign] = assign_effect!

function modify_effect!(diff::GenericDiff,
                        domain::Domain, state::State, effect::Term)
    term, val = effect.args[1], effect.args[2]
    op = GLOBAL_MODIFIERS[effect.name]
    diff.ops[term] = Compound(op, Term[term, val])
    return diff
end
for name in keys(GLOBAL_MODIFIERS)
    EFFECT_FUNCS[name] = modify_effect!
end

## Precondition differences ##

const PRECOND_FUNCS = Dict{Symbol,Function}()

"Convert precondition formula to a state difference."
function precond_diff(domain::Domain, state::State, precond::Term)
    precond_fn! = get(PRECOND_FUNCS, precond.name, add_precond!)
    return precond_fn!(GenericDiff(), domain, state, precond)
end

function add_precond!(diff::GenericDiff,
                      domain::Domain, state::State, precond::Term)
    push!(diff.add, precond)
    return diff
end

function del_precond!(diff::GenericDiff,
                      domain::Domain, state::State, precond::Term)
    push!(diff.del, precond.args[1])
    return diff
end
PRECOND_FUNCS[:not] = del_precond!

function and_precond!(diff::GenericDiff,
                      domain::Domain, state::State, precond::Term)
    for pre in get_args(precond)
        combine!(diff, precond_diff(domain, state, pre))
    end
    return diff
end
PRECOND_FUNCS[:and] = and_precond!

function forall_precond!(diff::GenericDiff,
                         domain::Domain, state::State, precond::Term)
    cond, pre = precond.args[1], precond.args[2]
    for s in satisfiers(domain, state, cond)
        combine!(diff, precond_diff(domain, state, substitute(pre, s)))
    end
    return diff
end
PRECOND_FUNCS[:forall] = forall_precond!


function communities_crime(repls=1, nfolds=6)
	communities_data = @load_communities_crime
	communities_data[1].Column8 = convert(Array{Float64},
										communities_data[1].Column8) .<= 0.18
	#0.18 is the mean value for racepctblack

	communities_crime_task = Task(communities_data[1], communities_data[2],
								grp=:Column8,
								debiasmeasures=[fpr, fnr])
	FairnessProblem(
		communities_crime_task,
		measures=[fnr, fpr],
		repls=repls, nfolds=nfolds, name="communities_crime")
end

function student(repls=1, nfolds=6)
   student_task = Task((@load_student_performance)...,
					   grp=:sex, debiasmeasures=[fnr])
   FairnessProblem(
	   student_task,
	   measures=[fnr, fpr],
	   repls=repls, nfolds=nfolds, name="StudentPerformance")
end

function hmda(repls=1, nfolds=6)
	hmda_df = CSV.read(joinpath(@__DIR__, "..", "data", "HMDA_2018_California.csv"))
	filter!(:derived_sex => x -> x != "Sex Not Available", hmda_df)
	categorical!(hmda_df, [:conforming_loan_limit, :derived_loan_product_type ,
			:derived_dwelling_category, :derived_ethnicity, :derived_race,
			:derived_sex, :action_taken, :purchaser_type, :loan_type,
			:applicant_age, ])
	select!(hmda_df, Not(:county_code))
	rename!(hmda_df, "open-end_line_of_credit" => "open_end_line_of_credit")

	hmda_task = Task(deletecols(hmda_df, :action_taken), hmda_df.action_taken,
				grp=:derived_sex, debiasmeasures=[false_omission_rate])

	FairnessProblem(hmda_task,
		measures=[false_omission_rate, fpr],
		repls=repls, nfolds=nfolds, name="HMDA_Mortgage")
end

function zafar(repls=1, nfolds=6)
   zafar_task = Task((genZafarData())...,
					   grp=:z, debiasmeasures=[])
   FairnessProblem(
	   zafar_task,
	   measures=[],
	   repls=repls, nfolds=nfolds, name="ZafarData")
end

function zafar2(repls=1, nfolds=6)
   zafar2_task = Task((genZafarData2())...,
					   grp=:z, debiasmeasures=[])
   FairnessProblem(
	   zafar2_task,
	   measures=[],
	   repls=repls, nfolds=nfolds, name="StudentPerformance")
end

function subgroup(repls=1, nfolds=6)
   subgroup_task = Task((genSubgroupData(setting="B01"))...,
					   grp=:z, debiasmeasures=[])
   FairnessProblem(
	   subgroup_task,
	   measures=[],
	   repls=repls, nfolds=nfolds, name="SubGroupData")
end

function biased(repls=1, nfolds=6)
   biased_task = Task((genBiasedSampleData())...,
					   grp=:z, debiasmeasures=[])
   biased = FairnessProblem(
	   biased_task,
	   measures=[],
	   repls=repls, nfolds=nfolds, name="BiasedSampleData")
end

function adult(repls=1, nfolds=6)
   adult_task = Task((@load_adult)..., grp=:sex, debiasmeasures=[ppr]) # could keep :sex also as protected attribute
   FairnessProblem(
	   adult_task,
	   refGrp="Male", measures=[ppr, fnr, fpr],
	   repls=repls, nfolds=nfolds, name="Adult")
end

function german(repls=1, nfolds=6)
   german_task = Task((@load_german)..., grp=:gender_status, debiasmeasures=[false_omission_rate])
   FairnessProblem(
	   german_task,
	   refGrp="male_single", measures=[foar, fnr, fpr],
	   repls=repls, nfolds=nfolds, name="German")
end

function portuguese(repls=1, nfolds=6)
   portuguese_task = Task((@load_bank_marketing)...,
					   grp=:marital, debiasmeasures=[false_omission_rate, ppr])
   FairnessProblem(
	   portuguese_task,
	   measures=[false_omission_rate, ppr, fpr],
	   repls=repls, nfolds=nfolds, name="Portuguese")
end

function framingham(repls=1, nfolds=6)
   framingham_df = CSV.read(joinpath(@__DIR__, "..", "data", "framingham.csv"))
   categorical!(framingham_df, [:male, :education, :currentSmoker, :BPMeds,
			   :prevalentStroke, :prevalentHyp, :diabetes, :TenYearCHD])

   framingham_task = Task(deletecols(framingham_df, :TenYearCHD),
			   framingham_df.TenYearCHD, grp=:male,
			   debiasmeasures=[false_omission_rate])

   FairnessProblem(framingham_task,
				   measures=[false_omission_rate],
				   repls=repls, nfolds=nfolds, name="Framingham")
end

function loan(repls=1, nfolds=6)
   loan_df = CSV.read(joinpath(@__DIR__, "..", "data", "loan_default.csv"))
   categorical!(loan_df, [:SEX, :EDUCATION, :MARRIAGE,
		   Symbol("default payment next month")])
   loan_task = Task(deletecols(loan_df, Symbol("default payment next month")),
			   loan_df[!, Symbol("default payment next month")],
			   grp=:SEX, debiasmeasures=[false_omission_rate])
   FairnessProblem(loan_task, measures=[false_omission_rate, fpr],
					   repls=repls, nfolds=nfolds, name="LoanDefault")
end

function medical(repls=1, nfolds=6)
   medical_df = CSV.read(joinpath(@__DIR__, "..", "data", "meps.csv"))

   categorical!(medical_df, ["REGION","SEX","MARRY","FTSTU","ACTDTY","HONRDC",
			   "RTHLTH","MNHLTH","HIBPDX","CHDDX","ANGIDX","MIDX",
			   "OHRTDX","STRKDX","EMPHDX","CHBRON","CHOLDX","CANCERDX","DIABDX",
			   "JTPAIN","ARTHDX","ARTHTYPE","ASTHDX","ADHDADDX","PREGNT","WLKLIM",
			   "ACTLIM","SOCLIM","COGLIM","DFHEAR42","DFSEE42", "ADSMOK42",
			   "PHQ242","EMPST","POVCAT","INSCOV", "UTILIZATION"])

   medical_task = Task(deletecols(medical_df, Symbol("UTILIZATION")),
			   medical_df[!, Symbol("UTILIZATION")],
			   grp=:RACE, debiasmeasures=[false_omission_rate])
   FairnessProblem(medical_task, measures=[false_omission_rate, fnr],
					   repls=repls, nfolds=nfolds, name="MedicalExpenditure")
end

fairness_problems() = [compas(), adult(), german(), portuguese(),
					   student(), framingham(), loan(), medical()]


mutable struct FileCachedFactor{T<:Unsigned, S<:AbstractFactor{T}} <: AbstractFactor{T}
    cachepath::String
    basefactor::S
    lockobj::Threads.TatasLock

    FileCachedFactor{T, S}(cachepath, basefactor, lockobj) where {S<:AbstractFactor{T}} where {T<:Unsigned} = (res = new(cachepath, basefactor, lockobj); finalizer(res, x -> rm(x.cachepath; force = true)); res)
end

Base.length(var::FileCachedFactor) = length(var.basefactor)

getname(var::FileCachedFactor) = getname(var.basefactor)

getlevels(var::FileCachedFactor) = getlevels(var.basefactor)

function FileCachedFactor(factor::S) where {S<:AbstractFactor{T}} where {T<:Unsigned}
    FileCachedFactor{T, S}("", factor, Threads.TatasLock())
end

function filecache(basefactor::S) where {S<:AbstractFactor{T}} where {T<:Unsigned}
    FileCachedFactor{T, S}("", basefactor, Threads.TatasLock())
end

function slice(factor::FileCachedFactor{T, S}, fromobs::Integer, toobs::Integer, slicelength::Integer) where {T<:Unsigned} where {S<:Union{FileCachedFactor, CachedFactor, Factor, FileFactor}}
    basefactor = factor.basefactor
    slicelength = verifyslicelength(fromobs, toobs, slicelength)
    slice(basefactor, fromobs, toobs, slicelength)
end

function slice(factor::FileCachedFactor{T, S}, fromobs::Integer, toobs::Integer, slicelength::Integer) where {S<:AbstractFactor{T}} where {T<:Unsigned}
    basefactor = factor.basefactor
    slicelength = verifyslicelength(fromobs, toobs, slicelength)
    lockobj = factor.lockobj
    lock(lockobj)
    try
        if !isfile(factor.cachepath)
            v = convert(Vector{T}, basefactor)
            tmpjuml = joinpath(tempdir(), JUMLDIR)
            mkpath(tmpjuml)
            outdatpath = joinpath(tmpjuml, "$(randstring(10)).dat")
            open(outdatpath, "w") do f
                write(f, v)
            end
            factor.cachepath = outdatpath
        end
    finally
        unlock(lockobj)
    end
    sizeT = sizeof(T)
    iostream = open(factor.cachepath, "r")
    seek(iostream, sizeT * (fromobs - 1))
    buffer = Vector{T}(undef, slicelength)
    map(Seq(Vector{T}, (buffer, iostream, fromobs, toobs, slicelength), nextdatachunk), SubArray{T,1,Array{T,1},Tuple{UnitRange{Int64}},true}) do slice
        view(slice, 1:length(slice))
    end
end

function Base.map(factor::FileCachedFactor{T}, dataframe::AbstractDataFrame; permute::Bool = false) where {T<:Unsigned}
    FileCachedFactor(map(factor.basefactor, dataframe; permute = permute))
end

function isordinal(factor::FileCachedFactor{T}) where {T<:Unsigned}
    isordinal(factor.basefactor)
end


# # Driving in a Stadium

#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/stadium.ipynb)

# This example demonstrates a 2D driving simulation where cars drive around a three-lane stadium.
# The entities are defined by the types:
#
# - `S` - `VehicleState`, containing the vehicle position (both globally and relative to the lane) and speed
# - `D` - `VehicleDef`, containing length, width, and class
# - `I` - `Symbol`, a unique label for each vehicle
#
# The environment is represented by a `Roadway` object which
# allows to define roads consisting of multiple lanes based on the RNDF format.

# We load relevant modules and generate a 3-lane stadium roadway:

using AutomotiveDrivingModels
using AutoViz
AutoViz.colortheme["background"] = colorant"white"; # hide
using Distributions

roadway = gen_stadium_roadway(3)
snapshot = render([roadway])
#md write("stadium.svg", snapshot) # hide
#md # ![three lane stadium](stadium.svg)

# As a next step, let's populate a scene with vehicles

w = DEFAULT_LANE_WIDTH
scene = Scene([
    Entity(VehicleState(VecSE2(10.0,  -w, 0.0), roadway, 29.0), VehicleDef(), :alice),
    Entity(VehicleState(VecSE2(40.0, 0.0, 0.0), roadway, 22.0), VehicleDef(), :bob),
    Entity(VehicleState(VecSE2(30.0, -2w, 0.0), roadway, 27.0), VehicleDef(), :charlie),
])
car_colors = get_pastel_car_colors(scene)
renderables = [
    roadway,
    (FancyCar(car=veh, color=car_colors[veh.id]) for veh in scene)...
]
snapshot = render(renderables)
#md write("stadium_with_cars.svg", snapshot) # hide
#md # ![stadium with cars](stadium_with_cars.svg)

# We can assign driver models to each agent and simulate the scenario.

timestep = 0.1
nticks = 300

models = Dict{Symbol, DriverModel}(
    :alice => LatLonSeparableDriver( # produces LatLonAccels
        ProportionalLaneTracker(), # lateral model
        IntelligentDriverModel(), # longitudinal model
    ),
    :bob => Tim2DDriver(
         mlane = MOBIL(),
    ),
    :charlie => StaticDriver{AccelTurnrate, MvNormal}(
        MvNormal([0.0,0.0], [1.0,0.1])
    )
)

set_desired_speed!(models[:alice],   12.0)
set_desired_speed!(models[:bob],     10.0)
set_desired_speed!(models[:charlie],  8.0)

scenes = simulate(scene, roadway, models, nticks, timestep)
# md nothing # hide

# An animation of the simulation can be rendered using the `Reel` package

using Reel
using Printf

camera = TargetFollowCamera(:alice; zoom=10.)

animation = roll(fps=1.0/timestep, duration=nticks*timestep) do t, dt
    i = Int(floor(t/dt)) + 1
    update_camera!(camera, scenes[i])
    renderables = [
        roadway,
        (FancyCar(car=veh, color=car_colors[veh.id]) for veh in scenes[i])...,
        IDOverlay(x_off=-2, y_off=1, scene=scenes[i]),
        TextOverlay(text=[@sprintf("time: %.1fs", t)], pos=VecE2(40,40), font_size=24)
    ]
    render(renderables, camera=camera)
end

#md write("animated_stadium.gif", animation) # hide
#md # ![animated stadium with cars](animated_stadium.gif)

# Alternatively, one can also use the `Interact` framework to inspect the simulation record interactively.

#md using Interact
#md using Reel
#md using Blink
#md
#md w = Window()
#md viz = @manipulate for step in 1 : length(scenes)
#md     render([roadway, scenes[step]])
#md end
#md body!(w, viz)

# The simulation results can be saved to a text file. Only entities using `VehicleState` and `VehicleDef` are supported out of the box.
# If you wish to do IO operation with different states and definition types you have to implement `Base.write` and `Base.read` for those new types.

open("2Dstadium_listrec.txt", "w") do io
    write(io, scenes)
end

# The trajectory data file can be loaded in a similar way. You need to specify the output type.

loaded_scenes = open("2Dstadium_listrec.txt", "r") do io
    read(io, Vector{EntityScene{VehicleState, VehicleDef, String}})
end
render([roadway, loaded_scenes[1]])
#md write("stadium_with_cars_loaded.svg", snapshot) # hide
#md # ![stadium with cars after load](stadium_with_cars_loaded.svg)


function sparse_cube(datacube)
    radiance = []
    for i in 1:size(datacube)[3]
        push!(radiance,sparse(datacube[:,:,i]))
    end
    return VectorOfArray(radiance)
end


function render(f::Function, cam::C, sampler = SingleSampler(T);
                save_interval::Number=Inf, filename::String="render.png",
                resolution::Integer=1) where {T,C<:Camera{T}}
    w,h = width(cam),height(cam)
    img = zeros(RGB{eltype(cam)}, (w,h))

    W = weight(sampler)
    t₀ = time()
    @showprogress "Rendering: " for i = 1:resolution:w
        for j = 1:resolution:h
            img[i,j] = W*mapreduce(ij -> f(cam[ij...]), +, sampler(i,j))
            if resolution > 1
                n = resolution-1
                for i′ = i:i+n
                    for j′ = j:j+n
                        img[i′,j′] = img[i,j]
                    end
                end
            end
            now = time()
            if now - t₀ > save_interval
                save_clamped(filename, img)
                t₀ = now
            end
        end
    end

    img
end

function render(cam::C, acc::A,
                shade::Function, background::RGB{T},
                args...; kwargs...) where {T,C<:Camera{T},O,A<:Accelerator{T,O}}
    render(cam, args...; kwargs...) do ray::Ray
        hit = Intersection(ray, O)
        intersect!(hit, acc)
        is_hit(hit) ? shade(hit) : background
    end
end

export render


#=
Cross-correlation plot for the paper.
Depending on the `absolute` variable, compare velocity with MTD or with
abs(MTD).

Integrates with DrWatson to not re-produce the data. Change `force` to
true to re-calculate correlations.
=#
using DrWatson
quickactivate(@__DIR__, "VelocityMTD")
using MusicManipulations, PyPlot, StatsBase, Statistics, Parameters
include(srcdir("correlations.jl"))
include(srcdir("extract_data.jl"))
include(srcdir("style.jl"))

datasets = ["tapping", "uwe", "pgmusic", "playalongs", "playalongs"]
datanames = ["tapping",  "Uwe Meile (piano)", "PG music (piano)", "drums playalongs", "drums playal. (lead)"]
subsets = ["all", "melody", "melody", "all", "lead"]
absolute = true

pcount = 0
fig = figure(figsize = (2figx, 3figx/4))
axρ = subplot(311)
axm = subplot(312)
axmn = subplot(313)
xlocs = zeros(length(datasets))
total = zeros(Int, length(datasets))
significant = zeros(Int, length(datasets))

for i in 1:length(datasets)
    # i == 3 && continue
    dataset = datasets[i]
    subset = subsets[i]
    c = @dict dataset subset absolute

    file, = produce_or_load(
        datadir("crosscor"), c, correlation_intervals;
        prefix = "crosscor",
        force = false
    )

    @unpack ρs, ρus, ρds, MIs, centers, MIus, MIds = file

    L = length(ρs)
    l = pcount + 1
    r = l + L - 1

    axρ.errorbar(l:r, ρs, yerr = hcat(ρus,ρds)',
    ecolor = "C0", color = "C$(i)", ls = "None", marker = "o")

    bot = mean(ρs) - mean(ρds)
    top = mean(ρs) + mean(ρus)
    rec = matplotlib.patches.Rectangle((l, bot), r-l, top-bot, color = "C$(i)",
    alpha = 0.2)
    # axρ.add_patch(rec)

    axρ.plot([l, r], [mean(ρs), mean(ρs)], color = "C$(i)", lw = 2, alpha = 0.75)

    # Plot default MI
    axm.plot([l, r], fill(mean(MIs),2), color = "C$(i)",
    lw = 2, alpha = 0.75)

    axm.errorbar(l:r, centers, yerr = hcat(MIus,MIds)',
    ecolor = "C0", color = "C0", ls = "None")
    axm.plot(l:r, MIs,  color = "C$(i)", ls = "None",
    marker = "o")
    y = 0.1
    # axm.axhline(0, color = "black", ls = "dashed", lw = 1.0)
    axm.set_ylim(0, 0.9)

    # plot normalized MI
    # make normalized mi
    no = zeros(length(MIs))
    for i in 1:length(no)
        mic = MIs[i] - centers[i]
        no[i] = if mic < 0
            mic/(MIds[i])
        else
            mic/(MIus[i])
        end
        no[i] = abs(no[i])
    end
    # plot it
    axmn.plot(l:r, no,  color = "C$(i)", ls = "None",
    marker = "o")
    # average
    axmn.plot([l, r], fill(mean(no), 2), color = "C$(i)",
    lw = 2, alpha = 0.75)
    axmn.axhline(1, color = "black", ls = "dashed", lw = 0.5)
    y = 3
    axmn.set_ylim(0, 4)

    axmn.text((r+l)/2, y, datanames[i], ha="center", va="center", color = "C$i")

    global pcount += L+1
    total[i] = length(no)
    significant[i] = count(n -> n > 1, no)
end

axρ.axhline(0, color = "black", ls = "dashed", lw = .5)
axρ.set_xticklabels([])
axρ.set_yticks(-0.3:0.2:0.3)
axρ.set_ylabel("\$\\rho\$")
axρ.set_title("correlation between velocity and $(absolute ? "abs(MTD)" : "MTD")")

axm.set_yticks(0:0.2:0.8)
axm.set_xticklabels([])
axm.set_ylabel("MI")

axmn.set_yticks(0:4)
axmn.set_ylabel("MI (norm.)")
axmn.set_xlabel("recording #")

#
for ax in fig.axes
    ax.minorticks_on()
    ax.tick_params(axis="y", which="minor", left=false)
    ax.set_xlim(0, pcount)
    ax.tick_params(axis="x", width =2, length=10)
    ax.tick_params(axis="x", which="minor", width = 1.5, length = 8)
end
add_identifiers!(fig; loc = (0.995, 0.975))
fig.align_ylabels()
fig.tight_layout()
fig.subplots_adjust(hspace = 0.2)
# fig.savefig(papersdir("figs", "crosscorrelations"))
fig.savefig(plotsdir("crosscorrelations", "absolute=$(absolute)"))


@testset "Maximum Gen Profit function" begin
    sys = c_sys_case4()
    bid0 = 0.3
    @test maxGenProfit_tcrdd(sys, bid0) == (true, 0.60892, 0.3, 1.0, 1, 1)
end