using Tokenize.Tokens: GREATER, LESS, GREATER_EQ, GREATER_THAN_OR_EQUAL_TO, LESS_EQ, LESS_THAN_OR_EQUAL_TO
using CSTParser: MacroName
begin
    struct ObsoleteVersionCheck; vers; end
    register(ObsoleteVersionCheck, Deprecation(
        "This version check is for a version of julia that is no longer supported by this package",
        "julia",
        typemin(VersionNumber), typemin(VersionNumber), typemax(VersionNumber)
    ))
    ObsoleteVersionCheck() = ObsoleteVersionCheck(Pkg.Reqs.parse(IOBuffer("julia 0.6"))["julia"])

    const comparisons = Dict(
         GREATER                  => >,
         LESS                     => <,
         GREATER_EQ               => >=,
         GREATER_THAN_OR_EQUAL_TO => ≥,
         LESS_EQ                  => <=,
         LESS_THAN_OR_EQUAL_TO    => ≤,
         #= TODO:
         EQEQ                     => ==,
         EQEQEQ                   => ===,
         IDENTICAL_TO             => ≡,
         NOT_EQ                   => !=,
         NOT_EQUAL_TO             => ≠,
         NOT_IS                   => !==
         =#)

    function detect_ver_arguments(VERSION_arg, v_arg)
        isexpr(VERSION_arg, CSTParser.IDENTIFIER) || return nothing
        VERSION_arg.val == "VERSION" || return nothing
        isexpr(v_arg, CSTParser.x_Str) || return nothing
        isexpr(children(v_arg)[1], CSTParser.IDENTIFIER) || return nothing
        isexpr(children(v_arg)[2], CSTParser.LITERAL, Tokens.STRING) || return nothing
        children(v_arg)[1].val == "v" || return nothing
        VersionNumber(children(v_arg)[2].val)
    end

    function dep_for_vers(::Type{ObsoleteVersionCheck}, vers)
        ObsoleteVersionCheck(vers["julia"])
    end

    opcode(x::CSTParser.OPERATOR) = x.kind
    iscomparison(x::CSTParser.OPERATOR) = CSTParser.precedence(x) == 6
    iscomparison(x::OverlayNode) = iscomparison(x.expr)

    function process_comparison(comparison, dep, context)
        isexpr(comparison, CSTParser.BinaryOpCall) || return nothing
        iscomparison(children(comparison)[2]) || return nothing
        comparison = comparison.expr
        opc = opcode(children(comparison)[2])
        haskey(comparisons, opc) || return nothing
        r1 = detect_ver_arguments(children(comparison)[1], children(comparison)[3])
        if r1 !== nothing
            f = comparisons[opc]
            alwaystrue = all(interval->(f(interval.lower, r1) && f(interval.upper, r1)), dep.vers.intervals)
            alwaysfalse = all(interval->(!f(interval.lower, r1) && !f(interval.upper, r1)), dep.vers.intervals)
            @assert !(alwaystrue && alwaysfalse)
            return (alwaystrue, alwaysfalse)
        end
        r2 = detect_ver_arguments(children(comparison)[3], children(comparison)[1])
        if r2 !== nothing
            f = comparisons[opc]
            alwaystrue = all(interval->(f(interval.lower, r2) && f(interval.upper, r2)), dep.vers.intervals)
            alwaysfalse = all(interval->(!f(interval.lower, r2) && !f(interval.upper, r2)), dep.vers.intervals)
            @assert !(alwaystrue && alwaysfalse)
            return (alwaystrue, alwaysfalse)
        end
        return nothing
    end

    function in_statement_position(expr)
        p = parent(expr)
        isexpr(p, CSTParser.FileH) && return true
        if isexpr(p, CSTParser.Begin) || isexpr(p, CSTParser.Block)
            # Could still be in value position if we're the last value in the
            # block and the block itself is in value position
            nextsibling(expr) !== nothing && return true
            return in_statement_position(p)
        end
        return false
    end

    match(ObsoleteVersionCheck, CSTParser.BinaryOpCall) do x
        dep, expr, resolutions, context = x
        result = process_comparison(expr, dep, context)
        result === nothing && return
        (alwaystrue, alwaysfalse) = result
        (alwaysfalse || alwaystrue) || return
        resolve_boolean(dep, resolutions, context, expr, alwaystrue)
    end
end


using CUDAnative, GPUArrays, CuArrays
using CUDAnative: device!, devices, CuDevice
using CuYao
import CuYao: cu
CuArrays.allowscalar(false)

function cu(chem::QuantumMPS)
    QuantumMPS(chem.nbit_measure, chem.nbit_virtual, chem.nbit_ancilla, chem.circuit, chem.initial_reg |> cu, chem.input_state)
end


## Type Hierarchy ##
abstract AbstractConverter
abstract FrameByFrameConverter <: AbstractConverter
abstract TrajectoryConverter <: AbstractConverter

# vc performs voice conversion based on specified converter.
function vc(c::FrameByFrameConverter,
            fm::AbstractMatrix{Float64} # feature matrix
            )
    # Split src feature matrix to power and spectral features
    power, src =  fm[1,:], fm[2:end,:]

    T = size(src, 2)
    converted = similar(fm)

    # Perform feature mapping for each time
    for t=1:T
        @inbounds converted[2:end,t] = fvconvert(c, src[:,t])
    end

    # keep original power
    # note that it is assumed that 0-th coef. represents power coef.
    converted[1,:] = power

    converted
end

# vc performs trajectory-based feature conversion. To reduce computational
# complexiy in practice, we split the input sequence to a set of sub-sequences
# and perform trajectory-based conversion for each sub-sequence.
function vc(c::TrajectoryConverter,
            fm::AbstractMatrix{Float64};  # feature matrix
            )
    # Split src feature matrix to power and spectral features
    power, src =  fm[1,:], fm[2:end,:]

    # D = length/2 + 1: order of static spectral features + power
    D, T = size(src, 1)>>1 + 1, size(src, 2)
    converted = Array(eltype(fm), D, T)

    # Perform Trajectory-based mapping
    count = 0
    limit = length(c) # maximum length of sub-sequence
    while true
        b = count * limit + 1
        e = (count+1) * limit
        if e > T
            e = T
        end
        phrase = src[:,b:e]
        converted[2:end,b:e] = fvconvert(c, phrase)

        if e == T
            break
        end
        count += 1
    end

    # keep original power
    converted[1,:] = power

    converted
end


@with_kw mutable struct AdversarialOffPolicySolver <: Solver
    𝒮_pro::OffPolicySolver # Solver parameters for the protagonist
    𝒮_ant::OffPolicySolver # Solver parameters for the antagonist
    px::PolicyParams # Nominal disturbance policy
    train_pro_every::Int = 1
    train_ant_every::Int = 1
    log::Union{Nothing, LoggerParams} = nothing # The logging parameters
    i::Int = 0 # The current number of environment interactions
end

function POMDPs.solve(𝒮::AdversarialOffPolicySolver, mdp)
    # make some assertions on the parameters that should be shared
    @assert 𝒮.𝒮_pro.buffer == 𝒮.𝒮_ant.buffer
    @assert 𝒮.𝒮_pro.buffer_init == 𝒮.𝒮_ant.buffer_init
    @assert 𝒮.𝒮_pro.required_columns == 𝒮.𝒮_ant.required_columns
    @assert 𝒮.𝒮_pro.N == 𝒮.𝒮_ant.N
    @assert 𝒮.𝒮_pro.S == 𝒮.𝒮_ant.S
    @assert 𝒮.𝒮_pro.max_steps == 𝒮.𝒮_ant.max_steps

    # define shared parameters for easy reference
    buffer = 𝒮.𝒮_pro.buffer
    buffer_init = 𝒮.𝒮_pro.buffer_init
    ΔN = 𝒮.𝒮_pro.ΔN + 𝒮.𝒮_ant.ΔN
    N = 𝒮.𝒮_pro.N
    S = 𝒮.𝒮_pro.S
    max_steps = 𝒮.𝒮_pro.max_steps

    # Construct the training buffern for the protagonist and antagonist
    𝒟_pro = buffer_like(buffer, capacity=𝒮.𝒮_protagonist.c_opt.batch_size, device=device(𝒮.𝒮_protagonist.π))
    𝒟_ant = buffer_like(buffer, capacity=𝒮.𝒮_protagonist.c_opt.batch_size, device=device(𝒮.𝒮_protagonist.π))

    # Get the discount factor
    γ = Float32(discount(mdp))

    # Two samplers, one for the traditional ap
    s_pro = Sampler(mdp, 𝒮.𝒮_pro.agent, adversary=𝒮.px, S=S, max_steps=max_steps, required_columns=extra_columns(buffer))
    s_ant = Sampler(mdp, 𝒮.𝒮_pro.agent, adversary=𝒮.𝒮_ant.agent, S=S, max_steps=max_steps, required_columns=extra_columns(buffer))

    s_nom = Sampler(mdp_nom, 𝒮.π, S=𝒮.S, A=𝒮.A, max_steps=𝒮.max_steps, required_columns=setdiff(extra_columns(𝒮.buffer), [:x]))
    push!(𝒮.log.fns, (;kwargs...) -> Dict("nominal_undiscounted_return" => undiscounted_return(s_nom, Neps=10)))

    isnothing(𝒮.log.sampler) && (𝒮.log.sampler = s_pro)

    # Log the pre-train performance
    log(𝒮.log, 𝒮.i)

    # Fill the buffer with initial observations before training
    𝒮.i += fillto!(buffer, s_pro, buffer_init, i=𝒮.i, explore=true)
    𝒮.𝒮_pro.i = 𝒮.i
    𝒮.𝒮_ant.i = 𝒮.i

    # Loop over the desired number of environment interactions
    for 𝒮.i in range(𝒮.i, stop=𝒮.i + N - ΔN, step=ΔN)
        # Update the per solver iteration
        𝒮.𝒮_pro.i = 𝒮.i
        𝒮.𝒮_ant.i = 𝒮.i

        # Sample transitions into the replay buffer
        push!(buffer, steps!(s_pro, Nsteps=𝒮.𝒮_pro.ΔN, explore=true, i=𝒮.i))
        push!(buffer, steps!(s_ant, Nsteps=𝒮.𝒮_ant.ΔN, explore=true, i=𝒮.i))

        # callback for potentially updating the buffer
        𝒮.post_experience_callback(buffer)

        # Train the networks
        infos = []
        elapsed(𝒮.i + 1:𝒮.i + ΔN, 𝒮.train_pro_every) && merge!(infos, train_step(𝒮.𝒮_pro, 𝒟_pro, γ))
        elapsed(𝒮.i + 1:𝒮.i + ΔN, 𝒮.train_ant_every) && merge!(infos, train_step(𝒮.𝒮_ant, 𝒟_ant, γ))

        # Log the results
        log(𝒮.log, 𝒮.i + 1:𝒮.i + ΔN, aggregate_info(infos))
    end
    𝒮.i += 𝒮.ΔN
    𝒮.𝒮_pro.i = 𝒮.i
    𝒮.𝒮_ant.i = 𝒮.i

    𝒮.𝒮_pro.π
end


function pancam(filter::Symbol)
   filter == :L2 ? TSpec(PancamL2[:,1],PancamL2[:,2]) :
   filter == :L3 ? TSpec(PancamL3[:,1],PancamL3[:,2]) :
   filter == :L4 ? TSpec(PancamL4[:,1],PancamL4[:,2]) :
   filter == :L5 ? TSpec(PancamL5[:,1],PancamL5[:,2]) :
   filter == :L6 ? TSpec(PancamL6[:,1],PancamL6[:,2]) :
   filter == :L7 ? TSpec(PancamL7[:,1],PancamL7[:,2]) : nothing
end

const PancamL2 =
   [729.893842887474	   0.0;
    735.83864118896 	   0.0402542372881356;
    739.235668789809	   0.120762711864407;
    741.783439490446	   0.347457627118644;
    743.481953290871	   0.591101694915254;
    744.755838641189	   0.75635593220339;
    746.454352441614	   0.925847457627119;
    747.303609341826	   0.995762711864407;
    748.152866242038	   1.0;
    752.3991507431  	   0.976694915254237;
    755.796178343949	   0.976694915254237;
    756.645435244161	   0.955508474576271;
    757.494692144374	   0.938559322033898;
    757.91932059448	   0.915254237288136;
    758.768577494692	   0.864406779661017;
    759.617834394904	   0.805084745762712;
    760.891719745223	   0.722457627118644;
    761.740976645435	   0.574152542372881;
    763.014861995754	   0.461864406779661;
    764.713375796178	   0.288135593220339;
    767.261146496815	   0.154661016949153;
    768.95966029724 	   0.0911016949152542;
    771.082802547771	   0.0444915254237288;
    774.055201698514	   0.0233050847457627;
    776.602972399151	   0.0105932203389831;
    780.0           	   0.0]

const PancamL3 =
   [653.460721868365	   0.0;
    656.433121019108	   0.0211864406779661;
    658.980891719745	   0.0741525423728814;
    661.104033970276	   0.199152542372881;
    662.802547770701	   0.415254237288136;
    664.076433121019	   0.586864406779661;
    666.19957537155	   0.720338983050847;
    669.171974522293	   0.775423728813559;
    671.295116772824	   0.832627118644068;
    672.569002123142	   0.889830508474576;
    673.842887473461	   0.959745762711864;
    674.692144373673	   0.991525423728814;
    675.541401273885	   1.0;
    676.815286624204	   0.99364406779661;
    677.23991507431	   0.955508474576271;
    677.664543524416	   0.925847457627119;
    678.513800424629	   0.796610169491525;
    680.212314225053	   0.544491525423729;
    682.335456475584	   0.307203389830509;
    684.033970276009	   0.163135593220339;
    685.307855626327	   0.116525423728814;
    687.855626326964	   0.0550847457627119;
    689.129511677282	   0.0338983050847458;
    691.677282377919	   0.0148305084745763;
    695.074309978769	   0.0]

const PancamL4 =
   [580.000000000000    0.0;
    584.670912951168    0.0105932203389831;
    588.067940552017    0.0444915254237288;
    590.615711252654    0.152542372881356;
    593.588110403397    0.563559322033898;
    595.711252653928    0.930084745762712;
    596.985138004246    0.98728813559322;
    597.409766454352    0.997881355932203;
    598.259023354565    1.0;
    600.382165605096    0.991525423728814;
    602.929936305732    0.974576271186441;
    603.779193205945    0.974576271186441;
    604.628450106157    0.98728813559322;
    605.477707006369    0.98728813559322;
    606.751592356688    0.963983050847458;
    607.600849256900    0.930084745762712;
    608.874734607219    0.802966101694915;
    610.148619957537    0.629237288135593;
    611.847133757962    0.425847457627119;
    612.696390658174    0.298728813559322;
    613.970276008493    0.199152542372881;
    616.093418259023    0.103813559322034;
    618.216560509554    0.0593220338983051;
    619.915074309979    0.0338983050847458;
    622.462845010616    0.0169491525423729;
    625.010615711253    0.00635593220338983;
    630.106157112527    0.0]

const PancamL5 =
   [518.004246284501	   0.0;
    520.127388535032	   0.010593220339;
    521.40127388535	   0.025423728814;
    522.250530785563	   0.055084745763;
    524.7983014862	   0.415254237288;
    525.647558386412	   0.616525423729;
    526.496815286624	   0.690677966102;
    526.92144373673	   0.718220338983;
    529.044585987261	   0.737288135593;
    530.743099787686	   0.783898305085;
    531.592356687898	   0.824152542373;
    532.44161358811	   0.862288135593;
    533.290870488323	   0.904661016949;
    534.989384288747	   0.940677966102;
    537.112526539278	   0.959745762712;
    539.660297239915	   1.0;
    540.509554140127	   1.0;
    541.35881104034	   0.98093220339;
    541.783439490446	   0.96186440678;
    542.208067940552	   0.940677966102;
    542.632696390658	   0.887711864407;
    544.331210191083	   0.580508474576;
    545.605095541401	   0.341101694915;
    546.87898089172	   0.16313559322;
    548.152866242038	   0.091101694915;
    549.426751592357	   0.03813559322;
    550.700636942675	   0.01906779661;
    552.823779193206	   0.006355932203;
    560.042462845011	   0.0]

const PancamL6 =
   [450.063694267516	   0.0;
    455.15923566879	   0.004237288136;
    458.556263269639	   0.014830508475;
    459.405520169851	   0.01906779661;
    460.254777070064	   0.033898305085;
    461.528662420382	   0.080508474576;
    464.501061571125	   0.322033898305;
    465.350318471338	   0.411016949153;
    466.19957537155	   0.506355932203;
    467.473460721868	   0.563559322034;
    468.747346072187	   0.584745762712;
    470.021231422505	   0.622881355932;
    471.295116772824	   0.66313559322;
    472.144373673036	   0.701271186441;
    474.692144373673	   0.866525423729;
    476.390658174098	   0.911016949153;
    477.664543524416	   0.930084745763;
    481.486199575372	   0.930084745763;
    483.184713375796	   0.940677966102;
    484.883227176221	   0.963983050847;
    486.581740976645	   0.987288135593;
    488.28025477707	   1.0;
    489.978768577495	   0.991525423729;
    490.828025477707	   0.972457627119;
    491.677282377919	   0.947033898305;
    492.101910828025	   0.915254237288;
    492.951167728238	   0.877118644068;
    493.375796178344	   0.834745762712;
    493.80042462845	   0.786016949153;
    494.225053078556	   0.716101694915;
    495.498938428875	   0.516949152542;
    496.348195329087	   0.463983050847;
    498.046709129512	   0.28813559322;
    498.895966029724	   0.228813559322;
    500.594479830149	   0.177966101695;
    501.443736730361	   0.14406779661;
    503.142250530786	   0.08686440678;
    505.690021231423	   0.025423728814;
    509.936305732484	   0.008474576271;
    515.031847133758	   0.002118644068;
    524.7983014862	   0.0]

const PancamL7 =
   [380.0     	         0.0
    387.643312101911	   0.023305084746
    392.314225053079	   0.046610169492
    397.409766454352	   0.088983050847
    402.08067940552	   0.150423728814
    410.148619957537	   0.300847457627
    419.915074309979	   0.459745762712
    421.188959660297	   0.483050847458
    422.887473460722	   0.5
    428.407643312102	   0.578389830508
    433.07855626327	   0.677966101695
    436.475583864119	   0.777542372881
    439.023354564756	   0.817796610169
    440.72186836518	   0.83686440678
    447.091295116773	   0.932203389831
    450.488322717622	   1.0
    452.186836518047	   0.987288135593
    453.036093418259	   0.966101694915
    453.885350318471	   0.906779661017
    455.583864118896	   0.819915254237
    457.282377919321	   0.580508474576
    458.980891719745	   0.411016949153
    461.104033970276	   0.218220338983
    462.802547770701	   0.156779661017
    464.925690021231	   0.080508474576
    466.19957537155	   0.048728813559
    468.322717622081	   0.027542372881
    471.295116772824	   0.012711864407
    473.842887473461	   0.004237288136
    479.787685774947	   0.0]


using Flux
using Transformers
using Transformers.Basic

struct TransformersModel
    embed::Union{WordPositionEmbed,CompositeEmbedding}
    encoder::Union{Encoder,Bert}
    decoder::Decoder
    generator::Generator
end

Flux.@functor TransformersModel
Flux.trainable(transformers::TransformersModel) = (transformers.embed, transformers.encoder, transformers.decoder, transformers.generator)

function (transformers::TransformersModel)(
    vocabulary::Vocabulary{T},
    inputs::Vector{T},
    outputs::Vector{T}
)::AbstractMatrix{<:AbstractFloat} where T
    # Encode inputs.
    input_indices = vocabulary(inputs) |> gpu
    if typeof(transformers.embed) <: CompositeEmbedding
        input_indices = (
            tok = input_indices,
            segment = fill(1, length(input_indices))
        )
    end
    input_embedding = transformers.embed(input_indices)
    encoded_embedding = transformers.encoder(input_embedding)

    # Decode outputs.
    output_indices = vocabulary(outputs) |> gpu
    if typeof(transformers.embed) <: CompositeEmbedding
        output_indices = (
            tok = output_indices,
            segment = fill(1, length(output_indices))
        )
    end
    output_embedding = transformers.embed(output_indices)
    decoded_embedding = transformers.decoder(output_embedding, encoded_embedding)

    # Calculate probabilities for next token.
    return transformers.generator(decoded_embedding)
end

function TransformersModel(
    embed::Union{WordPositionEmbed,CompositeEmbedding},
    size::Integer,
    head::Integer,
    hs::Integer,
    ps::Integer,
    layer::Integer,
    vocab_size::Integer;
    act=relu,
    pdrop=0.1,
    length_normalization::AbstractFloat=0.0
)
    return TransformersModel(
        embed,
        Encoder(size, head, hs, ps, layer; act=act,pdrop=pdrop),
        Decoder(size, head, hs, ps, layer; act=act,pdrop=pdrop),
        Generator(size, vocab_size)
    )
end


module Tracking
    using
        DocStringExtensions,
        GNSSSignals,
        StaticArrays,
        TrackingLoopFilters,
        StructArrays,
        LoopVectorization
    using Unitful: upreferred, Hz, dBHz, ms
    import Base.zero, Base.length, Base.resize!

    export
        get_early,
        get_prompt,
        get_late,
        get_correlator,
        get_accumulators,
        get_num_accumulators,
        get_early_index,
        get_prompt_index,
        get_late_index,
        get_accumulator,
        get_carrier_doppler,
        get_carrier_phase,
        get_code_doppler,
        get_code_phase,
        get_correlator_sample_shifts,
        get_early_late_sample_spacing,
        get_early_late_index_shift,
        get_secondary_code_or_bit_found,
        get_correlator_carrier_phase,
        get_correlator_carrier_frequency,
        get_state,
        get_system,
        get_cn0,
        track,
        TrackingState,
        NumAnts,
        NumAccumulators,
        MomentsCN0Estimator,
        AbstractCN0Estimator,
        get_bits,
        get_num_bits,
        EarlyPromptLateCorrelator,
        #VeryEarlyPromptLateCorrelator,
        SecondaryCodeOrBitDetector,
        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")
end


mutable struct CounterfactualExplanation
    x::AbstractArray
    target::Number
    target_encoded::Union{Number, AbstractVector, Nothing}
    x′::AbstractArray
    data::DataPreprocessing.CounterfactualData
    M::Models.AbstractFittedModel
    generator::Generators.AbstractGenerator
    params::Dict
    search::Union{Dict,Nothing}
end

using .Counterfactuals
function Base.show(io::IO, z::CounterfactualExplanation)

    printstyled(io, "Factual: ", bold=true)
    println(io, "x=$(z.x), y=$(Counterfactuals.factual_label(z)), p=$(Counterfactuals.factual_probability(z))")
    printstyled(io, "Target: ", bold=true)
    println(io, "target=$(z.target), γ=$(z.params[:γ])")

    if !isnothing(z.search)
        printstyled(io, "Counterfactual outcome: ", bold=true)
        println(io, "x′=$(z.x′), y′=$(Counterfactuals.counterfactual_label(z)), p′=$(Counterfactuals.counterfactual_probability(z))")
        printstyled(io, "Converged: $(Counterfactuals.converged(z) ? "✅"  : "❌") ", bold=true)
        println("after $(Counterfactuals.total_steps(z)) steps.")
    else
        @info "Search not yet initatiated."
    end

end

"""
    CounterfactualExplanation(
        x::Union{AbstractArray,Int},
        target::Union{AbstractFloat,Int},
        data::CounterfactualData,
        M::Models.AbstractFittedModel,
        generator::Generators.AbstractGenerator,
        γ::AbstractFloat,
        T::Int
    )

Outer method to construct a `CounterfactualExplanation` structure.
"""
# Outer constructor method:
function CounterfactualExplanation(
    x::Union{AbstractArray,Int},
    target::Union{AbstractFloat,Int},
    data::CounterfactualData,
    M::Models.AbstractFittedModel,
    generator::Generators.AbstractGenerator,
    γ::AbstractFloat,
    T::Int
)
    # Factual:
    x = typeof(x) == Int ? select_factual(data, x) : x
    # Counterfactual:
    x′ = copy(x)  # start from factual

    # Parameters:
    params = Dict(
        :γ => γ,
        :T => T,
        :mutability => DataPreprocessing.mutability_constraints(data)
    )

    return CounterfactualExplanation(x, target, nothing, x′, data, M, generator, params, nothing)

end

# Convenience methods:

# 0) Utils
"""
    output_dim(counterfactual_explanation::CounterfactualExplanation)

A convenience method that computes the output dimension of the predictive model.
"""
output_dim(counterfactual_explanation::CounterfactualExplanation) = size(Models.probs(counterfactual_explanation.M, counterfactual_explanation.x))[1]

using Flux
"""
    encode_target(counterfactual_explanation::CounterfactualExplanation)

A convenience method to encode the target variable, if necessary.
"""
function encode_target(counterfactual_explanation::CounterfactualExplanation)
    out_dim = output_dim(counterfactual_explanation)
    target = counterfactual_explanation.target
    return out_dim > 1 ? Flux.onehot(target, 1:out_dim) : target
end

# 1) Factual values
"""
    factual(counterfactual_explanation::CounterfactualExplanation)

A convenience method to get the factual value.
"""
factual(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.x

"""
    factual_probability(counterfactual_explanation::CounterfactualExplanation)

A convenience method to compute the class probabilities of the factual value.
"""
factual_probability(counterfactual_explanation::CounterfactualExplanation) = Models.probs(counterfactual_explanation.M, counterfactual_explanation.x)

"""
    factual_label(counterfactual_explanation::CounterfactualExplanation)

A convenience method to get the predicted label associated with the factual value.
"""
function factual_label(counterfactual_explanation::CounterfactualExplanation)
    p = factual_probability(counterfactual_explanation)
    out_dim = size(p)[1]
    y = out_dim == 1 ? round(p[1]) : Flux.onecold(p,1:out_dim)
    return y
end

# 2) Counterfactual values:
"""
    initialize!(counterfactual_explanation::CounterfactualExplanation)

A subroutine that intializes the counterfactual search.
"""
function initialize!(counterfactual_explanation::CounterfactualExplanation)

    # Encode target:
    counterfactual_explanation.target_encoded = encode_target(counterfactual_explanation)

    # Initialize search:
    counterfactual_explanation.search = Dict(
        :iteration_count => 1,
        :times_changed_features => zeros(length(counterfactual_explanation.x)),
        :path => [counterfactual_explanation.x′],
        :terminated => threshold_reached(counterfactual_explanation),
        :converged => threshold_reached(counterfactual_explanation),
        :other => nothing
    )

    if counterfactual_explanation.search[:terminated]
        @info "Factual already in target class and probability exceeds threshold γ."
    end

end

"""
    counterfactual(counterfactual_explanation::CounterfactualExplanation)

A convenience method to get the counterfactual value.
"""
counterfactual(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.x′

"""
    counterfactual_probability(counterfactual_explanation::CounterfactualExplanation)

A convenience method to compute the class probabilities of the counterfactual value.
"""
counterfactual_probability(counterfactual_explanation::CounterfactualExplanation) = Models.probs(counterfactual_explanation.M, counterfactual_explanation.x′)

"""
    counterfactual_label(counterfactual_explanation::CounterfactualExplanation)

A convenience method to get the predicted label associated with the counterfactual value.
"""
function counterfactual_label(counterfactual_explanation::CounterfactualExplanation)
    p = counterfactual_probability(counterfactual_explanation)
    out_dim = size(p)[1]
    y = out_dim == 1 ? round(p[1]) : Flux.onecold(p,1:out_dim)
    return y
end

# 3) Search related methods:
"""
    terminated(counterfactual_explanation::CounterfactualExplanation)

A convenience method to determine if the counterfactual search has terminated.
"""
terminated(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.search[:terminated]

"""
    converged(counterfactual_explanation::CounterfactualExplanation)

A convenience method to determine if the counterfactual search has converged.
"""
converged(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.search[:converged]

"""
    total_steps(counterfactual_explanation::CounterfactualExplanation)

A convenience method that returns the total number of steps of the counterfactual search.
"""
total_steps(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.search[:iteration_count]

"""
    path(counterfactual_explanation::CounterfactualExplanation)

A convenience method that returns the entire counterfactual path.
"""
path(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.search[:path]

"""
    target_probs(counterfactual_explanation::CounterfactualExplanation, x::Union{AbstractArray, Nothing}=nothing)

Returns the predicted probability of the target class for `x`. If `x` is `nothing`, the predicted probability corresponding to the counterfactual value is returned.
"""
function target_probs(counterfactual_explanation::CounterfactualExplanation, x::Union{AbstractArray, Nothing}=nothing)

    p = !isnothing(x) ? Models.probs(counterfactual_explanation.M, x) : counterfactual_probability(counterfactual_explanation)
    target = counterfactual_explanation.target

    if length(p) == 1
        h(x) = ifelse(x==-1,0,x)
        if target ∉ [0,1] && target ∉ [-1,1]
            throw(DomainError("For binary classification expecting target to be in {0,1} or {-1,1}."))
        end
        # If target is binary (i.e. outcome 1D from sigmoid), compute p(y=0):
        p = vcat(1.0 .- p, p)
        # Choose first (target+1) row if target=0, second row (target+1) if target=1:
        p_target = p[Int(h(target)+1),:]
    else
        if target < 1 || target % 1 !=0
            throw(DomainError("For multi-class classification expecting `target` ∈ ℕ⁺, i.e. {1,2,3,...}."))
        end
        # If target is multi-class, choose corresponding row (e.g. target=2 -> row 2)
        p_target = p[Int(target),:]
    end
    return p_target
end

"""
    apply_mutability(Δx′::AbstractArray, counterfactual_explanation::CounterfactualExplanation)

A subroutine that applies mutability constraints to the proposed vector of feature perturbations.
"""
function apply_mutability(Δx′::AbstractArray, counterfactual_explanation::CounterfactualExplanation)

    mutability = counterfactual_explanation.params[:mutability]
    # Helper functions:
    both(x) = x
    increase(x) = ifelse(x<0,0,x)
    decrease(x) = ifelse(x>0,0,x)
    none(x) = 0
    cases = (both = both, increase = increase, decrease = decrease, none = none)

    # Apply:
    Δx′ = [getfield(cases, mutability[d])(Δx′[d]) for d in 1:length(Δx′)]

    return Δx′

end

"""
    threshold_reached(counterfactual_explanation::CounterfactualExplanation)

A convenience method that determines of the predefined threshold for the target class probability has been reached.
"""
threshold_reached(counterfactual_explanation::CounterfactualExplanation) = target_probs(counterfactual_explanation)[1] >= counterfactual_explanation.params[:γ]

"""
    steps_exhausted(counterfactual_explanation::CounterfactualExplanation)

A convenience method that checks if the number of maximum iterations has been exhausted.
"""
steps_exhausted(counterfactual_explanation::CounterfactualExplanation) = counterfactual_explanation.search[:iteration_count] == counterfactual_explanation.params[:T]

"""
    get_counterfactual_state(counterfactual_explanation::CounterfactualExplanation)

A subroutine that is used to take a snapshot of the current counterfactual search state. This snapshot is passed to the counterfactual generator.
"""
function get_counterfactual_state(counterfactual_explanation::CounterfactualExplanation)
    counterfactual_state = Generators.CounterfactualState(
        counterfactual_explanation.x,
        counterfactual_explanation.target_encoded,
        counterfactual_explanation.x′,
        counterfactual_explanation.M,
        counterfactual_explanation.params,
        counterfactual_explanation.search
    )
    return counterfactual_state
end

"""
    update!(counterfactual_explanation::CounterfactualExplanation)

An important subroutine that updates the counterfactual explanation. It takes a snapshot of the current counterfactual search state and passes it to the generator. Based on the current state the generator generates perturbations. Various constraints are then applied to the proposed vector of feature perturbations. Finally, the counterfactual search state is updated.
"""
function update!(counterfactual_explanation::CounterfactualExplanation)

    counterfactual_state = get_counterfactual_state(counterfactual_explanation)

    # Generate peturbations:
    Δx′ = Generators.generate_perturbations(counterfactual_explanation.generator, counterfactual_state)
    Δx′ = apply_mutability(Δx′, counterfactual_explanation)
    Δx′ = reshape(Δx′, size(counterfactual_explanation.x′))
    x′ = counterfactual_explanation.x′ + Δx′
    x′ = DataPreprocessing.apply_domain_constraints(counterfactual_explanation.data, x′)
    counterfactual_explanation.x′ = x′ # update counterfactual

    # Updates:
    counterfactual_explanation.search[:path] = [counterfactual_explanation.search[:path]..., counterfactual_explanation.x′]
    counterfactual_explanation.search[:mutability] = Generators.mutability_constraints(counterfactual_explanation.generator, counterfactual_state)
    counterfactual_explanation.search[:times_changed_features] += reshape(Δx′ .!= 0, size(counterfactual_explanation.search[:times_changed_features])) # update number of times feature has been changed
    counterfactual_explanation.search[:iteration_count] += 1 # update iteration counter
    counterfactual_explanation.search[:converged] = threshold_reached(counterfactual_explanation)
    counterfactual_explanation.search[:terminated] = counterfactual_explanation.search[:converged] || steps_exhausted(counterfactual_explanation) || Generators.conditions_satisified(counterfactual_explanation.generator, counterfactual_state)
end


using .JSort

using SharedArrays

using JLD
using DelimitedFiles
using Interpolations
using Polynomials
import Base.show
import YAML

function constructzrange(path)
    data = readdlm(path, skipstart=2)
    itp = LinearInterpolation(convert(Array{Float64, 1}, 1000data[:, 1]),
                              convert(Array{Float64, 1}, data[:, 2]),  # Convert to keV
                              extrapolation_bc=Line())
    return itp
end

function readgainshift(path::Nothing)
    gainshift = [Float32.([0.0, 1.0]) for f in 1:8, b in 1:8]
end

function readgainshift(path)
    println(path)
    gainshift, header = readdlm(path, ',', header=true)
    # Check the file
    b = f = 1
    gainshiftarray = [Float32[] for f in 1:8, b in 1:8]
    for (i, row) in enumerate(eachrow(gainshift))
        # if Int(row[1]) ≠ f || Int(row[2]) ≠ b
        #     error("Gainshift file is corrupt at row $i = $row")
        # end
        f = Int(row[1])
        b = Int(row[2])
        gainshiftarray[f, b] = row[3:end]
    #     f += 1
    #     if f > 8
    #         f = 1
    #         b += 1
    #     end
    end
    gainshiftarray
end

readgainshift_gamma(paths::Vector{Nothing}) = []
readgainshift_gamma(paths::Nothing) = []
function readgainshift_gamma(paths::Vector{T}) where T
    [readgainshift_gamma(path) for path in paths]
end

function readgainshift_gamma(path::AbstractString)
    gainshift, header = readdlm(path, ',', header=true)
    gainshiftarray = zeros((33, 3))
    channels = Set(collect(1:33))
    for (i, row) in enumerate(eachrow(gainshift))
        try
            channel = round(Int, row[1])
            gainshiftarray[channel, :] = row[2:end]
            pop!(channels, channel)
        catch e
            println(e)
            error("Gamma gainshift file corrupt at row $i = $row")
        end
    end
    if length(channels) > 0
        @warn "Channels $channels are unspecified $path"
    end
    gainshiftarray
end

struct Parameters
    parameterpath::String
    raw  # The YAML file
    calibrator::Dict{Symbol, AbstractArray{Poly}}
    isotope
    datafiles
    exfromeΔe::Array{Float64, 2}
    excorr::Array{Float64, 2}
    thickrange::Array{Float64, 1}
    tnaicorresi::Array{T, 1} where T<:Number
    eΔerect
    percent
    range::Interpolations.Extrapolation{Float64,1,Interpolations.GriddedInterpolation{Float64,1,Float64,Interpolations.Gridded{Interpolations.Linear},Tuple{Array{Float64,1}}},Interpolations.Gridded{Interpolations.Linear},Interpolations.Line{Nothing}}
    cores
    savepath::String
    gates::Dict{Symbol,Gate}

    function Parameters(parameterpath::T) where T <: AbstractString
        join(path) = joinpath(splitdir(parameterpath)[1], path)
        join(path::Nothing) = nothing
        parameters    = YAML.load_file(parameterpath)
        isotope       = getparameter(parameters, "isotope")
        datafiles     = getparameter(parameters, "datafiles", T=Vector{String}) .|> join
        exfromeΔe     = getparameter(parameters, "ex from ede", repeat([0.0 0.0 0.0], 8))
        excorr        = getparameter(parameters, "excitation energy correction", repeat([0.0 1.0], 8))
        thickrange    = getparameter(parameters, "thick range", T=Array{Float64})
        tnaicorresi   = getparameter(parameters, "tnai corr enai", [0, 0, 1, 0])
        eΔerect       = getparameter(parameters, "ede rectangle")
        percent       = getparameter(parameters, "percent to read", 1.0)
        cores         = getparameter(parameters, "cpu cores", 1)
        savepath      = getparameter(parameters, "save path", "sirius") |> join

        # Calibration
        egainshift  = readgainshift(join(parameters["e gainshift file"]))
        Δegainshift = readgainshift(join(parameters["de gainshift file"]))
        γgainshift  = parameters["gamma gainshift file"] .|> join |> readgainshift_gamma
        tgainshift  = parameters["time gainshift file"]  .|> join |> readgainshift_gamma
        calibrator  = Dict{Symbol, Calibrator}(
            :e  => combinegainshift(egainshift, 1:8, 1:8),
            :Δe => combinegainshift(Δegainshift, 1:8, 1:8, initial=[0.0, 2.5]),
            :γ  => combinegainshift(γgainshift, 1:33),
            :t  => combinegainshift(tgainshift, 1:33, initial=[0.0, 1.0])
        )
        # Gates # Fix case for empty gates
        gates = getgates(parameters["gates"])
        range = getparameter(parameters, "range file") |> join |> constructzrange

        new(parameterpath, parameters, calibrator, isotope, datafiles, exfromeΔe, excorr, thickrange,
            tnaicorresi, eΔerect, percent, range, cores, savepath, gates)
    end
end

function show(io::IO, a::Parameters)
    println("Parameters for $(a.isotope)")
    println("Loaded from $(a.parameterpath).")
    println("e-calibration coefficients:")
    show(io, a.calibrator[:e])
    println("\nΔe-calibration coefficients:")
    show(io, a.calibrator[:Δe])
    println("\ngamma calibration coefficients:")
    show(io, a.calibrator[:γ])
    println("\ntime calibration coefficients:")
    show(io, a.calibrator[:t])
    println("Gates:")
    for gate in a.gates
        println("$(gate.first) : $(gate.second.low) to $(gate.second.high)")
    end
end

function show(io::IO, a::AbstractVector{T} where T<: Poly)
    for i in eachindex(a)
        print("$i: ")
        println(a[i])
    end
end

function show(io::IO, a::AbstractMatrix{T} where T<: Poly)
    for row in 1:size(a, 1)
        for col in 1:size(a, 2)
            print("$row, $col: ")
            println(a[row, col])
        end
    end
end

function getparameter(parameters, parameterkey, defaultvalue; constraint=nothing)
    if parameterkey ∉ keys(parameters)
        parameter = defaultvalue
    else
        parameter = parameters[parameterkey]
        # Reshape arrays
        if typeof(defaultvalue) <: Array{Float64}
            parameter = reshape(parameter, size(defaultvalue))
        else
            parameter = convert(typeof(defaultvalue), parameter)
        end
        if !isnothing(constraint)
            if !constraint(parameter)
                error("Parameter $parameterkey does not satisfy $constraint")
            end
        end
    end
    return parameter
end

function getparameter(parameters, parameterkey; T=nothing, constraint=nothing)
    if parameterkey ∉ keys(parameters)
        error("Parameter $parameterkey is required")
    else
        parameter = parameters[parameterkey]
        if !isnothing(T)
            parameter = convert(T, parameter)
        end
        if !isnothing(constraint)
            if !constraint(parameter)
                error("Parameter $parameterkey does not satisfy $constraint")
            end
        end
    end
    return parameter
end

function getpath(parameters::Parameters, key::Symbol)
    basepath = splitdir(parameters.parameterpath)[1]
    path = joinpath(basepath, getfield(parameters, key))
end

function getpath(parameters::Parameters, key)
    basepath = splitdir(parameters.parameterpath)[1]
    filepath = parameters.raw[key]
    path = joinpath(basepath, filepath)
end

### HANDLE GAINSHIFT ###
function combinegainshift(gainshifts::AbstractVector, channels;
                          initial = [0.0, 5.0],
                          docalibrate = true)::Vector{Poly}
    #@assert length(channels) == length(gainshift)
    # Combine all gainshift operations into a single
    # loop for a given channel. Gives a 3x speedup
    # Try to unroll the loop
    calibrators = Poly[]
    for channel in channels
        p = Poly(initial)
        if docalibrate
            for gainshift in gainshifts
                p = Poly(gainshift[channel, :])(p)
            end
        end
        # c = coeffs(p)
        # n = length(c)
        # f = let n=n, c=c
            # x -> begin
                # s = 0.0
                # @inbounds for i in 1:n
                    # s += c[i]*x^(i-1)
                # end
                # s
            # end
        # end
        push!(calibrators, p)
    end
    calibrators
end

function combinegainshift(gainshifts::AbstractMatrix, channelsf, channelsb;
                          initial = [0.0, 5.0],
                          docalibrate = true)::Matrix{Poly}
    combinegainshift([gainshifts], channelsf, channelsb, initial=initial)
end
function combinegainshift(gainshifts::AbstractVector, channelsf, channelsb;
                          initial = [0.0, 5.0],
                          docalibrate = true)::Matrix{Poly}
    # Combine all gainshift operations into a single
    # loop for a given channel. Gives a 3x speedup
    # Try to unroll the loop
    calibrators  = Matrix{Poly}(undef, length(channelsf), length(channelsb))
    for f in channelsf, b in channelsf
        p = Poly(initial)
        if docalibrate
            for gainshift in gainshifts
                p = Poly(gainshift[f, b])(p)
            end
        end
        # TODO why not just use Poly(x)??
        # c = coeffs(p)
        # n = length(c)
        # func = let n=n, c=c
            # x -> begin
                # s = 0.0
                # @inbounds for i in 1:n
                    # s += c[i]*x^(i-1)
                # end
                # s
            # end
        # end
        calibrators[f, b] = p
    end
    calibrators
end

function savegainshift_gamma(path, coefficients)
    println(coefficients)
end

function savegainshift_gamma(parameters::Parameters, coefficients)
    path = getpath(parameters, "gamma gainshift file")
    savegainshift_gamma(path, coefficients)
end

function savegainshift_gamma(parameters::Parameters, filename::String, coefficients)
    path = getpath(parameters, :savepath)
    path = joinpath(path, filename)
    # pad different lengthed polynomials
    N = maximum(map(length, coefficients))
    coefficients = map(x -> begin
            diff = length(x) - N
            if diff < 0
                x = [x..., zeros(abs(diff))...]
            else
                x
            end
        end, coefficients)

    open(path, "w") do io
        header = ["channel" "shift" "gain" "sqgain"]
        @show coefficients
        writedlm(io, [header, coefficients...], ',', header=true)
    end
end

function loadlabr(parameters::Parameters)::Vector{TrueEvent}
    return loadlabr(parameters.savepath)
end

function loadlabr(path::AbstractString)::Vector{TrueEvent}
    fhandle = open(joinpath(path, "labr.bin"))
    data = Base.read(fhandle) |> x -> reinterpret(UInt32, x)
    i = 1
    labr = TrueEvent[]
    while i ≤ length(data)
        f, b = data[i:i+1]
        i += 2
        de, e = convert.(Float32, data[i:i+1])
        i += 2
        labrcount = data[i]
        if labrcount + i >= length(data)
            println("Reading failed!")
        end
        i += 1
        adctdcs = ADCTDC[]
        for j in 1:labrcount
            channel = data[i]
            i += 1
            adc, tdc = convert.(Float32, data[i:i+1])
            push!(adctdcs, ADCTDC(channel, adc, tdc))
            i += 2
        end
        push!(labr, TrueEvent(f, b, de, e, adctdcs))
    end
    labr
end

function savecoefficients(parameters, path, coefficients, channels)
    all_coefficients = [[0.0, 1.0, 0.0] for i in 1:33]
    all_coefficients[channels] = coefficients
    # Tag each coefficient with corresponding channel
    all_coefficients = [[i, val...] for (i, val) in enumerate(all_coefficients)]
    savegainshift_gamma(parameters, path * ".csv", all_coefficients)
end

function savecoefficients_eΔe(parameters, coefficients_e, coefficients_Δe, path::Tuple{String, String}=("egainshift.csv", "degainshift.csv"))
    # Coefficients formated as  front, back, x0   x1   y0   y1
    savepath = getpath(parameters, :savepath)
    spath = joinpath(savepath, path[1])
    open(spath, "w") do io
        header = ["front" "back" "shift" "gain"]
        @show coefficients_e
        writedlm(io, [header, coefficients_e...], ',', header=true)
    end

    spath = joinpath(savepath, path[2])
    open(spath, "w") do io
        header = ["front" "back" "shift" "gain"]
        @show coefficients_Δe
        writedlm(io, [header, coefficients_Δe...], ',', header=true)
    end
end

function readede(path::String; T=Float32)
    data = open(path, "r") do io
        len = read(io, T) |> Int
        data = Matrix{T}(undef, len, 2)
        read!(io, data)
        data
    end
end


# -*- coding: utf-8 -*-
# AES tests from:
# https://git.lysator.liu.se/nettle/nettle/blob/master/testsuite/aes-test.c

# AES 128
for (key,text,encrypted) in [
    (
        "00010203050607080A0B0C0D0F101112",
        "506812A45F08C889B97F5980038B8359",
        "D8F532538289EF7D06B506A4FD5BE9C9"
    ),(
        "14151617191A1B1C1E1F202123242526",
        "5C6D71CA30DE8B8B00549984D2EC7D4B",
        "59AB30F4D4EE6E4FF9907EF65B1FB68C"
    ),(
        "28292A2B2D2E2F30323334353738393A",
        "53F3F4C64F8616E4E7C56199F48F21F6",
        "BF1ED2FCB2AF3FD41443B56D85025CB1",
    ),(
        "A0A1A2A3A5A6A7A8AAABACADAFB0B1B2",
        "F5F4F7F684878689A6A7A0A1D2CDCCCF",
        "CE52AF650D088CA559425223F4D32694",
    ),(
        "2b7e151628aed2a6abf7158809cf4f3c",
        "6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710",
        "3ad77bb40d7a3660a89ecaf32466ef97f5d3d58503b9699de785895a96fdbaaf43b1cd7f598ece23881b00e3ed0306887b0c785e27e8ad3f8223207104725dd4",
    )
]
    @test encrypt("aes128", hex2bytes(key), hex2bytes(text)) == hex2bytes(encrypted)
    @test decrypt("aes128", hex2bytes(key), hex2bytes(encrypted)) == hex2bytes(text)
end

# AES192
for (key,text,encrypted) in [
    (
        "00010203050607080A0B0C0D0F10111214151617191A1B1C",
        "2D33EEF2C0430A8A9EBF45E809C40BB6",
        "DFF4945E0336DF4C1C56BC700EFF837F",
    ),(
        "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
        "6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710",
        "bd334f1d6e45f25ff712a214571fa5cc974104846d0ad3ad7734ecb3ecee4eefef7afd2270e2e60adce0ba2face6444e9a4b41ba738d6c72fb16691603c18e0e",
    )
]
    @test encrypt("aes192", hex2bytes(key), hex2bytes(text)) == hex2bytes(encrypted)
    @test decrypt("aes192", hex2bytes(key), hex2bytes(encrypted)) == hex2bytes(text)
end

# AES256
for (key,text,encrypted) in [
    (
        "00010203050607080A0B0C0D0F10111214151617191A1B1C1E1F202123242526",
        "834EADFCCAC7E1B30664B1ABA44815AB",
        "1946DABF6A03A2A2C3D0B05080AED6FC",
    ),(
        "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
        "6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710",
        "f3eed1bdb5d2a03c064b5a7e3db181f8591ccb10d410ed26dc5ba74a31362870b6ed21b99ca6f4f9f153e7b1beafed1d23304b7a39f9f3ff067d8d8f9e24ecc7",
    )
]
    @test encrypt("aes256", hex2bytes(key), hex2bytes(text)) == hex2bytes(encrypted)
    @test decrypt("aes256", hex2bytes(key), hex2bytes(encrypted)) == hex2bytes(text)
end

# Test lower-level API
key = "this key's exactly 32 bytes long"
enc = Encryptor("AES256", key)
plaintext = "this is 16 chars"
ciphertext = encrypt(enc, Vector{UInt8}(plaintext))

dec = Decryptor("AES256", key)
deciphertext = decrypt(dec, ciphertext)
@test Vector{UInt8}(plaintext) == deciphertext # no bytestring

willcauseassertion = "this is 16 (∀).." # case of length(::String) == 16
@test length(willcauseassertion) == 16
@test sizeof(willcauseassertion) == 18
# @test_throws AssertionError Vector{UInt8}(willcauseassertion) == decrypt(dec, encrypt(enc, Vector{UInt8}(willcauseassertion))) # can not catch this c assertion
@test_throws ArgumentError Vector{UInt8}(willcauseassertion) == decrypt(dec, encrypt(enc, Vector{UInt8}(willcauseassertion)))
# @test_throws AssertionError Vector{UInt8}(willcauseassertion) == decrypt(dec, encrypt(enc, Vector{UInt8}(willcauseassertion))) # can not catch this c assertion
@test_throws ArgumentError Vector{UInt8}(willcauseassertion) == decrypt(dec, encrypt(enc, Vector{UInt8}(willcauseassertion)))

willbebroken = "this is 16 (∀)" # case of length(::String) != 16
@test length(willbebroken) == 14
@test sizeof(willbebroken) == 16
@test Vector{UInt8}(willbebroken) == decrypt(dec, encrypt(enc, Vector{UInt8}(willbebroken)))
@test willbebroken == String(decrypt(dec, encrypt(enc, Vector{UInt8}(willbebroken))))
@test Vector{UInt8}(willbebroken) == decrypt(dec, encrypt(enc, Vector{UInt8}(willbebroken)))
@test willbebroken == String(decrypt(dec, encrypt(enc, Vector{UInt8}(willbebroken))))

criticalbytes = hex2bytes("6e6f74555446382855aa552de2888029")
@test length(criticalbytes) == 16
@test sizeof(criticalbytes) == 16
@test criticalbytes == decrypt(dec, encrypt(enc, criticalbytes))

# This one will pass, but may be caught UnicodeError exception when evaluate it by julia ide.
dummy = String(decrypt(dec, encrypt(enc, criticalbytes)))
@test isa(dummy, AbstractString)
if !isdefined(Core, :String) || !isdefined(Core, :AbstractString)
    @test !isa(dummy, ASCIIString)
end
@test isa(dummy, String) # gray zone
@test sizeof(dummy) == 16
@test length(Vector{UInt8}(dummy)) == 16
@test length(dummy) != 16

# Test errors
@test_throws ArgumentError Encryptor("this is not a cipher", key)
@test_throws ArgumentError Decryptor("this is not a cipher either", key)

@test_throws ArgumentError Encryptor("AES256", "bad key length")
@test_throws ArgumentError Decryptor("AES256", "bad key length")

enc = Encryptor("AES256", key)
dec = Decryptor("AES256", key)
result = Vector{UInt8}(undef, sizeof(plaintext) - 1)
@test_throws ArgumentError encrypt!(enc, result, plaintext)
@test_throws ArgumentError decrypt!(dec, result, plaintext)

# Test show methods
println("Testing cipher show methods:")
println(get_cipher_types()["AES256"])
println(Encryptor("AES256", key))
println(Decryptor("AES256", key))

println("Testing cipher AES256CBC:")
# AES256CBC
for (pw,salt,iv,key,text,encrypted) in [
    (
        Vector(b"Secret Passphrase"),
        "a3e550e89e70996c",
        "7c7ed9434ddb9c2d1e1fcc38b4bf4667",
        "e299ff9d8e4831f07e5323913c53e5f0fec3a040a211d6562fa47607244d0051",
        "4d657373616765",
        "da8aab1b904205a7e49c1ecc7118a8f4",
    ),(
        Vector(b"Secret Passphrase"),
        "a3e550e89e70996c",
        "7c7ed9434ddb9c2d1e1fcc38b4bf4667",
        "e299ff9d8e4831f07e5323913c53e5f0fec3a040a211d6562fa47607244d0051",
        "4d657373616765090909090909090909",
        "da8aab1b904205a7e49c1ecc7118a8f4804bef7be79216196739de7845da182d",
    )
]
    (key32, iv16) = (hex2bytes(key), hex2bytes(iv))
    @test gen_key32_iv16(pw, hex2bytes(salt)) == (key32, iv16)
    @test encrypt("aes256", :CBC, iv16, key32, add_padding_PKCS5(hex2bytes(text), 16)) == hex2bytes(encrypted)
    @test trim_padding_PKCS5(decrypt("aes256", :CBC, iv16, key32, hex2bytes(encrypted))) == hex2bytes(text)
end

# Test errors
badkey = "this key's exactly 32(∪∩∀ДЯ)...."
@test length(badkey) == 32
@test sizeof(badkey) == 40
@test_throws ArgumentError Encryptor("AES256", badkey)
@test_throws ArgumentError Decryptor("AES256", badkey)

iv = Vector(b"this is 16 chars")
key = Vector(b"this key's exactly 32 bytes long")
enc = Encryptor("AES256", key)
dec = Decryptor("AES256", key)
shortresult = Vector{UInt8}(undef, sizeof(plaintext) - 1)
@test_throws ArgumentError encrypt!(enc, :CBC, iv, shortresult, plaintext)
@test_throws ArgumentError decrypt!(dec, :CBC, iv, shortresult, plaintext)
result = Vector{UInt8}(undef, sizeof(plaintext))
shorttext = Vector{UInt8}(undef, sizeof(plaintext) - 1)
@test_throws ArgumentError encrypt!(enc, :CBC, iv, result, shorttext)
@test_throws ArgumentError decrypt!(dec, :CBC, iv, result, shorttext)
longresult = Vector{UInt8}(undef, sizeof(plaintext) + 1)
longtext = Vector{UInt8}(undef, sizeof(plaintext) + 1)
@test_throws ArgumentError encrypt!(enc, :CBC, iv, longresult, longtext)
@test_throws ArgumentError decrypt!(dec, :CBC, iv, longresult, longtext)
shortiv = Vector{UInt8}(undef, sizeof(iv) - 1)
@test_throws ArgumentError encrypt!(enc, :CBC, shortiv, result, plaintext)
@test_throws ArgumentError decrypt!(dec, :CBC, shortiv, result, plaintext)
longiv = Vector{UInt8}(undef, sizeof(iv) + 1)
@test_throws ArgumentError encrypt!(enc, :CBC, longiv, result, plaintext)
@test_throws ArgumentError decrypt!(dec, :CBC, longiv, result, plaintext)

# encrypt!(enc, :GCM, iv, result, plaintext)
# decrypt!(dec, :CCM, iv, result, plaintext)
@test_throws ArgumentError encrypt!(enc, :UNKNOWN, iv, result, plaintext)
@test_throws ArgumentError decrypt!(dec, :UNKNOWN, iv, result, plaintext)

println("Cipher AES256CBC OK.")


include("../src/Dataset.jl")
import .Dataset:
  get_MINST,
  TrainTest,
  split_X_by_batches

include("../src/Model.jl")
import .Model:
  split_encoder_result,
  create_vae,
  model_sample,
  conv_MINST_model

include("../src/Utils.jl")
import .Utils:
  gen_images,
  save_model,
  load_model

using Test
using Flux
using Flux.Tracker: TrackedReal
using Printf
using Distributions: Uniform

temp_dir = tempdir()

@testset "Model save/load" begin
  n_sample = Int(floor(rand(Uniform(10,100))))
  n_latent = Int(floor(rand(Uniform(1,50)))) * 2

  ps, loss_fn, f, g = create_vae(n_sample, n_latent)
  file = tempname()
  save_model(f, g, file)

  fp, gp = load_model
  @test true
end

@testset "Image Utilities" begin
  n_sample = Int(floor(rand(Uniform(10,100))))
  n_latent = Int(floor(rand(Uniform(1,50)))) * 2
  dataset = get_MINST(n_sample)
  ps, loss_fn, f, g = create_vae(n_latent, n_sample)

  # zero initialized μ̂, logσ̂
  outfile_zero_latent = joinpath(temp_dir, "zero_latent.png")
  @test typeof(model_sample(f)) == BitArray{4}
  gen_images(outfile_zero_latent, f)
  @test isfile(outfile_zero_latent)
  @sprintf("saved file to %s", outfile_zero_latent)

  # f(g(x)) test
  outfile_X_latent = joinpath(temp_dir, "X_latent.png")
  gen_images(outfile_X_latent, g, f, dataset.test_x)
  @test isfile(outfile_X_latent)
  @sprintf("saved file to %s", outfile_X_latent)
end

@testset "ADAM optimization can run" begin
  n_sample = Int(floor(rand(Uniform(10,100))))
  n_latent = Int(floor(rand(Uniform(1,50)))) * 2
  dataset = get_MINST(n_sample)
  X = dataset.train_x
  ps, loss_fn, f, g = create_vae(n_latent, n_sample)
  opt = ADAM()

  @test typeof(loss_fn(X)) == TrackedReal{Float32}
  X = float.(X .> 0.5)
  Flux.train!(loss_fn, ps, zip([X]), opt)
  @test true == true
end

@testset "Convolution and transpose is isomorphic" begin
  dataset = get_MINST()

  n_sample = 100
  n_latent = 10
  X = dataset.train_x[:,:,:,1:n_sample]
  f, g = conv_MINST_model(n_latent)

  X_transformed = g(reshape(X, 28, 28, 1, n_sample))

  x_mean, x_std = split_encoder_result(X_transformed, n_latent)

  Xp = f(x_mean)
  @test size(Xp) == size(X)
end

@testset "MINST dataset: size/shape okay" begin
  ds = get_MINST()
  img_shape = (28, 28, 1)
  n_train = 60000
  n_test = 10000
  @test size(ds.train_x) == (img_shape..., n_train)
  @test size(ds.train_y) == (n_train, )
  @test size(ds.test_x) ==  (img_shape..., n_test)
  @test size(ds.test_y) == (n_test, )
end


using GPUifyLoops, Cthulhu, CuArrays, CUDAnative

function kernel!(::Val{nbcol}, A, B) where {nbcol}
    @inbounds @loop for i in (1:size(A,1);
                              (blockIdx().x-1)*blockDim().x + threadIdx().x)
        a = zero(eltype(B))
        for j = 1:nbcol
            a += B[i, j]
        end
        A[i] = a
    end
    nothing
end

@eval function kernel!(A::CuArray, B::CuArray)
  threads = 1024
  blocks = ceil(Int, size(A,1)/threads)
  @launch(CUDA(), threads=threads, blocks=blocks, kernel!(Val(size(B,2)), A, B))
end

b = rand(Float32, 10^8, 10)
a = similar(b, size(b,1))
cb = CuArray(a)
ca = similar(cb, size(b,1))

@static if isdefined(Base, :active_repl)
    # Example of how to decend using Cthulhu
    args = (Val(size(cb,2)), ca, cb)
    cuargs = map(cudaconvert, args)
    tt = map(typeof, cuargs)
    descend(kernel!, Tuple{tt...})
end

# Run the kernel
kernel!(ca, cb)
kernel!(ca, cb)
kernel!(ca, cb)
kernel!(ca, cb)

# Dump intermediate files into the tmp directory
@device_code dir="tmp" kernel!(ca, cb)


export pileup

using Base.Collections

type Pileuper{T}
    reads::T
    window::PriorityQueue{Read, Int32, Base.Order.ForwardOrdering} # Read -> ref pos of last matching base
    muts::PriorityQueue{Mut, Int32, Base.Order.ForwardOrdering}    # Mut  -> pos
    chr::Int32
    Pileuper(x::T) = new(x, PriorityQueue(Read, Int32), PriorityQueue(Mut, Int32), -2)
end

pileup{T}(x::T) = pileup(Pileuper{T}(x))

function pileup{T}(p::Pileuper{T})
    for r in p.reads
        if r.refID != p.chr
            flush_muts!(p)
            p.chr = r.refID
        end
        add_muts!(p, r)
    end
    flush_muts!(p)
end

function produce_muts!{T}(p::Pileuper{T})
    mut = dequeue!(p.muts)

    while !isempty(p.window) && peek(p.window).second < mut.pos
        dequeue!(p.window)
    end

    produce((keys(p.window), p.chr, mut))
end

function flush_muts!{T}(p::Pileuper{T})
    while !isempty(p.muts)
        produce_muts!(p)
    end

    # clear p.window by hack; not sure if this is faster than just allocating a new one
    empty!(p.window.xs)
    empty!(p.window.index)
end

function add_muts!{T}(p::Pileuper{T}, r::Read)
    while !isempty(p.muts) && peek(p.muts).second < r.pos
        produce_muts!(p)
    end

    enqueue!(p.window, r, r.pos + calc_distance(r) - 1)

    for mut in r.muts
        mut = map_to_ref(mut, r)
        haskey(p.muts, mut) || enqueue!(p.muts, mut, mut.pos)
    end
end


export buildnet

function buildnet(;name="temp",model=NetworkModel(),method="Euler1",duration=0.1)
  iteration=round(Int64,duration/model.dt)
  syncIter=model.synctime>0?round(Int64,model.synctime÷model.dt):1
  iteration=iteration÷syncIter

  # open a file
  fid=open(name*".jl","w")

  str="""
  println("\\nExecuting $(name).jl")

  using JLD

  if $(model.nCores)>0 && (addCores=$(model.nCores)-nworkers())>0
    nworkers()==1?addprocs(addCores):addprocs(addCores-1)
  end
  println("Using \$(nworkers()) workers")
  """
  println(fid,str)

  # units
  unit="""
  metre=1e2;meter=1e2;cm=metre/1e2;mm=metre/1e3;um=metre/1e6;nm=metre/1e9;
  second=1e3;ms=second/1e3;
  Hz=1/second;
  voltage=1e3;mV=voltage/1e3;
  ampere=1e6;mA=ampere/1e3;uA=ampere/1e6;nA=ampere/1e9;
  farad=1e6;uF=ufarad=farad/1e6;
  siemens=1e3;mS=msiemens=siemens/1e3;nS=nsiemens=siemens/1e9"""
  println(fid,"@everywhere begin\n",unit,"\nend\n")

  # parameters
  println(fid,"@everywhere begin\n",model.parameters,"\nend\n")

  # equations
  eqs=map(strip,split(replace(model.equations,";","\n"),"\n"));
  vars=[]
  diffEqs=[]

  for k=1:length(eqs)
    if eqs[k][1]=='d'
      vars=[vars;match(r"(?<=d).*(?=/dt)",eqs[k]).match]
      diffEq=vectorize(match(r"(?<==).*$",eqs[k]).match)
      diffEqs=[diffEqs;diffEq]
    else
      funcName=match(r"^.*(?==)",eqs[k]).match
      funcExpr=vectorize(match(r"(?<==).*$",eqs[k]).match)
      diffEqs=map(diffEq->replace(diffEq,funcName,funcExpr),diffEqs)
    end
  end

  varsType=map(var->var*"::Array{Float64}",vars)
  argumentType=join([varsType;"t::Float64"],",");
  println(fid,"# function")
  for k=1:length(vars)
    println(fid,"@everywhere F",vars[k],"(",argumentType,")=",diffEqs[k])
  end
  println(fid)

  # function
  println(fid,"function $name(iteration::Int64=10)\n")

  # Memory use before allocation
  str="""
  if Sys.OS_NAME == :Linux
    println("Memory use before allocation:")
    run(pipeline(`free -og`, `grep -v Swap`))
  elseif Sys.OS_NAME == :Windows
    # println("Memory use before allocation:")
    # run(`wmic OS get FreePhysicalMemory /Value`)
  end
  """
  println(fid,str)

  println(fid,"t=0.0")
  println(fid,"dt=",model.dt)
  println(fid,"N=",model.groups[:N])
  println(fid,"syncIter=",syncIter,"\n")

  # initialize variables
  println(fid,"# initialize variables")
  for var in vars
    println(fid,var,"=SharedArray(Float64,N)")
  end
  println(fid,model.initialization,"\n")

  # connection
  println(fid,"# initialize connections\n",model.connection,"\n")

  # report
  report="""# report
  if k%(1000÷syncIter)==0
     @printf("%s%4.2f finished\\n","%",100*k/iteration)
  end
  """

  # integrate

  argument=join([map(var->var*"[localindex]",vars);"t"],",")
  splitArg=split(argument,",")
  argumentTemp=join([splitArg[1]*"+Δ$(vars[1])";splitArg[2:end]],",")

  func=map(var->"F$var($argument)",vars)
  tempFunc="F$(vars[1])($argumentTemp)"

  if model.solver=="Euler1"
    integrate="Δ$(vars[1])=dt*$(func[1])\n"
  elseif model.solver=="Euler2"
    integrate="Δ$(vars[1])=dt/2.0*$(func[1])\n"
    integrate*="Δ$(vars[1])+=dt/2.0*$tempFunc\n"
  elseif model.solver=="Euler4"
    integrate="Δ$(vars[1])=dt/4.0*$(func[1])\n"
    integrate*="Δ$(vars[1])+=dt/4.0*$tempFunc\n"
    integrate*="Δ$(vars[1])+=dt/4.0*$tempFunc\n"
    integrate*="Δ$(vars[1])+=dt/4.0*$tempFunc\n"
  end

  for k=2:length(vars)
    integrate*="Δ$(vars[k])=dt*$(func[k])\n"
  end

  for k=1:length(vars)
	integrate*="$(vars[k])[localindex]+=Δ$(vars[k])\n"
  end
  integrate="""# integrate
  @sync for p in workers()
  @spawnat p begin
  for idx=1:syncIter
  localindex=localindexes($(vars[1]))
  $integrate
  end #end of sync
  end #end of begin
  end #end of workers
  """

  for var in vars
  reg= Regex("(?<=\\b)($var)(?=\\b)")
  diffEqs=map(diffEq->replace(diffEq,reg,s"\1[localindex]"),diffEqs)
  end

  # spike
  if !isempty(model.spike)
    spike=split(replace(model.spike,";","\n"),"\n")
    threshold=spike[1]
    length(spike)==2? refractory=spike[2]:refractory=""

    println(fid,"# initialize tSpike")
    println(fid,"tSpike=zeros(round(Int64,N*iteration*dt))")
    println(fid,"ts=zeros(N)")
    println(fid,"spikeNeuron=zeros(Int64,round(Int64,N*iteration*dt))")
    println(fid,"state=1","\n")

    spike="# spike\n"
    spike*="isExceedThreshold=$threshold\n"
    if !isempty(refractory)
      spike*="isNotRefractory=$refractory\n"
      spike*="isSpike=isExceedThreshold & isNotRefractory\n"
    else
      spike*="isSpike=isExceedThreshold\n"
    end
    spike*="spikeIndex=find(isSpike)\n"
    spike*="nSpike=length(spikeIndex)\n"
    spike*="tSpike[state:state+nSpike-1]=t\n"
    spike*="ts[spikeIndex]=t\n"
    spike*="spikeNeuron[state:state+nSpike-1]=spikeIndex\n"
    spike*="state+=nSpike\n"
  else spike=""
  end

  # reset
  if !isempty(model.reset)
    reset="# reset\n"
    for line in split(replace(model.reset,";","\n"),"\n")
      var=match(r"^.*(?==)",line).match
      expr=match(r"(?<==).*",line).match
      reset*="$(var)[spikeIndex]=$(expr)\n"
    end
  else reset=""
  end

  # synapse
  if !isempty(model.synapse)
    synapse="# synapse\n"
    for line in split(replace(model.synapse,";","\n"),"\n")
      var1=match(r"^.*(?=\+)",line).match
      var2=match(r"(?<==).*",line).match
      if model.nCores>0
      synapse*="$(var1)[:]+=sum($(var2)[:,spikeIndex],2)\n"
      else synapse*="$(var1)+=sum($(var2)[:,spikeIndex],2)\n"
      end
    end
  else synapse=""
  end

  # plasticity
  if !isempty(model.plasticity)
    plasticity="# plasticity\n"
    plasticity*=replace(replace(model.plasticity,"j",":"),"i","spikeIndex")*"\n"
  else plasticity=""
  end

  # record
  if !isempty(model.record)
    println(fid,"# initialize Record")
    for k=1:length(vars)
      println(fid,"$(vars[k])Record=zeros($(length(model.record)),iteration)")
    end
    println(fid)

    record="# record\n"
    for k=1:length(vars)
      record*="$(vars[k])Record[:,k]=$(vars[k])[$(model.record)]\n"
    end
  else
    record=""
  end

  # Memory use after allocation
  str="""
  if Sys.OS_NAME == :Linux
    println("Memory use after allocation:")
    run(pipeline(`free -og`, `grep -v Swap`))
  elseif Sys.OS_NAME == :Windows
    # println("Memory use before allocation:")
    # run(`wmic OS get FreePhysicalMemory /Value`)
  end
  """
  println(fid,str)
  # print status
  println(fid,"# print status")
  str="""
  if iteration>=50
      println("simulation starts...")
  else
    println("compiling...")
  end"""
  println(fid,str,"\n")

  # for-loop
  println(fid,"tic()")
  println(fid,"@fastmath @inbounds for k=1:iteration\n")
  println(fid,report)
  println(fid,integrate)
  println(fid,spike)
  println(fid,reset)
  println(fid,synapse)
  println(fid,plasticity)
  println(fid,record)
  println(fid,"t+=syncIter*dt")
  println(fid,"end # end of for loop ")
  println(fid,"toc()","\n")

  # save variables
  str="""# save variables
  try jldopen("$name.jld","w") do fid
    try  write(fid,"spikeNeuron",spikeNeuron[1:state-1])  end
    try  write(fid,"tSpike",tSpike[1:state-1])            end
    try  write(fid,"$(vars[1])Record",$(vars[1])Record)   end
    """
  for var in model.saveVars
  str*="try write(fid,\"$(var)\",$var) end\n"
  end
  str*="end #end of do\n"
  str*="end #end of try jldopen\n"
  println(fid,str)
    # end of function
  println(fid,"end # end of function\n")

  # precompile and run
  println(fid,"# precompile and run")
  println(fid,"precompile($name,($iteration,))")
  println(fid,"$name(20)")
  println(fid,"$name(20)")
  println(fid,"$name($iteration)")

  # post-processing
  println(fid,"\n",model.postProcs)

  # close file, open file
  close(fid)
  try run(`atom $(abspath("$name.jl"))`) end
  # include("$name.jl")
end


###################################################
## network/iterative/state.jl
## Different iterative procedures
###################################################

"""
    `IterativeState`
    Abstract type, represents all information of a state of path/timingscomputation
    in the iterative process
    Must contain:
    - `data::NetworkData`           timing data in the network
    - `timings::NetworkTimings`     the current timings solution
    - `trips::Vector{NetworkTrip}`  trips that we currently optimize on
    - `paths::Vector{Vector{Dict{Edge, Float64}}}` path subset for each current trip datapoint. Each trip has a vector of paths associated to it. Each path is represented as a dictionary mapping edges to weights. This ensures that the path representation is consistent for all types of methods (including graph reduction). In most cases, the weights are all ones and the representation is equivalent to a vector of edges.
    - `pathDiff::Float64`           measure of difference in paths from previous step. Full definition is in paper. The lower it is, the closer we are to convergence.
    Must implement:
    - `updateState!` updates the state when new link-times are computed
"""
abstract type IterativeState end

function Base.show(io::IO, s::IterativeState)
    typeName = split(string(typeof(s)),".")[end]
    nTrips = length(s.trips)
    nPaths = sum([length(paths) for paths in s.paths])
    println(io,"Iterative heuristic: $(typeName)")
    println(io,"Optimizing on $nTrips trips and $nPaths paths")
end

NetworkTimings(it::IterativeState) = it.timings

"""
    `doIteration!`, shortcut to update state with optimizer's output
"""
function doIteration!(it::IterativeState; method::AbstractString="lp", continuityParam::Float64 = 0.1, solverArgs...)
    if method=="lp"
        updateState!(it, lpTimes(it; solverArgs...))
    elseif method=="fraclp"
        updateState!(it, fraclpTimes(it; solverArgs...))
    elseif method=="mip"
        updateState!(it, mipTimes(it; solverArgs...))
    elseif method == "heuristic"
        updateState!(it, heuristicTimes(it))
    elseif method == "socp"
        updateState!(it, socpTimes(it; solverArgs...))
    elseif method == "roadtype"
        updateState!(it, constantSpeedTimes(it; solverArgs...), updatePaths=false)
    elseif method == "constant"
        updateState!(it, constantSpeedTimes(it; uniqueSpeed=true, solverArgs...),
                     updatePaths=false)
    elseif method == "lpCo"
        updateState!(it, lpTimes(it, continuityConstraint="simple",
                                     velocityBound=continuityParam; solverArgs...))
    elseif method == "lpCoNbhd"
        updateState!(it, lpTimes(it, continuityConstraint="neighborhoods",
                                     velocityBound=continuityParam; solverArgs...))
    elseif method == "socpCo"
        updateState!(it, socpTimes(it, continuityConstraint="simple",
                                       continuityParam=continuityParam; solverArgs...))
    elseif method == "socpCoNbhd"
        updateState!(it, socpTimes(it, continuityConstraint="neighborhoods",
                                       continuityParam=continuityParam; solverArgs...))
    elseif method == "socpCoReg"
        updateState!(it, socpTimes(it, continuityConstraint="regularized",
                                       continuityParam=continuityParam; solverArgs...))
    elseif method == "pathchoice"
        updateState!(it, pathChoice(it), updatePaths=false)
    else
        error("Unknown optimizer")
    end
    return it
end


# Return a random DNA/RNA sequence of the given length.
function random_seq(n::Integer, nts, probs, outtype = String)
    cumprobs = cumsum(probs)
    x = Vector{Char}(undef, n)
    for i in 1:n
        x[i] = nts[searchsorted(cumprobs, rand()).start]
    end
    return outtype(x)
end

function random_seq(::Type{A}, n::Integer) where {A<:Alphabet}
    # TODO: Resolve the use of symbols(A()).
    nts = symbols(A())
    probs = Vector{Float64}(undef, length(nts))
    fill!(probs, 1 / length(nts))
    return LongSequence{A}(random_seq(n, nts, probs))
end

function random_dna(n, probs=[0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'T', 'N'], probs)
end

function random_rna(n, probs=[0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'U', 'N'], probs)
end

function random_aa(len)
    return random_seq(len,
        ['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I',
         'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X' ],
        push!(fill(0.049, 20), 0.02))
end

function random_dna_symbols(n, probs=[0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'T', 'N'], probs, Vector{DNA})
end

function random_rna_symbols(n, probs=[0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'U', 'N'], probs, Vector{RNA})
end

function random_rna_symbols(n, probs=[0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'C', 'G', 'U', 'N'], probs, Vector{RNA})
end

function random_aa_symbols(n, probs=[0.24, 0.24, 0.24, 0.24, 0.04])
    return random_seq(n, ['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I',
     'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X' ], probs, Vector{AminoAcid})
end

function random_dna_kmer(len)
    return random_dna(len, [0.25, 0.25, 0.25, 0.25])
end

function random_rna_kmer(len)
    return random_rna(len, [0.25, 0.25, 0.25, 0.25])
end

function dna_complement(seq::AbstractString)
    seqc = Vector{Char}(undef, length(seq))
    complementer = Dict(zip("-ACGTSWYRKMDVHBN", "-TGCASWRYMKHBDVN"))
    for (i, c) in enumerate(seq)
        seqc[i] = complementer[c]
    end
    return String(seqc)
end

function rna_complement(seq::AbstractString)
    seqc = Vector{Char}(undef, length(seq))
    complementer = Dict(zip("-ACGUSWYRKMDVHBN", "-UGCASWRYMKHBDVN"))
    for (i, c) in enumerate(seq)
        seqc[i] = complementer[c]
    end
    return String(seqc)
end


include("Include.jl")

# load the data dictionary -
data_dictionary = DataDictionary(0,0,0)
TXTL_dictionary = TXTLDictionary(0,0,0)
number_of_fluxes = length(data_dictionary["default_flux_bounds_array"][:,1])
number_of_species = length(data_dictionary["species_bounds_array"][:,1])
#Set objective reaction
data_dictionary["objective_coefficient_array"][171] = -1;
#-------------------------------------------------------------------------------------------------#Define case number
# 1 = Amino Acid Uptake & Synthesis
# 2 = Amino Acid Uptake w/o Synthesis
# 3 = Amino Acid Synthesis w/o Uptake
case = 1
if case == 1
  data_dictionary["AASyn"] = 100;
  data_dictionary["AAUptake"] = 30
  data_dictionary["AASecretion"] = 0;
end
if case == 2
  data_dictionary["AASyn"] = 0;
  data_dictionary["AAUptake"] = 30
  data_dictionary["AASecretion"] = 0;
end
if case == 3
  data_dictionary["AASyn"] = 100;
  data_dictionary["AAUptake"] = 0
  data_dictionary["AASecretion"] = 100;
end
#-------------------------------------------------------------------------------------------------
volume = TXTL_dictionary["volume"]
plasmid_concentration = 5;
gene_copy_number = (volume/1e9)*(6.02e23)*plasmid_concentration;
TXTL_dictionary["gene_copies"] = gene_copy_number
data_dictionary = DataDictionary(0,0,0)
#Set objective reaction
data_dictionary["objective_coefficient_array"][194] = -1;
data_dictionary["AASyn"] = 100;
data_dictionary["AAUptake"] = 30
data_dictionary["AASecretion"] = 0;
data_dictionary["GlcUptake"] = 30#*rand(1)[1];
data_dictionary["Oxygen"] = 30#+2*rand(1)[1];
TXTL_dictionary["RNAP_concentration_nM"] = 75 #+ (80-60)*rand(1)[1];
TXTL_dictionary["RNAP_elongation_rate"] = 25 #+ (30-20)*rand(1)[1];
TXTL_dictionary["RIBOSOME_concentration"] = 0.0016# + (0.0018-0.0012)*rand(1)[1];
TXTL_dictionary["RIBOSOME_elongation_rate"] = 2# + (3-1.5)*rand(1)[1];
data_dictionary["Oxygen"] = 100;
data_dictionary["AAUptake"] = 1 #+ (1.5-0.5)*rand(1)[1];
data_dictionary["GlcUptake"] = 30# + (38-28)*rand(1)[1];
data_dictionary["AC"] = 8# + (12-6)*rand(1)[1];
data_dictionary["SUCC"] = 2# + (3-1)*rand(1)[1];
data_dictionary["PYR"] = 7# + (10-5)*rand(1)[1];
data_dictionary["MAL"] = 3 #+ (5-2)*rand(1)[1];
data_dictionary["LAC"] = 8 #+ (12-8)*rand(1)[1];
data_dictionary["ENERGY"] = 1#1*rand(1)[1];
data_dictionary["ALA"] = 2 #+ (4-3)*rand(1)[1];
data_dictionary["ASP"] = 1#+rand(1)[1];
data_dictionary["GLN"] = 0.25#*rand(1)[1];
data_dictionary["LysUptake"] = 1.5#*rand(1)[1];
data_dictionary["LysSecretion"] = 0.5#1*rand(1)[1];
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------------------------------------
t0_exp = 0
tf_exp = 3
tstep_exp = 0.01
experimental_time_exp = [t0_exp:tstep_exp:tf_exp;]

species_list = readdlm("./data/species_list")
Mean_raw = readdlm("./data/Mean")
Std_raw = readdlm("./data/Std")
Upper_raw = readdlm("./data/Upper")
Lower_raw = readdlm("./data/Lower")
Mean = Dict()
Std = Dict()
Upper = Dict()
Lower = Dict()
for i in 1:length(species_list)
	rxn_idx = species_list[i]
	Mean[rxn_idx] = Mean_raw[:,i]
	Std[rxn_idx] = Std_raw[:,i]
	Upper[rxn_idx] = Upper_raw[:,i]
	Lower[rxn_idx] = Lower_raw[:,i]
end

data_dictionary["species_drift_amplitude"] = 2 # percent allowed to drift for unconstraint species for each iteration (0.01 hr)
data_dictionary["species_drift_decay_rate"] = 1
metabolite_list = data_dictionary["list_of_metabolite_symbols"]
rxn_list = data_dictionary["list_of_reaction_strings"]

(n_species,n_rxn) = size(data_dictionary["stoichiometric_matrix"])
t0 = 0
tf = 3
tstep = 0.01
experimental_time = t0:tstep:tf
length_time = length(experimental_time)
data_dictionary["tstep"] = tstep

tRNA_index = Int64[]
for i in 1:length(metabolite_list)
	rxn_idx = metabolite_list[i]
	if length(rxn_idx) > 5
		if rxn_idx[end-4:end] == "_tRNA"
			push!(tRNA_index,i)
		end
	end
end

initial_conditions = zeros(n_species,1)
for species_index in 1:n_species
	if in(species_index,tRNA_index)
		initial_conditions[species_index] = Mean[metabolite_list[species_index]][2]
	else
		initial_conditions[species_index] = Mean[metabolite_list[species_index]][1]
	end
end
#initial_conditions[142] = .001
initial_conditions[143] = .001
initial_conditions[144] = .0001
data_dictionary["initial_conditions"] = initial_conditions

# this determines which parameters are fitted to data
#dataset_index = [15; 18; 22; 24; 27; 29; 31; 34; 38; 41; 42; 59; 60; 61; 66; 69; 70; 76; 79; 83; 84; 86; 88; 89; 106; 110; 114; 121; 123; 126; 128; 130; 132; 133; 134; 135; 140]
dataset_index = [15; 18; 22; 24; 27; 29; 31; 34; 38; 41; 42; 59; 60; 61; 66; 69; 70; 76; 79; 83; 84; 86; 88; 89; 106; 110; 114; 121; 123; 126; 128; 130; 132; 133; 134; 135]
species_constraint_index = copy(dataset_index)

species_list = data_dictionary["list_of_metabolite_symbols"]

syn_data_upper = zeros(n_species,length_time)
syn_data_lower = zeros(n_species,length_time)
syn_mean = zeros(n_species,length_time)
for syn_index in species_constraint_index
	itp = interpolate((experimental_time_exp,),Mean[metabolite_list[syn_index]],Gridded(Linear()))
	Mean_itp = itp[experimental_time]
	itp = interpolate((experimental_time_exp,),Std[metabolite_list[syn_index]],Gridded(Linear()))
	Std_itp = itp[experimental_time]
	itp = interpolate((experimental_time_exp,),Lower[metabolite_list[syn_index]],Gridded(Linear()))
	Lower_itp = itp[experimental_time]
	itp = interpolate((experimental_time_exp,),Upper[metabolite_list[syn_index]],Gridded(Linear()))
	Upper_itp = itp[experimental_time]
	syn_data_upper[syn_index,:] = Upper_itp
	syn_data_lower[syn_index,:] = Lower_itp
	syn_mean[syn_index,:] = Mean_itp
end
syn_data = Dict{AbstractString,Any}()
syn_data["mean"] = syn_mean
syn_data["upper"] = syn_data_upper
syn_data["lower"] = syn_data_lower

#--------------------------------------------------------------------------------------------------
#--------------------------------------------------------------------------------------------------
#--------------------------------------------------------------------------------------------------
#--------------------------------------------------------------------------------------------------
#--------------------------------------------------------------------------------------------------
if !isdir("TXTL")
	mkdir("TXTL")
end

if !isfile("TXTL/num_dir")
	num_dir = 0
else
	num_dir = Int64(readdlm("TXTL/num_dir")[1])
end

# Original values:
RNAP_IC = 1e-3
RIBO_IC = 0.002
RNAP_elongation_rate = 25
plasmid_saturation_coefficient = 116e-6
polysome_amplification = 5
RIBOSOME_elongation_rate = 2
mRNA_saturation_coefficient = 0.045
mRNA_degradation_rate = 10

params = vcat(RNAP_IC, RIBO_IC, RNAP_elongation_rate, plasmid_saturation_coefficient, polysome_amplification, RIBOSOME_elongation_rate, mRNA_saturation_coefficient, mRNA_degradation_rate)

param_bounds = [8e-4 12e-4
				0.0015 0.003
				20 30
				116e-7 116e-5
				5 15
				1.5 3
				0.0045 0.45
				8 12]

data_dictionary["initial_conditions"][143] = params[1]
data_dictionary["initial_conditions"][141] = params[2]
TXTL_dictionary["RNAP_elongation_rate"] = params[3]
TXTL_dictionary["plasmid_saturation_coefficient"] = params[4]
TXTL_dictionary["polysome_amplification"] = params[5]
TXTL_dictionary["RIBOSOME_elongation_rate"] = params[6]
TXTL_dictionary["mRNA_saturation_coefficient"] = params[7]
TXTL_dictionary["mRNA_degradation_rate"] = params[8]

alpha = 25
#variance_min = .1
#variance_max = 10
failed_tps_threshold = 0
CAT_threshold = 5/1000 # base-case is 0.02326

for i in 1:1e10

	params_new = copy(params)
	for p in 1:length(params_new)
		if in(p,[4;7])
			params_new[p] = exp(log(param_bounds[p,1])+rand().*(log(param_bounds[p,2])-log(param_bounds[p,1])))
		else
			params_new[p] = param_bounds[p,1]+rand().*(param_bounds[p,2]-param_bounds[p,1])
		end
	end

	data_dictionary_new = deepcopy(data_dictionary)
	TXTL_dictionary_new = deepcopy(TXTL_dictionary)
	data_dictionary_new["initial_conditions"][143] = params_new[1]
	data_dictionary_new["initial_conditions"][141] = params_new[2]
	TXTL_dictionary_new["RNAP_elongation_rate"] = params_new[3]
	TXTL_dictionary_new["plasmid_saturation_coefficient"] = params_new[4]
	TXTL_dictionary_new["polysome_amplification"] = params_new[5]
	TXTL_dictionary_new["RIBOSOME_elongation_rate"] = params_new[6]
	TXTL_dictionary_new["mRNA_saturation_coefficient"] = params_new[7]
	TXTL_dictionary_new["mRNA_degradation_rate"] = params_new[8]

	avg_perturbation = exp(mean(abs(log(params_new./params))))
	avg_perturbation_round = round(avg_perturbation,3)

	time_state_array_new,error_new,num_failed_tps = RunFBA(data_dictionary_new, TXTL_dictionary_new)
	cost_new = sum(error_new)

	CAT_end = time_state_array_new[140,end]
	CAT_end_round = round(CAT_end,5)

	if (num_failed_tps <= failed_tps_threshold) && (CAT_end > CAT_threshold)
		num_dir += 1
		mkdir("TXTL/$num_dir")
		writedlm("TXTL/$num_dir/params",params_new)
		writedlm("TXTL/$num_dir/time_state_array",time_state_array_new)
		writedlm("TXTL/$num_dir/num_failed_tps",num_failed_tps)
		writedlm("TXTL/num_dir",num_dir)
		println("num_failed_tps = $num_failed_tps, CAT_end = $CAT_end_round, ACCEPTED")
	else
		println("num_failed_tps = $num_failed_tps, CAT_end = $CAT_end_round")
	end
end


type LaserBounds{P<:LaserTagPOMDP}
    ubp::QMDPPolicy{P, Int}
    lbp::RandomPolicy
    max_depth::Int
    lb_not_same::Float64

    LaserBounds{P}() where P = new()
end

#=
function LaserBounds(p::LaserTagPOMDP)
    sol = QMDPSolver(max_iterations=1000)
    pol = solve(sol, p)
    lb_not_same = -p.step_cost/(1-discount(p))
    return LaserBounds{typeof(p)}(pol, lb)
end
=#

function bounds(l::LaserBounds, p::LaserTagPOMDP, b::ScenarioBelief)
    ps = particles(b)
    if !all(isterminal(p, s) for s in ps) && !isnull(previous_obs(b))
        if get(previous_obs(b)) == LaserTag.D_SAME_LOC
            @assert all(!isterminal(p, s) && s.robot==s.opponent for s in ps)
        else
            @assert any(isterminal(p, s) || s.robot!=s.opponent for s in ps)
        end
    end
    steps = l.max_depth-b.depth
    lb = ARDESPOT.branching_sim(p, l.lbp, b, steps)/length(b.scenarios)
    return bounds(l, p, ps, lb)
end

function bounds(l::LaserBounds, p::LaserTagPOMDP, particles, provided_lb=nothing)
    if all(isterminal(p, s) for s in particles)
        lb = 0.0
    elseif all(!isterminal(p, s) && s.robot==s.opponent for s in particles)
        lb = l.ubp.pomdp.tag_reward
    elseif !isa(provided_lb, Void)
        lb = provided_lb
    elseif length(particles) == 1
        lb = state_value(l.ubp, first(particles))
    else
        lb = l.lb_not_same
    end

    vsum = 0.0
    for s in particles
        vsum += state_value(l.ubp, s)
    end
    ub = vsum/length(particles)
    # if lb > ub
    #     warn("lb > ub")
    #     @show particles
    # end
    return lb, vsum/length(particles)
end

function init_bounds(l::LaserBounds, p::LaserTagPOMDP, max_depth::Int, rng=Base.GLOBAL_RNG)
    sol = QMDPSolver(max_iterations=1000)
    l.ubp = solve(sol, p)
    l.lbp = solve(RandomSolver(rng), p)
    l.max_depth = max_depth
    l.lb_not_same = -p.step_cost/(1-discount(p))
    return l
end

init_bounds(l::LaserBounds, p::LaserTagPOMDP, s::DESPOTSolver) = init_bounds(l, p, s.D, s.rng)
# DESPOT.init_bounds(l::LaserBounds, p::LaserTagPOMDP, c::DESPOT.DESPOTConfig) = init_bounds(l, p, c.search_depth)
#
# function DESPOT.bounds{S}(l::LaserBounds, p::LaserTagPOMDP, b::Vector{DESPOT.DESPOTParticle{S}}, ::DESPOT.DESPOTConfig)
#     return bounds(l, p, p.state for p in b)
# end
#
# DESPOT.default_action(l::LaserBounds, pomdp::POMDP, particles, c) = default_action(particles)

function default_action(particles)
    if all(s.robot==s.opponent for s in particles)
        return LaserTag.TAG_ACTION
    else
        warn("non-tag default")
        return rand(1:4)
    end
end

struct NoGapTag end

function ARDESPOT.default_action(ngt::NoGapTag, b, ex)
    if isa(ex, NoGap)
        return default_action(particles(b))
    else
        rethrow(ex)
    end
    return -1
end

function nogap_tag(b, ex)
    if isa(ex, NoGap)
        @assert all(!s.terminal && s.robot==s.opponent for s in iterator(b))
        # @assert ng.value == l.ubp.pomdp.tag_reward
        return LaserTag.TAG_ACTION
    else
        rethrow(ex)
    end
end

immutable InevitableInit end

POMCPOW.init_N(m::InevitableInit, pomdp::LaserTagPOMDP, h::BeliefNode, a::Int) = h.node == 1 ? 0 : 1

function POMCPOW.init_V(m::InevitableInit, pomdp::LaserTagPOMDP, h::BeliefNode, a::Int)
    # only works for POMCPOW now because of node access
    if a == LaserTag.TAG_ACTION
        if h.node == 1
            return 0.0
        elseif h.tree.o_labels[h.node] == LaserTag.C_SAME_LOC
            return pomdp.tag_reward
        else
            return -pomdp.tag_reward - pomdp.step_cost
        end
    else
        return -pomdp.step_cost
    end
end

struct RootToNextMLFirst
    rng::MersenneTwister
end

function MCTS.next_action(gen::RootToNextMLFirst, p::VDPTagPOMDP, b, node)
    if isroot(node) && n_children(node) < 1
        target_sum=MVector(0.0, 0.0)
        agent_sum=MVector(0.0, 0.0)
        for s in particles(b::ParticleCollection)
            target_sum += s.target
            agent_sum += s.agent
        end
        next = VDPTag.next_ml_target(mdp(p), target_sum/n_particles(b))
        diff = next-agent_sum/n_particles(b)
        return TagAction(false, atan2(diff[2], diff[1]))
    else
        return rand(gen.rng, actions(p))
    end
end

MCTS.next_action(gen::RootToNextMLFirst, p::Union{GenerativeBeliefMDP,MeanRewardBeliefMDP}, b, node) = next_action(gen, p.pomdp, b, node)

# XXXD
# function MCTS.next_action(gen::RootToNextMLFirst, p::VDPTag.DiscreteVDPTagProblem, b, node)
#     if isroot(node) && n_children(node) < 1
#         cpc = ParticleCollection([convert_s(state_type(cproblem(p)), s, p) for s in b.particles])
#         ca = next_action(gen, cproblem(p), cpc, node)
#         return convert_a(action_type(p), ca, p)
#     else
#         return rand(gen.rng, actions(p))
#     end
# end

struct ModeMDP <: Policy
    vi::ValueIterationPolicy
end

struct ModeMDPSolver <: Solver
    vi::ValueIterationSolver
end

ModeMDPSolver(;kwargs...) = ModeMDPSolver(ValueIterationSolver(;kwargs...))

POMDPs.solve(sol::ModeMDPSolver, pomdp::POMDP) = ModeMDP(solve(sol.vi, pomdp))

POMDPs.action(p::ModeMDP, b) = action(p.vi, mode(b))

struct VDPUpper end

function ARDESPOT.ubound(ub::VDPUpper, pomdp::POMDP, b::ScenarioBelief)
    if all(isterminal(pomdp, s) for s in particles(b))
        return 0.0
    else
        return mdp(cproblem(pomdp)).tag_reward
    end
end

struct SampleRollout{P<:Policy, RNG<:AbstractRNG}
    policy::P
    rng::RNG
end

function MCTS.estimate_value(ro::SampleRollout, mdp::Union{GenerativeBeliefMDP, MeanRewardBeliefMDP}, b, steps::Int)
    start_state = rand(ro.rng, b)
    sim = POMDPToolbox.RolloutSimulator(ro.rng,
                                        Nullable{Any}(start_state),
                                        Nullable{Float64}(),
                                        Nullable{Int}(steps))
    return POMDPToolbox.simulate(sim, mdp.pomdp, ro.policy, start_state)
end


#****************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#****************************************************************************

function genbo3(foldername, filename, u, xij, xik, qi, qj, qk, ti, tj, tk, ai, aj, ak, rho, mu, eta, kappa, natomtype)

# n,m = size(u);
# if (u[1]==0)
#     m = 0;
# else
#     n,m = size(u);
#     if (n==0) || (n == 1)
#         m = 0;
#     elseif (n == natomtype+2)
#         if u[n,1] == 0
#             m = 0;
#         end
#     else
#         error("Number of functions for the two-body bond order potential must be equal to 1 plus the number of atom types");
#     end
# end

if length(u)>0
    if (u[1]==0)
        m = 0;
        n = 0;
    else
        n,m = size(u);
    end
else
    n = 0;
    m = 0;
end

gen = 1;
ifile = 0;
stropu = "";
strcpu = "";
strgpu = "";
potnum = 0;
fp = "TripletcPair";
fn = "TripletcDensity";

for i = 1:m
    ui = u[:,i];
    if ui[n] != 0
        potnum = potnum + 1;
        stropu0, strcpu0, strgpu0 = gentriplet(foldername, filename * string(potnum), ui[1:(n-2)], xij, xik, qi, qj, qk, ti, tj, tk, ai, aj, ak, mu, eta, kappa, gen, ifile);
        stropu0 = modifystr3(stropu0, natomtype);
        strcpu0 = modifystr3(strcpu0, natomtype);
        strgpu0 = modifystr3(strgpu0, natomtype);
        stropu1, strcpu1, strgpu1 = genpair(foldername, fp * string(potnum), ui[n-1], xij, qi, qj, ti, tj, ai, aj, mu, eta, kappa, gen, ifile);
        stropu2, strcpu2, strgpu2 = gendensity(foldername, fn * string(potnum), ui[end], rho, mu, eta, kappa, gen, ifile);
        stropu = stropu * stropu0 * "\n" * stropu1 * "\n" * stropu2 * "\n";
        strcpu = strcpu * strcpu0 * "\n" * strcpu1 * "\n" * strcpu2 * "\n";
        strgpu = strgpu * strgpu0 * "\n" * strgpu1 * "\n" * strgpu2 * "\n";
    end
end

if potnum==0
    ifile = 0;
    gen = 0;
    stropu, strcpu, strgpu = gentriplet(foldername, filename, u, xij, xik, qi, qj, qk, ti, tj, tk, ai, aj, ak, mu, eta, kappa, gen, ifile);
    tmqopu, tmqcpu, tmqgpu = genpair(foldername, fp, u, xij, qi, qj, ti, tj, ai, aj, mu, eta, kappa, gen, ifile);
    tmpopu, tmpcpu, tmpgpu = gendensity(foldername, fn, u, mu, mu, eta, kappa, gen, ifile);
    stropu = stropu  * "\n" * tmqopu * "\n" * tmpopu * "\n";
    strcpu = strcpu  * "\n" * tmqcpu * "\n" * tmpcpu * "\n";
    strgpu = strgpu  * "\n" * tmqgpu * "\n" * tmpgpu * "\n";
    cppfiles(foldername, filename, stropu, strcpu, strgpu, 1);
else
    sp0, sp1, sp2, sp3 = getpotstr(3);
    sp0 = replace(sp0, "int ng)" => "int ng, int potnum)");
    sp1 = replace(sp1, "int);" => "int, int);");
    sp2 = replace(sp2, "int);" => "int, int);");
    tmpopu, tmpcpu, tmpgpu = genpotnum(filename, potnum, sp0, sp1, sp2, sp3);
    stropu = stropu * tmpopu * "\n";
    strcpu = strcpu * tmpcpu * "\n";
    strgpu = strgpu * tmpgpu * "\n";

    # here
    tmpgpu = replace(tmpgpu, "(T *u," => "Gradient(T *u, T *du, T *u_xij, T *u_xik,");
    tmpgpu = replace(tmpgpu, "(u," => "Gradient(u, du, u_xij, u_xik,");
    tmpgpu = replace(tmpgpu, "(double *" => "Gradient(double *, double *, double*, double*");
    tmpgpu = replace(tmpgpu, "(float *" => "Gradient(float *, float *, float*, float*");
    strgpu = strgpu * "\n" * tmpgpu;

    sp0, sp1, sp2, sp3 = getpotstr(2);
    sp0 = replace(sp0, "int ng)" => "int ng, int potnum)");
    sp1 = replace(sp1, "int);" => "int, int);");
    sp2 = replace(sp2, "int);" => "int, int);");
    tmpopu, tmpcpu, tmpgpu = genpotnum(fp, potnum, sp0, sp1, sp2, sp3);
    stropu = stropu * tmpopu * "\n";
    strcpu = strcpu * tmpcpu * "\n";
    strgpu = strgpu * tmpgpu * "\n";

    # here
    tmpgpu = replace(tmpgpu, "(T *u," => "Gradient(T *u, T *du, T *u_xij,");
    tmpgpu = replace(tmpgpu, "(u," => "Gradient(u, du, u_xij,");
    tmpgpu = replace(tmpgpu, "(double *" => "Gradient(double *, double *, double*");
    tmpgpu = replace(tmpgpu, "(float *" => "Gradient(float *, float *, float*");
    strgpu = strgpu * "\n" * tmpgpu;

    sp0, sp1, sp2, sp3 = getpotstr(0);
    sp0 = replace(sp0, "int ng)" => "int ng, int potnum)");
    sp1 = replace(sp1, "int);" => "int, int);");
    sp2 = replace(sp2, "int);" => "int, int);");
    tmpopu, tmpcpu, tmpgpu = genpotnum(fn, potnum, sp0, sp1, sp2, sp3);
    stropu = stropu * tmpopu;
    strcpu = strcpu * tmpcpu;
    strgpu = strgpu * tmpgpu;

    # here
    tmpgpu = replace(tmpgpu, "(T *u," => "Gradient(T *u, T *du, T *u_rho,");
    tmpgpu = replace(tmpgpu, "(u," => "Gradient(u, du, u_rho,");
    tmpgpu = replace(tmpgpu, "(double *" => "Gradient(double *, double *, double*");
    tmpgpu = replace(tmpgpu, "(float *" => "Gradient(float *, float *, float*");
    strgpu = strgpu * "\n" * tmpgpu;

    cppfiles(foldername, filename, stropu, strcpu, strgpu, 0);
end

end

function modifystr3(oldstr, natomtype)

newstr = oldstr;
for i = 1:natomtype
    s = "u[" * string(i-1) * " * i*" * string(natomtype) * "]";
    if contains(newstr, s)
        news = "if (tk[i] == " * string(i) * ") \n\t\t\t" * "u[i]";
        newstr = replace(newstr, s, news);
    end
end

return newstr

end


using Revise, Pkg, DataFrames, Finometrics, BenchmarkTools

if  pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

Pkg.activate(pwd())
using Capacity

#need this since we are uing generated functions
if @isdefined(FIRST_RUN)
  revise(Capacity)
end

const FIRST_RUN = true
#=throw("1) Determine which spec is better for the tables (see tables0 vs tables)
2) verify the SEs for the AR(1) approach- NOTE- verify the time series are no longer stationary
- maybe put in a check to prevent this from occuring in thef uture? like an INF if the coef is ≥1?
3) Convert the multiple solution simulation graphs into slides
4) Outline the anlytical solutions")=#
#throw("Need to run the funds with refresh
#also, need to do full refresh runs to account for the error in 6/9 month
#momentum")
#throw("Reset the size acceptance threshold to 20% (or smaller)")
#=throw( "- why are MF assets higher than total assets? Doesn't seem right- maybe use the corridor method,
 or try changing up the factors?
   -Try the new domestic only factor
   -Determine if we really need time FE, or an intercept, or nothing. Update the deck accordingly")=#

#throw("now run the new and imrpoved AR(1)")

@warn("
check LARS todo
check timing of shocks to momentum assets

  Why are g and ng producing the same results? repkicate and test
")

@warn("
 - check the mutual fund file and make sure the betas and assets are ok- consider trimming the dates
 via the global setting if needed, may need to aquire new mf file if necessary. Check and
   compare the winsorized results. If it doesn't make a difference, may want to disable
   or shortcircuit the winsorization
 -compute standard errors on G (weighted abs std deviation betas * A / (total assets))
  - then maybe use the delta method to compute the SE of the ratio
 - try the mutual fund assets approach, either just the growth or the bootstrap approach
 - try OOS regressions
 - consider alternate measures of participation
 - check correlation of D/P ratio w/ the lmt measure
 - long run resutls are suspect due to the persistence of the log measure
 - longer horizons- try running the monthyl data on 35/60 years and see what happens
 - short run results seem pretty good- note using G instead of aG makes sense as
 prices will rise with demand
 - the rise in prices is tough to capture because contemperaneous outflows may
 coincide with crisis periods that often have high momentum returns

")
@info("
# Important note: CUDA v2.6.3 runs fastest, v3.x.x runs slow so far.
The easiest way to transition between version is to, in the main environment,
switch to the later version.
The results seem basically equivelent if allowed to run to completion,
  although in one instance I didn't get convergence with the new a version
wrt to the reason for the gap, I tried 1) checking the broadcast
  2) the regression. Implicity, the projection isn't responsible for the
  performance regression either since the gap still exists in the
  level version. Its not limit either.
  It is stable, at leas twith the old version.
  Note the latest run is microscopicalyl different from a bit back,
  but the difference is so tiny w/ correlation 1.0, beta ≈ 1.0. More improtantly,
  the answer is stable as written with a full pass.")
@info "WISHLIST (in rough priority)
REPLICATION and VERIFICATION!!! Mutual funds or momentum?
Mutual funds- some decent correlations. Could try some randomization for placebo test
Ultimately however, if the momentum emasure works well, that is strong evidence that
outweighs even placebo tests
"
#=NOTE OLD @info "I recced this against Lewellen 2015. Returns in Lewellen are a bit lower, and the b/m metric
sd is a bit high. Size and investment (using inv4) are ok. N seems out of control, but who knows
how they are calculating that."=#

#=@ NOTE OLD @info "MCMC- Seems like A₀ is mostly fixed- probably opportunities for performance improvements
  via profiling as well as additional verificaiton that the differences between the rep
  and actual fall with sample size"=#
#@ NOTE OLD @info "a GMM style EIV model will be much more efficient (check rossi slides)"
#note- the uncontrolled version greatly increases power and correlation

function capacityscript()
  #Capacity.gz2zstd("data\\crsp\\crspd35y")
  #Capacity.gz2zstd("data\\crsp\\crspd30y")
  #Capacity.gz2zstd("data\\crsp\\crspd23y")
  #Capacity.gz2zstd("data\\crsp\\crspd20y")
  #Capacity.gz2zstd("data\\crsp\\crspd10y")
  #Capacity.gz2zstd("data\\preliminary\\prelimcwarspm")5
  #Capacity.gz2zstd("data\\crsp\\crspd2y")
  #Capacity.gz2zstd("data\\comp\\compq")
  #Capacity.gz2zstd("data\\ibes\\ibes")
  #Capacity.gz2zstd("data\\hfr\\hfralive")
  #Capacity.gz2zstd("data\\hfr\\hfrdead")

  Capacity.getparameters!() #WARNING- mutates parameter dictionary to match excel parameter file
  #Capacity.formpreliminaryfunds()z
  #Capacity.formpreliminarybeta()

  #Capacity.computemeasure()
  Capacity.analyze()

  #Capacity.capacityio()

  #Capacity.comomentumtable()
  #Capacity.comomentumgraphs()
  #Capacity.fundgraphs()

end

@time capacityscript()