#THIS CODE BASED ON N. REDDY'S THICWEED.JL DISPLAY SCRIPT

#CONSTANTS
# each tuple specifies the x, y, fontsize, xscale to print the character
# in a 100x100 box at position 100,100.
charvals_dict = Dict{Char,Tuple}('A'=> (88.5,199.,135.,1.09),
                 'C'=> (79.,196.,129.,1.19),
                 'G'=> (83.,196.,129.,1.14),
                 'T'=> (80.,199.,135.,1.24))

colour_dict = Dict{Char,String}('A'=>"(0,128,0)", #green
                 'C'=>"(0,0,128)", #blue
                 'G'=>"(255,69,0)", #yellow-brown
                 'T'=>"(150,0,0)")  #red

#char output params for string display of PWMs in nested sampler instrumentation
lwcs_vec=['a','c','g','t']
upcs_vec=['A','C','G','T']
cs_vec=[:green,:blue,:yellow,:red]
thresh_dict=Dict([(0.,'_'),(.25,'.'),(.5,"lwcs"),(1,"upcs")])

source_left_x = 10
xlen = 20
yscale = 250
scale_factor = 0.9
ypixels_per_source = 250
ypad = 10
xpad = 20
xpixels_per_position = 20
fontsize1=60
fontsize2=40

#function to convert ICA_PWM_model sources lists to PWM sequence logo diagrams
function logo_from_model(model::ICA_PWM_model,svg_output::String;freq_sort::Bool=false)
    source_tups = Vector{Tuple{AbstractFloat, Integer, Matrix{AbstractFloat}}}() #(%of sequences w/ source, prior index, weight matrix)
    mix = model.mix_matrix #o x s
    for (prior, source) in enumerate(model.sources)
        push!(source_tups, (sum(model.mix_matrix[:,prior])/size(model.mix_matrix)[1], prior, exp.(source[1])))
    end

    freq_sort && sort(source_tups)

    file = open(svg_output, "w")
    write(file, svg_header(xpad+xpixels_per_position*maximum([length(source[1]) for source in model.sources]),ypixels_per_source*length(model.sources)+ypad))

    curry = 0
    for (frequency, index, source) in source_tups
        curry += yscale
        font1y = curry-190
        ndig = ndigits(index+1)
        write(file, pwm_to_logo(source,source_left_x,curry,xlen,yscale*scale_factor))
        write(file, "<text x=\"10\" y=\"$font1y\" font-family=\"Helvetica\" font-size=\"$fontsize1\" font-weight=\"bold\" >$(index)</text>")
        write(file, "<text x=\"$(20+fontsize1*0.6*ndig)\" y=\"$(font1y-fontsize1+1.5*fontsize2)\" font-family=\"Helvetica\" font-size=\"$fontsize2\" font-weight=\"bold\" >$(frequency*100) % of sequences</text>\n")
    end

    write(file, svg_footer())
    @info "Logo written."
    close(file)
end

function svg_header(canvas_size_x, canvas_size_y)
    return """
    <?xml version="1.0" standalone="no"?>
    <!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
    "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
    <svg width="$canvas_size_x" height="$canvas_size_y"
     version="1.1" xmlns="http://www.w3.org/2000/svg">
    """
end

function pwm_to_logo(source,xpos,ypos,xlen,yscale)
    outstr = ""
    for n in 1:size(source,1)
        outstr *= print_weightvec_at_x_y(source[n,:],xpos+xlen*n,ypos,xlen,yscale)
    end
    outstr *= "\n"
    return outstr
end

function pwmstr_to_io(io,source;log=true)
    log && (source=exp.(source))
    for pos in 1:size(source,1)
        char,color=uni_wvec_params(source[pos,:])
        printstyled(io, char; color=color)
    end
end

function uni_wvec_params(wvec) #assign a single unicode char and color symbol for the most informational position in a position weight vector
    wvec.+=10^-99
    wvec = [x/sum(wvec) for x in wvec]
    infscore=(2.0 + sum([x*log(2,x) for x in wvec]))
    infvec = [x*infscore for x in wvec]
    val,idx=findmax(infvec)
    return char,color=get_char_by_thresh(idx,val)
end
                function get_char_by_thresh(idx,val)
                    char = '?'; seen_thresh=0.
                    for (thresh,threshchar) in thresh_dict
                        val >= thresh && thresh >= seen_thresh && (char=threshchar; seen_thresh=thresh)
                    end
                    char=='?' && println(val)
                    char=="lwcs" && (char=lwcs_vec[idx])
                    char=="upcs" && (char=upcs_vec[idx])
                    color=cs_vec[idx]

                    return char,color
                end

function print_weightvec_at_x_y(wvec, xpos, ypos, xlen, yscale)
    # xlen is the length occupied by that column
    # yscale is the total height for 2 information bits i.e. the
    # maximum height available for a "perfect" nucleotide
    outstr = ""
    basestr = "ACGT"
    wvec.+=10^-99  #prevent log(0) = -Inf
    wvec = [x/sum(wvec) for x in wvec] #renorm
    infscore = (2.0 + sum([x*log(2,x) for x in wvec]))
    if infscore==0.0
        return ""
    end
    wvec = [x*infscore*yscale/2.0 for x in wvec]
    # at this point, the sum of all wvec is a maximum of yscale
    wveclist = [(wvec[n],basestr[n]) for n in 1:4]
    wveclist = sort(wveclist)
    curr_ypos = ypos
    for n in 1:4
        curr_ypos -= wveclist[n][1]
        outstr *= print_char_in_rect(wveclist[n][2],xpos,curr_ypos,xlen,wveclist[n][1])
    end
    return outstr
end

blo = 1

function print_char_in_rect(c,x,y,width,height)
    raw_x, raw_y, raw_fontsize, raw_xscale = charvals_dict[c]
    raw_x = (raw_x*raw_xscale-100)/raw_xscale
    raw_y = raw_y-100

    xscale = width/100.0 * raw_xscale
    yscale = height/100.0

    scaled_x = x/xscale + raw_x
    scaled_y = y/yscale + raw_y

    return "<text x=\"$scaled_x\" y=\"$scaled_y\" font-size=\"$raw_fontsize\" font-weight=\"bold\" font-family=\"Helvetica\" fill=\"rgb"*colour_dict[c]*"\" transform=\"scale($xscale,$yscale)\">$c</text>\n"
end

function svg_footer()
    return "</svg>"
end


# This file is a part of SimilaritySearch.jl

using Dates

"""
    abstract type NeighborhoodReduction end

Overrides `Base.reduce(::NeighborhoodReduction, res::KnnResult, index::SearchGraph)` to postprocess `res` using some criteria.
Called from `find_neighborhood`, and returns a new KnnResult struct (perhaps a copy of res) since `push_neighborhood` captures
the reference of its output.
"""
abstract type NeighborhoodReduction end
abstract type LocalSearchAlgorithm end

"""
    abstract type Callback end

Abstract type to trigger callbacks after some number of insertions.
SearchGraph stores the callbacks in `callbacks` (a dictionary that associates symbols and callback objects);
A SearchGraph object controls when callbacks are fired using `callback_logbase` and `callback_starting`

"""
abstract type Callback end

### Basic operations on the index

"""
    @with_kw mutable struct Neighborhood

Determines the size of the neighborhood, \$k\$ is adjusted as a callback, and it is intended to affect previously inserted vertices.
The neighborhood is designed to consider two components \$k=in+out\$, i.e. _in_coming and _out_going edges for each vertex.
- The \$out\$ size is computed as \$minsize + \\log(logbase, n)\$ where \$n\$ is the current number of indexed elements; this is computed searching
for \$out\$  elements in the current index.
- The \$in\$ size is unbounded.
- reduce is intended to postprocess neighbors (after search process, i.e., once out edges are computed); do not change \$k\$ but always must return a copy of the reduced result set.

Note: Set \$logbase=Inf\$ to obtain a fixed number of \$in\$ nodes; and set \$minsize=0\$ to obtain a pure logarithmic growing neighborhood.

"""
@with_kw mutable struct Neighborhood
    ksearch::Int32 = 2
    logbase::Float32 = 2
    minsize::Int32 = 2
    reduce::NeighborhoodReduction = SatNeighborhood()
end

Base.copy(N::Neighborhood; ksearch=N.ksearch, logbase=N.logbase, minsize=N.minsize, reduce=copy(N.reduce)) =
    Neighborhood(; ksearch, logbase, minsize, reduce)

struct NeighborhoodSize <: Callback end

"""
    struct SearchGraph <: AbstractSearchContext

SearchGraph index. It stores a set of points that can be compared through a distance function `dist`.
The performance is determined by the search algorithm `search_algo` and the neighborhood policy.
It supports callbacks to adjust parameters as insertions are made.

- `hints`: Initial points for exploration (empty hints imply using random points)

Note: Parallel insertions should be made through `append!` or `index!` function with `parallel_block > 1`

"""
@with_kw struct SearchGraph{DistType<:SemiMetric, DataType<:AbstractDatabase, SType<:LocalSearchAlgorithm}<:AbstractSearchContext
    dist::DistType = SqL2Distance()
    db::DataType = VectorDatabase()
    links::Vector{Vector{Int32}} = Vector{Int32}[]
    locks::Vector{Threads.SpinLock} = Threads.SpinLock[]
    hints::Vector{Int32} = Int32[]
    search_algo::SType = BeamSearch()
    neighborhood::Neighborhood = Neighborhood()
    verbose::Bool = true
end

@with_kw struct SearchGraphCallbacks
    hints::Union{Nothing,Callback} = DisjointHints()
    neighborhood::Union{Nothing,Callback} = NeighborhoodSize()
    hyperparameters::Union{Nothing,Callback} = OptimizeParameters(kind=ParetoRecall())
    logbase::Float32 = 1.5
    starting::Int32 = 8
end

@inline Base.length(g::SearchGraph) = length(g.locks)
include("visitedvertices.jl")

Base.copy(g::SearchGraph;
        dist=g.dist,
        db=g.db,
        links=g.links,
        locks=g.locks,
        hints=g.hints,
        search_algo=copy(g.search_algo),
        neighborhood=copy(g.neighborhood),
        verbose=true
) = SearchGraph(; dist, db, links, locks, hints, search_algo, neighborhood, verbose)

## search algorithms

"""
    SearchGraphPools(results=GlobalKnnResult, vstates=GlobalVisitedVertices, beams=GlobalBeamKnnResult)

A set of pools to alleviate memory allocations in `SearchGraph` construction and searching. Relevant on multithreading scenarious where distance functions, `evaluate`
can call other metric indexes that can use these shared resources (globally defined).

Each pool is a vector of `Threads.nthreads()` preallocated objects of the required type.
"""
struct SearchGraphPools{VisitedVerticesType}
    beams::Vector{KnnResultShift}
    satnears::Vector{KnnResult}
    vstates::VisitedVerticesType
end

@inline function getvstate(len, pools::SearchGraphPools)
    @inbounds v = pools.vstates[Threads.threadid()]
    _init_vv(v, len)
end

@inline function getbeam(bsize::Integer, pools::SearchGraphPools)
    @inbounds reuse!(pools.beams[Threads.threadid()], bsize)
end

@inline function getsatknnresult(pools::SearchGraphPools)
    reuse!(pools.satnears[Threads.threadid()], 1)
end

getpools(::SearchGraph; beams=GlobalBeamKnnResult, satnears=GlobalSatKnnResult, vstates=GlobalVisitedVertices) = SearchGraphPools(beams, satnears, vstates)

include("beamsearch.jl")
## parameter optimization and neighborhood definitions
include("optbs.jl")
include("neighborhood.jl")
include("hints.jl")

"""
    search(index::SearchGraph, q, res; hints=index.hints, pools=getpools(index))

Solves the specified query `res` for the query object `q`.
"""
function search(index::SearchGraph, q, res::KnnResult; hints=index.hints, pools=getpools(index))
    if length(index) > 0
        search(index.search_algo, index, q, res, hints, pools)
    else
        (res=res, cost=0)
    end
end

"""
    execute_callbacks(index::SearchGraph, n=length(index), m=n+1)

Process all registered callbacks
"""
function execute_callbacks(callbacks::SearchGraphCallbacks, index::SearchGraph, n=length(index), m=n+1; force=false)
    if force || (n >= callbacks.starting && ceil(Int, log(callbacks.logbase, n)) != ceil(Int, log(callbacks.logbase, m)))
        callbacks.hints !== nothing && execute_callback(callbacks.hints, index)
        callbacks.neighborhood !== nothing && execute_callback(callbacks.neighborhood, index)
        callbacks.hyperparameters !== nothing && execute_callback(callbacks.hyperparameters, index)
    end
end


"""
    EmulatedABCRejection(
        reference_data,
        simulator_function,
        priors,
        threshold,
        n_particles,
        n_design_points;
        summary_statistic               = "keep_all",
        distance_function               = Distances.euclidean,
        batch_size                      = 10*n_particles,
        max_iter                        = 1000,
        emulator_training               = DefaultEmulatorTraining(),
        emulated_particle_selection     = MeanEmulatedParticleSelection(),
        write_progress                  = true,
        progress_every                  = 1000,
        )

Run emulation-based rejection ABC algorithm. Model simulation results are used to train the
emulator, which is then used to get the approximated distance for each particle. The rest
of the workflow is similar to [`SimulatedABCRejection`](@ref).

See [ABC Overview](@ref abc-overview) for more details.

# Mandatory arguments
- `reference_data::AbstractArray{Float,2}`: Observed data to which the simulated model output will be compared. Array dimensions sould match that of the simulator function result.
- `simulator_function::Function`: A function that takes a parameter vector as an argument and outputs model results.
- `priors::AbstractArray{ContinuousUnivariateDistribution,1}`: Continuous univariate distributions, from which candidate parameters will be sampled. Array size should match the number of parameters.
- `threshold::Float`: The ``\\varepsilon`` threshold to be used in ABC algorithm. Only those particles that produce emulated results that are within this threshold from the reference data are included into the posterior.
- `n_particles::Int`: The number of parameter vectors (particles) that will be included in the posterior.
- `n_design_points::Int`: The number of design particles that will be simulated to traing the emulator.

# Optional keyword arguments
- `summary_statistic::Union{String,AbstractArray{String,1},Function}`: Summary statistics that will be applied to the data before computing the distances. Defaults to `keep_all`. See [detailed documentation of summary statistics](@ref summary_stats).
- `distance_function::Function`: A function that will be used to compute the distance between the summary statistic of the simulated data and that of reference data. Defaults to `Distances.euclidean`.
- `batch_size::Int`: The number of particles that will be emulated on each iteration. Defaults to `1000 * n_particles`.
- `max_iter::Int`: The maximum number of emulations that will be run. Defaults to 1000.
- `emulator_training<:AbstractEmulatorTraining`: This determines how the emulator will be trained. See [`AbstractEmulatorTraining`](@ref) for more details.
- `emulated_particle_selection<:AbstractEmulatedParticleSelection`: This determines how the particles that will be added to the posterior are selected after each emulation run. See [`AbstractEmulatedParticleSelection`](@ref) for details. Defaults to [`MeanEmulatedParticleSelection`](@ref).
- `write_progress::Bool`: Whether algorithm progress should be printed on standard output. Defaults to `true`.

# Returns
An [`ABCRejectionOutput`](@ref) object.
"""
function EmulatedABCRejection(reference_data::AbstractArray{AF,2},
    simulator_function::Function,
    priors::AbstractArray{D,1},
    threshold::AF,
    n_particles::Int,
    n_design_points::Int;
    summary_statistic::Union{String,AbstractArray{String,1},Function}="keep_all",
    distance_function::Function=Distances.euclidean,
    emulation_type::AbstractEmulatorTraining = DefaultEmulatorTraining(),
    batch_size::Int=10*n_particles,
    max_iter::Int=1000,
    emulator_training::ET = DefaultEmulatorTraining(),
    emulated_particle_selection::EPS = MeanEmulatedParticleSelection(),
    kwargs...) where {
    AF<:AbstractFloat,
    D<:ContinuousUnivariateDistribution,
    EPS<:AbstractEmulatedParticleSelection,
    ET<:AbstractEmulatorTraining
    }

    summary_statistic = build_summary_statistic(summary_statistic)
    reference_summary_statistic = summary_statistic(reference_data)

    emulator_training_input = EmulatorTrainingInput(
        DistanceSimulationInput(reference_summary_statistic, simulator_function, summary_statistic, distance_function),
        n_design_points,
        emulator_training)

    input = EmulatedABCRejectionInput(length(priors), n_particles, threshold,
        priors, batch_size, max_iter, emulator_training_input, emulated_particle_selection)

    return ABCrejection(input; kwargs...)
end

"""
    EmulatedABCSMC(
        reference_data,
        simulator_function,
        priors,
        threshold_schedule,
        n_particles,
        n_design_points;
        summary_statistic               = "keep_all",
        distance_function               = Distances.euclidean,
        batch_size                      = 10*n_particles,
        max_iter                        = 1000,
        emulator_training               = DefaultEmulatorTraining(),
        emulator_retraining             = NoopRetraining(),
        emulated_particle_selection     = MeanEmulatedParticleSelection(),
        write_progress                  = true,
        progress_every                  = 1000,
        )

Run emulation-based ABC-SMC algorithm. This is similar to [`EmulatedABCRejection`](@ref),
the main difference being that an array of thresholds is provided instead of a single threshold.
It is assumed that thresholds are sorted in decreasing order.

An emulation based ABC iteration is executed for each threshold. For the first threshold, the provided prior is used.
For each subsequent threshold, the posterior from the previous iteration is used as a prior.

See [ABC Overview](@ref abc-overview) for more details.

# Mandatory arguments
- `reference_data::AbstractArray{Float,2}`: Observed data to which the simulated model output will be compared. Array dimensions sould match that of the simulator function result.
- `simulator_function::Function`: A function that takes a parameter vector as an argument and outputs model results.
- `priors::AbstractArray{ContinuousUnivariateDistribution,1}`: Continuous univariate distributions, from which candidate parameters will be sampled during the first iteration. Array size should match the number of parameters.
- `threshold_schedule::AbstractArray{Float,1}`: The threshold schedule to be used in ABC algorithm. An ABC iteration is executed for each threshold. It is assumed that thresholds are sorted in decreasing order.
- `n_particles::Int`: The number of parameter vectors (particles) that will be included in the posterior.
- `n_design_points::Int`: The number of design particles that will be simulated to traing the emulator.

# Optional keyword arguments
- `summary_statistic::Union{String,AbstractArray{String,1},Function}`: Summary statistics that will be applied to the data before computing the distances. Defaults to `keep_all`. See [detailed documentation of summary statistics](@ref summary_stats).
- `distance_function::Function`: A function that will be used to compute the distance between the summary statistic of the simulated data and that of reference data. Defaults to `Distances.euclidean`.
- `batch_size::Int`: The number of particles that will be emulated on each iteration. Defaults to `1000 * n_particles`.
- `max_iter::Int`: The maximum number of emulations that will be run. Defaults to 1000.
- `emulator_training<:AbstractEmulatorTraining`: This determines how the emulator will be trained for each iteration. See [`AbstractEmulatorTraining`](@ref) for more details.
- `emulator_retraining<:AbstractEmulatorRetraining`: This is used to specify parameters of additional emulator retraining that can be done for each iteration. By default this retraining is switched off ([`NoopRetraining`](@ref)). See [`AbstractEmulatorRetraining`] for more details.
- `emulated_particle_selection<:AbstractEmulatedParticleSelection`: This determines how the particles that will be added to the posterior are selected after each emulation run. See [`AbstractEmulatedParticleSelection`](@ref) for details. Defaults to [`MeanEmulatedParticleSelection`](@ref).
- `write_progress::Bool`: Whether algorithm progress should be printed on standard output. Defaults to `true`.

# Returns
An [`ABCSMCOutput`](@ref) object.
"""
function EmulatedABCSMC(reference_data::AbstractArray{AF,2},
    simulator_function::Function,
    priors::AbstractArray{D,1},
    threshold_schedule::AbstractArray{AF,1},
    n_particles::Int,
    n_design_points::Int;
    summary_statistic::Union{String,AbstractArray{String,1},Function}="keep_all",
    emulator_training::ET = DefaultEmulatorTraining(),
    distance_metric::Function=Distances.euclidean,
    batch_size::Int=10*n_particles,
    max_iter::Int=20,
    emulator_retraining::ER = NoopRetraining(),
    emulated_particle_selection::EPS = MeanEmulatedParticleSelection(),
    kwargs...) where {
    AF<:AbstractFloat,
    ET<:AbstractEmulatorTraining,
    ER<:AbstractEmulatorRetraining,
    EPS<:AbstractEmulatedParticleSelection,
    D<:ContinuousUnivariateDistribution
    }

    summary_statistic = build_summary_statistic(summary_statistic)
    reference_summary_statistic = summary_statistic(reference_data)

    emulator_training_input = EmulatorTrainingInput(
        DistanceSimulationInput(reference_summary_statistic, simulator_function, summary_statistic, distance_metric),
        n_design_points,
        emulator_training)

    input = EmulatedABCSMCInput(length(priors), n_particles, threshold_schedule,
        priors, batch_size, max_iter, emulator_training_input, emulator_retraining, emulated_particle_selection)

    return ABCSMC(input; kwargs...)
end

"""
    EmulatedModelSelection

Perform model selection using emulation-based ABC.

# Arguments
- `n_design_points::Int64`: The number of parameter vectors used to train the Gaussian process emulator.
- `reference_data::AbstractArray{Float64,2}`: The observed data to which the simulated model output will be compared. Size: (n_model_trajectories, n_time_points)
- `n_particles::Int64`: The number of parameter vectors (particles) that will be included in the final posterior.
- `threshold_schedule::AbstractArray{Float64}`: A set of maximum distances from the summarised model output to summarised observed data for a parameter vector to be included in the posterior. Each distance will be used in a single run of the ABC-SMC algorithm.
- `parameter_priors::AbstractArray{AbstractArray{ContinuousUnivariateDistribution},1}`: Priors for the parameters of each model. The length of the outer array is the number of models.
- `summary_statistic::Union{String,AbstractArray{String,1},Function}`: Either: 1. A `String` or 1D Array of strings that Or 2. A function that outputs a 1D Array of Floats that summarises model output. Defaults to `keep_all`. See [detailed documentation of summary statistics](@ref summary_stats).
- `simulator_functions::AbstractArray{Function,1}`: An array of functions that take a parameter vector as an argument and outputs model results (one per model).
- 'model_prior::DiscreteUnivariateDistribution': The prior from which models are sampled. Default is a discrete, uniform distribution.
- `distance_function::Function`: Any function that computes the distance between 2 1D Arrays. Optional argument (default is to use the Euclidean distance).
- `max_iter::Int`: The maximum number of simulations that will be run. The default is 1000*`n_particles`. Each iteration samples a single model and performs ABC using a single particle.
- `max_batch_size::Int`: The maximum batch size for the emulator when making predictions.

# Returns
A [`ModelSelectionOutput`](@ref) object that contains which models are supported by the observed data.
"""
function EmulatedModelSelection(
    reference_data::AbstractArray{AF,2},
    simulator_functions::AbstractArray{Function,1},
    parameter_priors::AbstractArray{AD,1},
    threshold_schedule::AbstractArray{AF,1},
    n_particles::Int,
    n_design_points::Int,
    model_prior::DiscreteUnivariateDistribution=Distributions.DiscreteUniform(1,length(parameter_priors));
    summary_statistic::Union{String,AbstractArray{String,1},Function}="keep_all",
    distance_function::Function=Distances.euclidean,
    max_iter::Int=1000,
    max_batch_size::Int=1000,
    emulator_training::ET = DefaultEmulatorTraining(),
    emulator_retraining::ER = NoopRetraining(),
    emulated_particle_selection::EPS = MeanEmulatedParticleSelection(),
    ) where {
    D<:ContinuousUnivariateDistribution,
    AD<:AbstractArray{D,1},
    AF<:AbstractFloat,
    ET<:AbstractEmulatorTraining,
    ER<:AbstractEmulatorRetraining,
    EPS<:AbstractEmulatedParticleSelection,
    }

    summary_statistic = build_summary_statistic(summary_statistic)
    reference_summary_statistic = summary_statistic(reference_data)
    emulator_training_input = [EmulatorTrainingInput(
        DistanceSimulationInput(reference_summary_statistic, simulator_functions[m], summary_statistic, distance_function),
        n_design_points, emulator_training) for m in 1:length(parameter_priors)]

    input = EmulatedModelSelectionInput(length(parameter_priors),
        n_particles,
        threshold_schedule,
        model_prior,
        parameter_priors,
        emulator_training_input,
        emulator_retraining,
        emulated_particle_selection,
        max_batch_size,
        max_iter)

    return model_selection(input)
end


using RCall, MixedModels, Test
using StatsBase: zscore
const LMM = LinearMixedModel
const GLMM = GeneralizedLinearMixedModel

@testset "merMod" begin
    # this is available in MixedModels.dataset(:sleepstudy) but with different
    # capitalization than in R
    sleepstudy = rcopy(R"sleepstudy")
    jlmm = fit!(LMM(@formula(Reaction ~ 1 + round(Days) + (1|Subject)),sleepstudy), REML=false)
    @testset "bare model" begin
        @test_throws ArgumentError (@rput jlmm)
    end
    @testset "reversed tuple" begin
        jm = (sleepstudy, jlmm);
        @test_throws ArgumentError (@rput jm)
    end
end


#!/usr/local/bin/julia
# Assumes agent goes first always for now
using OpenAIGym
using go

const debug = false

model_dir = joinpath(splitdir(@__FILE__)[1], "../models")

function usage()
    println(STDOUT, "Usage: ./openai_gym.jl exp_name rounds {random|keras [model name]}")
    exit()
end

nargs = length(ARGS)
exp_name = ARGS[1]
nrounds = parse(Int, ARGS[2])

if nargs != 3
    if ARGS[3] == "random"
        policy = go.RandomPolicy()
    elseif ARGS[3] == "keras"
        if nargs != 4
            usage()
        end
        policy = go.KerasNetwork(model_dir, ARGS[4])
    else
        usage()
    end
else
    usage()
end

function getempty(board)
    reshape(board[3, :, :], (go.N, go.N)).'
end

function get_opponent_move(last_empty, cur_empty, last_move::Int=0)
    # @show last_empty
    # @show cur_empty
    move = 0
    for i = 1:(go.N * go.N)
        if last_empty[i] == 1 && cur_empty[i] == 0 && i != last_move
            move = i
        end
    end
    if move == 0
        @show cur_empty
        @show last_empty
        @show last_move
        return go.PASS_MOVE
    else
        return from_gym_coords(go.pointindex(move))
    end
end

"""
gym indexes in row order from upper left of the board
 while we index from lower left along columns.
Together with from_gym_coords, transform between them.
"""
function to_gym_coords(p::go.Point)
    go.Point(p[2], go.N-p[1] + 1)
end

function from_gym_coords(p::go.Point)
    go.Point(go.N - p[2] + 1, p[1])
end

function run_episode()
    state = reset(env)
    localboard = go.Board()

    # TODO: The coordinate transforms are hacky and uncommented, but they work.
    last_empty = getempty(state.observation)
    opponent_move = go.EMPTY_MOVE # Init here so it stays in scope for after loop
    while !state.done
        # Our move
        if opponent_move == go.PASS_MOVE
            move = go.PASS_MOVE
            state = step(env, 81)
            break
        else
            move = go.choose_move(localboard, policy)
            state = step(env, go.linearindex(to_gym_coords(move)) - 1)
        end
        go.play_move(localboard, move)

        # Opponent move
        board = state.observation
        empty = getempty(board)
        opponent_move = get_opponent_move(last_empty, empty, go.linearindex(to_gym_coords(move)))

        if debug
            @show move
            @show opponent_move
        end

        go.play_move(localboard, opponent_move)
        last_empty = empty

        if debug
            println(localboard)
            display(env)
        end
        flush(STDOUT)
        flush(STDERR)
    end
    println(localboard)
    display(env)
    @show opponent_move
    @show state.reward
    @show go.calculate_score(localboard)
end

env = Env("Go$(go.N)x$(go.N)-v0")
env.env[:monitor][:start](exp_name)

for episode in 1:nrounds
    println("Episode $(episode)/$(nrounds)")
    run_episode()
end

env.env[:monitor][:close]()


module DemoProject

export calcHypotenuse
include("pythagoras.jl")
end


type XYZ <: AtomFileType end

file_exts(::Type{XYZ}) = (".xyz",)

function Atoms(::Type{XYZ}, num_atoms::Int)
    Atoms(num_atoms,
        :type => Vector{Symbol}(num_atoms),
        :coords => Matrix{Float64}(3, num_atoms),
        :molecule => Vector{Int}(num_atoms)
    )
end

function read_frame_header{T}(::Type{T}, ::Type{XYZ}, chunk::Chunk)
    frame_start = chunk.i

    num_atoms = read_int(chunk)
    read_newlines(chunk, 2)

    atom_data_start = chunk.i
    read_newlines(chunk, num_atoms)
    atom_data_size = chunk.i - atom_data_start

    frame_size = chunk.i - frame_start

    frame_header = FrameHeader(frame_start, frame_size, num_atoms, atom_data_start, atom_data_size)

    frame = Frame()

    return frame_header, frame
end

function read_atoms{T}(::Type{T}, ::Type{XYZ}, chunk::Chunk, atom::Atoms, indices::UnitRange{Int})
    types = atom.props[:type]::Vector{Symbol}
    coords = atom.props[:coords]::Matrix{T}
    molecules = atom.props[:molecule]::Vector{Int}

    for i in indices
        atom_type = read_symbol(chunk)
        read_whitespace(chunk)

        x = read_float(T, chunk)
        read_whitespace(chunk)

        y = read_float(T, chunk)
        read_whitespace(chunk)

        z = read_float(T, chunk)
        read_whitespace(chunk)

        molecule = read_int(chunk)
        read_newline(chunk)

        types[i] = atom_type
        coords[1, i] = x
        coords[2, i] = y
        coords[3, i] = z
        molecules[i] = molecule
    end
end

function write_atoms(::Type{XYZ}, io::IO, frame::Frame)
    haskey(frame, :atoms) || throw(ArgumentError("Need frame atoms"))

    atoms = frame[:atoms]
    haskey(atoms, :type) || throw(ArgumentError("Need atom types"))
    haskey(atoms, :coords) || throw(ArgumentError("Need atom coords"))

    num_atoms = length(atoms)
    println(io, num_atoms)

    haskey(frame, :timestep) && print(io, "t = ", frame[:timestep], ' ')

    if haskey(frame, :box_max)
        xhi, yhi, zhi = frame[:box_max]
        xlo, ylo, zlo = 0, 0, 0

        if haskey(frame, :box_min)
            xlo, ylo, zlo = frame[:box_min]
        end

        print(io, "box = ", xhi - xlo, ' ', yhi - ylo, ' ', zhi - zlo, ' ')
    end

    println(io)

    has_extra = haskey(atoms, :molecule) || haskey(atoms, :energy)

    for (i, atom) in enumerate(atoms)
        print(io, atom[:type], ' ', atom[:coords][1], ' ', atom[:coords][2], ' ', atom[:coords][3])

        haskey(atoms, :charge) && print(io, ' ', atom[:charge])

        if has_extra
            print(io, " #")
            haskey(atoms, :molecule) && print(io, ' ', atom[:molecule])
            haskey(atoms, :energy) && print(io, ' ', atom[:energy])
        end

        println(io)
    end
end


export RandomIdGenerator

using Random

"""
Customized id generators must implement the following two methods:

  - `generate_span_id(::AbstractIdGenerator)`
  - `generate_trace_id(::AbstractIdGenerator)`
"""
abstract type AbstractIdGenerator end

"""
    RandomIdGenerator(rng=Random.GLOBAL_RNG)

Use the `rng` to generate a random trace id or span id.
"""
Base.@kwdef struct RandomIdGenerator{R} <: AbstractIdGenerator
    rng::R = Random.GLOBAL_RNG
end

generate_trace_id(rng::RandomIdGenerator) = rand(rng.rng, TraceIdType)
generate_span_id(rng::RandomIdGenerator) = rand(rng.rng, SpanIdType)


"""
    get_3D_field(product_year, tracer, period, resolution, field)

Downloads and returns the 3D field of a tracer from the World Ocean Atlas.

Availabe product years are 2009, 2013, and 2018.
(2018 more likely to work.)
Tracers are either "phosphate", "nitrate", or "silicate".
Resolution is either "5°", "1°", or "0.25°".
(1° most likely to work.)
Fields are "mean" or "objectively analyzed climatology".
(mean most likely to work.)
Note WOA's nomenclature should work too, e.g.,
"p" for "phosphate", "mn" for "mean", and so on.
"""
function get_3D_field(product_year, tracer, period, resolution, field)
    println("Getting the 3D field of WOA$(my_product_year(product_year)) $(my_averaging_period(period)) $(WOA_path_varname(tracer)) $(surface_grid_size(resolution)) data")
    ds = WOA_Dataset(product_year, tracer, period, resolution)
    field3D = ds[WOA_varname(tracer, field)][:][:, :, :, 1]
    println("  Rearranging data")
    field3D = permutedims(field3D, [2 1 3])
    return field3D
end

function get_gridded_3D_field(ds, tracer, field)
    println("  Reading NetCDF file")
    field3D = ds[WOA_varname(tracer, field)][:][:, :, :, 1]
    lon   = ds["lon"][:]   .|> Float64
    lat   = ds["lat"][:]   .|> Float64
    depth = ds["depth"][:] .|> Float64
    println("  Rearranging data")
    # Reorder the variable index order (lat <-> lon from WOA to OCIM)
    field3D = permutedims(field3D, [2 1 3])
    # Rearrange longitude range (from WOA data, which is -180:180, to 0:360)
    lon = mod.(lon, 360)
    lon_reordering = sortperm(lon)
    lon = lon[lon_reordering]
    field3D .= field3D[:, lon_reordering, :]
    return field3D, lat, lon, depth
end

"""
    mean_std_and_number_obs(ds, tracer)

Returns a DataFrame containing the following columns
lat, lon, depth, mean, std, nobs
"""
function mean_std_and_number_obs(ds, tracer)
    println("  Reading NetCDF file")
    lon   = mod.(ds["lon"][:] .|> Float64, 360)
    lat   = ds["lat"][:]   .|> Float64
    depth = ds["depth"][:] .|> Float64
    μ3D    = ds[WOA_varname(tracer, "mn")][:][:, :, :, 1]
    σ3D    = ds[WOA_varname(tracer, "sd")][:][:, :, :, 1]
    nobs3D = ds[WOA_varname(tracer, "dd")][:][:, :, :, 1]
    println("  Filtering missing data")
    CI = findall(@. !ismissing(μ3D) & !ismissing(nobs3D) & !iszero(nobs3D))
    lon1D   = lon[map(x -> x.I[1], CI)]
    lat1D   = lat[map(x -> x.I[2], CI)]
    depth1D = depth[map(x -> x.I[3], CI)]
    μ1D     = μ3D[CI] .|> Float64
    σ1D     = σ3D[CI] .|> Float64
    nobs1D  = nobs3D[CI] .|> Int64
    return DataFrame(
                     :Latitude => lat1D,
                     :Longitude => lon1D,
                     :Depth => depth1D,
                     :Mean => μ1D,
                     :Std => σ1D,
                     :n_obs => nobs1D,
                    )
end

function filter_gridded_3D_field(field3D, lat, lon, depth)
    # Find where there is data for both mean and std
    println("  Filtering data")
    CI = findall(x -> !ismissing(x) && !iszero(x), field3D) # filter out fill-values and 0's
    fieldvec = field3D[CI]
    latvec = lat[map(x -> x.I[1], CI)]
    lonvec = lon[map(x -> x.I[2], CI)]
    depthvec = depth[map(x -> x.I[3], CI)]
    return fieldvec, latvec, lonvec, depthvec, CI
end

get_unit(ds, tracer, field) = convert_to_Unitful(ds[WOA_varname(tracer, field)].attrib["units"])

function mean_and_variance_gridded_3d_field(grid::OceanGrid, field3D, lat, lon, depth)
    fieldvec, latvec, lonvec, depthvec, CI = filter_gridded_3D_field(field3D, lat, lon, depth)
    println("  Averaging data over each grid box")
    χ_3D = zeros(size(grid))
    σ²_3D = zeros(size(grid))
    n_3D = zeros(size(grid))
    # Use NearestNeighbors to bin into OceanGrid
    gridbox_centers = [ustrip.(vec(grid.lat_3D)) ustrip.(vec(grid.lon_3D)) ustrip.(vec(grid.depth_3D))] # TODO Maybe add option for using iwet instead of full vec
    gridbox_centers = permutedims(gridbox_centers, [2, 1])
    kdtree = KDTree(gridbox_centers)
    for i in eachindex(CI)
        idx = knn(kdtree, [lonvec[i], latvec[i], depthvec[i]], 1, true)[1][1]
        χ_3D[idx] += fieldvec[i]      # μ = Σᵢ μᵢ / n
        σ²_3D[idx] += fieldvec[i]^2   # σ² = Σᵢ μᵢ² / n - μ²
        n_3D[idx] += 1
    end
    χ_3D .= χ_3D ./ n_3D               # μ = Σᵢ μᵢ / n
    σ²_3D .= σ²_3D ./ n_3D .- χ_3D.^2  # σ² = Σᵢ μᵢ² / n - μ²
    # Enforce μ = 0 and σ = ∞ where no observations
    # (Instead of NaNs)
    println("  Setting μ = 0 and σ² = ∞ where no obs")
    χ_3D[findall(n_3D .== 0)] .= 0.0
    σ²_3D[findall(n_3D .== 0)] .= Inf
    println("  Setting a realistic minimum for σ²")
    meanχ = mean(χ_3D, weights(n_3D))
    σ²_3D .= max.(σ²_3D, 1e-4meanχ^2)
    return χ_3D, σ²_3D
end

function convert_to_SI_unit!(χ_3D, σ²_3D, ds, tracer, field)
    # Convert to SI units
    χ_unit = get_unit(ds, tracer, field)
    χ_3D .*= ustrip(upreferred(1.0χ_unit))
    σ²_3D .*= ustrip(upreferred(1.0χ_unit^2))
end

function fit_to_grid(grid::OceanGrid, product_year, tracer, period, resolution, field)
    ds = WOA_Dataset(product_year, tracer, period, resolution)
    field3D, lat, lon, depth = get_gridded_3D_field(ds, tracer, field)
    χ_3D, σ²_3D = mean_and_variance_gridded_3d_field(grid, field3D, lat, lon, depth)
    convert_to_SI_unit!(χ_3D, σ²_3D, ds, tracer, field)
    return χ_3D, σ²_3D
end

#==================================
Helper functions
==================================#

function WOA_Dataset(product_year, tracer, period, resolution; verbose=false)
    println("Registering WOA data with DataDeps")
    verbose && print(": ", url_WOA_THREDDS(product_year, tracer, period, resolution))
    register_WOA(product_year, tracer, period, resolution)
    ddstr = @datadep_str string(my_DataDeps_name(product_year, tracer, period, resolution),
                        "/",
                        WOA_NetCDF_filename(product_year, tracer, period, resolution))
    Dataset(ddstr)
end

function remove_DataDep(product_year, tracer, period, resolution)
    println("Removing WOA$(my_product_year(product_year)) $(my_averaging_period(period)) $(WOA_path_varname(tracer)) $(surface_grid_size(resolution)) data")
    nc_file = @datadep_str string(my_DataDeps_name(product_year, tracer, period, resolution), "/", WOA_NetCDF_filename(product_year, tracer, period, resolution))
    rm(nc_file; recursive=true, force=true)
end


##################################################
# Raw apply/combine
##################################################

unwrap(x) = x
function unwrap(x::Sym{TAG}; context::Context = @__context__) where {TAG}
    hashead(x, :object) && return firstarg(x)
    if hashead(x, :call)
        op = firstarg(x)
        if hashead(op, :fn)
            xs = tailargs(x)
            f = firstarg(op)
            tags = map(tagof, xs)
            if all(tag -> tag isa Type, tags) && hasmethod(f, Tuple{tags...})
                if all(hashead(:object), xs) && all(isfaithful, xs)
                    return invoke(f, Tuple{tags...}, map(firstarg, xs)...)
                elseif all(ismocking, xs) && all(isabstracttype, tags)
                    return mock(f, xs...; context = context)
                end
            end
        end
    end
    return x
end

_apply(args...) = _apply(map(arg -> convert(Sym, arg), args)...)
function _apply(op::Sym, xs::Sym...)
    tags = map(tagof, xs)
    TAG = promote_tag(:call, op, tags...)
    x = Sym{TAG}(:call, op, xs...)
    return x::Sym{TAG}
end

function _combine(head::Symbol, xs...)
    TAG = promote_tag(head, map(tagof, xs)...)
    return Sym{TAG}(head, xs...)
end

##################################################
# Querying apply/combine
##################################################

function apply(op, xs...; context::Context = @__context__)
    return unwrap(query!(_apply(op, xs...), context = context), context = context)
end

function combine(head::Symbol, xs...; context::Context = @__context__)
    unwrap(query!(_combine(head, xs...), context = context), context = context)
end


#Подсчитать число и вертикальных и горизонтальных прямолинейных перегородок (прямоугольных - нет)
include("library.jl")
function count_obstacles(r)
    back_path = BackPath(r, (Sud, West))
    print(counterPartition!(r), "\tHorizontal\n")
    BackPath(r, (Nord, West))
    print(counterPartition!(r, Sud), "\tVertical\n")
    BackPath(r, (Sud, West))
    back!(r, back_path)
end


"""
    MOKernel

An abstract type for multi-output kernels.
"""
abstract type MOKernel <: Kernel end


# module links
# """
#     Links: state,input -> output for each subsystem, of type DiffEqFlux.FastLayer

#     interface includes function of form h(xu, p) where xu is a concatenation of state and input, parameters p
# """
# using DiffEqFlux

# abstract type Link <: DiffEqFlux.FastLayer end

# include("affine_canonical.jl")
# export AffineFeedthrough, LinearFeedthrough, AffineCanonical, LinearCanonical, AffineResistive, LinearResistive
# export paramlength, initial_params, convert_verilogA

# end # module

using DiffEqFlux

abstract type Link <: DiffEqFlux.FastLayer end

include("affine_layer.jl")

""" Required as a FastLayer, define a paramsum field """
paramlength(f::AffineLayer) = f.paramsum

""" Required as a FastLayer """
initial_params(f::AffineLayer) = f.initial_params()

""" convert xu concatenation to x,u in 1d case """
@inline function extract_inputs(f::L, xu::AbstractArray{T,1}) where {T} where {L<:Link}
    x = view(xu, 1:f.state_dim)
    u = view(xu, f.state_dim+1:f.state_dim+f.input_dim)
    return x, u
end

""" convert xu concatenation to x,u in 2d case """
@inline function extract_inputs(f::L, xu::AbstractArray{<:T,2}) where {T} where {L<:Link}
    x = view(xu, 1:f.state_dim, :)
    u = view(xu, f.state_dim+1:f.state_dim+f.input_dim, :)
    return x, u
end

""" convert xu concatenation to x,u in 3d case """
@inline function extract_inputs(f::L, xu::AbstractArray{T,3}) where {T} where {L<:Link}
    x = view(xu, 1:f.state_dim, :, :)
    u = view(xu, f.state_dim+1:f.state_dim+f.input_dim, :, :)
    return x, u
end

""" Identity Link"""
struct IdentityLink{I} <: Link
    initial_params::I
    paramsum::Int
    function IdentityLink(initb = Flux.zeros)
        initial_params() = initb(0)
        new{typeof(initial_params)}(initial_params, 0)
    end
end

function (f::IdentityLink)(x, p)
    return x
end


# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Script to compile data for some descriptive statistics
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

using DataFrames, CSV, XLSX, LinearAlgebra, Statistics, StatsBase, StatsPlots, Dates

# other scripts
dir_home = "x:/VIVES/1-Personal/Florian/git/2022_01_comext/src/"
include(dir_home * "functions.jl")

# location of data input/output (io)
dir_io = "C:/Users/u0148308/data/comext/"
dir_dropbox = "C:/Users/u0148308/Dropbox/BREXIT/"

# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# Notes:
#   - figure 1: evolution of total imports/exports from BE to UK vis-a-vis FR, NL, DE and EU
#       + HP filter available, what do chose for λ?
#   - figure 2: evolution of imports/exports from BE to UK at 2digit level
#       + HP filter available, what do chose for λ?
#   - figure 3: evolution of share of imports/exports from BE to UK of total imports/exports from BE
#       + also add FR, DE, NL, EU27
#   - table 1: evolution of most important products at 2digit level
#       + shares per product available for bilteral trade (i.e. sum to 100%) or for total trade (i.e. sum to figure 3)
#       + could compute figure 3 from data of table 1 as well!
#
#   - all figures in values and unit values (unit prices):
#       + compute a price index as using a weighted sum
#   - include vertical lines for different Brexit events

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Brexit events
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# vertical line
# vline!([Date(2016,6,23)], label="", color=:black, lw=2) # referendum
# vline!([Date(2016,12,07)], label="") # vote to trigger Article 50
# vline!([Date(2017,03,29)], label="") # Article 50
# vline!([Date(2017,06,19)], label="") # negotiations commence
# vline!([Date(2019,03,21)], label="") # first extension of Article 50
# vline!([Date(2019,04,10)], label="") # second extension of Article 50
# vline!([Date(2019,10,28)], label="") # third extension of Article 50
# vline!([Date(2019,12,20)], label="") # withdrawl bill passes UK
# vline!([Date(2020,01,31)], label="") # exit of UK from EU
# vline!([Date(2020,12,20)], label="") # UK passes Trade and Cooperation Agreement
# vline!([Date(2020,12,31)], label="") # transition period ends
# vline!([Date(2021,04,27)], label="") # EU passes Trade and Cooperation Agreement

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Raw data
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

df_comext = CSV.read(dir_dropbox * "rawdata/" * "df_comext_BE" * ".csv", DataFrame)

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Figure 1
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# prepare data
df = transform(df_comext, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> ifelse(iszero(x) | ismissing(x), missing, x)), renamecols=false) # need to remove missing also from EU
subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> !ismissing(x)))

# remove outliers based on medians abs dev.
# Notes:
#   - where to put MAD (need to drop EU/WORLD for sure)
outlier_cutoff = 3
transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)
gdf = groupby(df, "PRODUCT_NC")
df = transform(gdf, :UNIT_PRICE => MAD_method => :MAD)
subset!(df, :MAD => ByRow(x -> x < outlier_cutoff)) # lose further 1-42083391/50055769 ~16% observations

# produce data for plotting
df = data_fig1(df, ["BE"], ["GB", "EU", "DE", "FR", "NL"])

# -----------------------

# formatting
transform!(df, [:DECLARANT_ISO, :PARTNER_ISO, :FLOW] .=> ByRow(string), renamecols=false)
transform!(df, :PERIOD => ByRow(x -> Date(string(x), DateFormat("yyyymm"))) => :DATE)

# HP filter
λ = 20
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW"]
gdf = groupby(df, cols_grouping)
transform!(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :PRICE_INDEX] .=> (x -> HP(x, λ)) .=> [:VALUE_IN_EUROS_HP, :QUANTITY_IN_KG_HP, :PRICE_INDEX_HP])

# create unweighted unit prices
transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)
transform!(df, [:VALUE_IN_EUROS_HP, :QUANTITY_IN_KG_HP] => ByRow((v,q) -> v/q) => :UNIT_PRICE_HP)

# create 3MMA as observations are very volatile
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW"]
gdf = groupby(df, cols_grouping)
df = transform(gdf, :PRICE_INDEX => (x -> movingaverage(x,3)) => :PRICE_INDEX_3MMA)

# plotting
for flow in ["imports", "exports"]

    # values
    # Notes:
    #   - take out EU for scale
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> x != "EU")) plot(:DATE, :VALUE_IN_EUROS,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="euros", title="Belgian "* flow*": values")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/" * "fig1_" * flow * "_values" * ".png") # export image dropbox

    # price index (weighted unit prices)
    # Notes:
    #   - unit prices are extremely volatile, 3 MMA would be useful
    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :PRICE_INDEX,
        group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="price index (weighted)", title="Belgian "*flow*": price index (weighted)")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/" * "fig1_" * flow * "_prices_weighted" * ".png")

    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :PRICE_INDEX_3MMA,
        group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="price index (weighted, 3MMA)", title="Belgian "*flow*": price index (weighted, 3MMA)")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/" * "fig1_" * flow * "_prices_weighted_3MMA" * ".png")

    # price index (unweighted unit prices)
    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :UNIT_PRICE,
        group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="price index (unweighted)", title="Belgian "*flow*": price index (unweighted)")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/" * "fig1_" * flow * "_prices_unweighted" * ".png")

end

# plotting of HP
for flow in ["imports", "exports"]

    # values
    # Notes:
    #   - take out EU for scale
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> x != "EU")) plot(:DATE, :VALUE_IN_EUROS_HP,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="euros", title="Belgian "*flow*": values (HP, λ=$λ)")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/HP/" * "fig1_" * flow * "_values" * "_HP" * ".png") # export image dropbox

    # price index (weighted unit prices)
    # Notes:
    #   - unit prices are extremely volatile, 3 MMA would be useful
    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :PRICE_INDEX_HP,
        group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="price index (weighted)", title="Belgian "*flow*": price index (weighted, λ=$λ)")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/HP/" * "fig1_" * flow * "_prices_weighted" * "_HP" * ".png")

    # price index (unweighted unit prices)
    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :UNIT_PRICE_HP,
        group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="price index (unweighted)", title="Belgian "*flow*": price index (unweighted, λ=$λ)")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig1/HP/" * "fig1_" * flow * "_prices_unweighted" * "_HP" * ".png")

end

# BE vis-a-vis GB
p = @df subset(df, :PARTNER_ISO => ByRow(x -> x == "GB")) plot(:DATE, :UNIT_PRICE_HP,
        group=:FLOW, lw=2, legend=:topleft, ylabel="euros", title="Price index (unweighted, λ=$λ): \n Belgium vis-a-vis Great Britain")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
savefig(p, dir_dropbox * "results/images/comext/fig1/HP/" * "fig1_" * "BE_GB" * "_prices_unweighted" * "_HP" * ".png") # export image

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Figure 4
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# rolling standard deviation
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW"]
gdf = groupby(df, cols_grouping)
df = transform(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :PRICE_INDEX, :UNIT_PRICE, :VALUE_IN_EUROS_HP, :QUANTITY_IN_KG_HP, :PRICE_INDEX_HP, :UNIT_PRICE_HP]
                .=> (x -> rolling_std(x, 6)) .=>
                [:STD_VALUE, :STD_QUANTITY, :STD_PRICE, :STD_UNIT, :STD_VALUE_HP, :STD_QUANTITY_HP, :STD_PRICE_HP, :STD_UNIT_HP])

# plotting HP
for flow in ["imports", "exports"]

    # values
    # Notes:
    #   - take out EU for scale
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> x != "EU")) plot(:DATE, :STD_UNIT_HP,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="euros", title="Belgian "* flow*": std unit prices (unweighted, λ=$λ)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig4/HP/" * "fig4_" * flow * "_std" * "_prices_unweighted" * "_HP" * ".png") # export image dropbox

end

# BE vis-a-vis GB
p = @df subset(df, :PARTNER_ISO => ByRow(x -> x == "GB")) plot(:DATE, :STD_UNIT_HP,
        group=:FLOW, lw=2, legend=:topleft, ylabel="euros", title="6 months STD price index (unweighted, λ=$λ): \n Belgium vis-a-vis Great Britain")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
savefig(p, dir_dropbox * "results/images/comext/fig4/HP/" * "fig4_" * "BE_GB" * "_STD_prices_unweighted" * "_HP" * ".png") # export image

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Figure 2
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# prepare data
df = transform(df_comext, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> ifelse(iszero(x) | ismissing(x), missing, x)), renamecols=false) # need to remove missing also from EU
subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> !ismissing(x)))
subset!(df, :PARTNER_ISO => ByRow(x -> !(x in ["EU", "WORLD"])))

# remove outliers based on medians abs dev.
outlier_cutoff = 3
transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)
gdf = groupby(df, "PRODUCT_NC")
df = transform(gdf, :UNIT_PRICE => MAD_method => :MAD)
subset!(df, :MAD => ByRow(x -> x < outlier_cutoff))

# produce data for plotting
df = data_fig2(df, ["BE"], ["GB"], 2)

# ------------------

# formatting
transform!(df, [:DECLARANT_ISO, :PARTNER_ISO, :FLOW] .=> ByRow(string), renamecols=false)
transform!(df, :PERIOD => ByRow(x -> Date(string(x), DateFormat("yyyymm"))) => :DATE)

df.PRODUCT_NC_digits = parse.(Int64, String.(df.PRODUCT_NC_digits)) # transform into Int64

# HP filter
λ = 20
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PRODUCT_NC_digits"]
gdf = groupby(df, cols_grouping)
transform!(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG, :PRICE_INDEX] .=> (x -> HP(x, λ)) .=> [:VALUE_IN_EUROS_HP, :QUANTITY_IN_KG_HP, :PRICE_INDEX_HP])

# create unweighted unit prices
transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)
transform!(df, [:VALUE_IN_EUROS_HP, :QUANTITY_IN_KG_HP] => ByRow((v,q) -> v/q) => :UNIT_PRICE_HP)

# plotting
# subsetting products for figures
sort!(df, :DATE) # for some reason needed to have correct DATE axis
products = 1:8:89

for flow in ["imports", "exports"]
    for i in products

    product_range = i:i+7
    end_range = string(product_range[end])

    # values
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PRODUCT_NC_digits => ByRow(x -> x in product_range)) plot(:DATE, :VALUE_IN_EUROS,
            group=:PRODUCT_NC_digits, lw=2, legend=:bottomleft, ylabel="euros", title="Belgian "*flow*": values")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig2/values/" * "fig2_" * flow * "_" * end_range * "_values" * ".png") # export image

    # price index (weighted unit prices)
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PRODUCT_NC_digits => ByRow(x -> x in product_range)) plot(:DATE, :PRICE_INDEX,
            group=:PRODUCT_NC_digits, lw=2, legend=:bottomleft, ylabel="unit prices (weighted)", title="Belgian "*flow*": unit prices (weighted)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig2/prices_weighted/" * "fig2_" * flow * "_" * end_range * "_prices_weighted" * ".png")

    # price index (unweighted unit prices)
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PRODUCT_NC_digits => ByRow(x -> x in product_range)) plot(:DATE, :UNIT_PRICE,
            group=:PRODUCT_NC_digits, lw=2, legend=:bottomleft, ylabel="unit prices (unweighted)", title="Belgian "*flow*": unit prices (unweighted)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig2/prices_unweighted/" * "fig2_" * flow * "_" * end_range * "_prices_unweighted" * ".png")

    end
end

for flow in ["imports", "exports"]
    for i in products

    product_range = i:i+7
    end_range = string(product_range[end])

    # values
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PRODUCT_NC_digits => ByRow(x -> x in product_range)) plot(:DATE, :VALUE_IN_EUROS_HP,
            group=:PRODUCT_NC_digits, lw=2, legend=:bottomleft, ylabel="euros", title="Belgian "*flow*": values (HP, λ=$λ)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig2/HP/values/" * "fig2_" * flow * "_" * end_range * "_values" * "_HP" * ".png") # export image

    # price index (weighted unit prices)
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PRODUCT_NC_digits => ByRow(x -> x in product_range)) plot(:DATE, :PRICE_INDEX_HP,
            group=:PRODUCT_NC_digits, lw=2, legend=:bottomleft, ylabel="unit prices (weighted)", title="Belgian "*flow*": unit prices (weighted, λ=$λ)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig2/HP/prices_weighted/" * "fig2_" * flow * "_" * end_range * "_prices_weighted" * "_HP" * ".png")

    # price index (unweighted unit prices)
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PRODUCT_NC_digits => ByRow(x -> x in product_range)) plot(:DATE, :UNIT_PRICE_HP,
            group=:PRODUCT_NC_digits, lw=2, legend=:bottomleft, ylabel="unit prices (unweighted)", title="Belgian "*flow*": unit prices (unweighted, λ=$λ)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig2/HP/prices_unweighted/" * "fig2_" * flow * "_" * end_range * "_prices_unweighted" * "_HP" * ".png")

    end
end

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Figure 3
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# prepare data
df = transform(df_comext, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> ifelse(iszero(x) | ismissing(x), missing, x)), renamecols=false) # need to remove missing also from EU
subset!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> !ismissing(x)))
#subset!(df, :PARTNER_ISO => ByRow(x -> !(x in ["EU", "WORLD"])))

# remove outliers based on medians abs dev.
outlier_cutoff = 3
transform!(df, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] => ByRow((v,q) -> v/q) => :UNIT_PRICE)
gdf = groupby(df, "PRODUCT_NC")
df = transform(gdf, :UNIT_PRICE => MAD_method => :MAD)
subset!(df, :MAD => ByRow(x -> x < outlier_cutoff))

# produce data for plotting
df = data_fig3(df, ["BE"], ["GB", "EU", "DE", "FR", "NL"])

# ------------------

# formatting
transform!(df, [:DECLARANT_ISO, :PARTNER_ISO, :FLOW] .=> ByRow(string), renamecols=false)
transform!(df, :PERIOD => ByRow(x -> Date(string(x), DateFormat("yyyymm"))) => :DATE)

# HP filter
λ = 20
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW"]
gdf = groupby(df, cols_grouping)
transform!(gdf, [:VALUE_SHARE, :QUANTITY_SHARE] .=> (x -> HP(x, λ)) .=> [:VALUE_SHARE_HP, :QUANTITY_SHARE_HP])

# plotting
sort!(df, :DATE)

for flow in ["imports", "exports"]

    # value share
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> !(x == "EU"))) plot(:DATE, :VALUE_SHARE,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="percentages", title="Belgian "*flow*": value share")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig3/" * "fig3_" * flow * "_value_share" * ".png") # export image

    # quantity share
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> !(x == "EU"))) plot(:DATE, :QUANTITY_SHARE,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="kg", title="Belgian "*flow*": quantity share")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig3/" * "fig3_" * flow * "_quantity_share" * ".png") # export image

end

# plotting HP
for flow in ["imports", "exports"]

    # value share
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> !(x == "EU"))) plot(:DATE, :VALUE_SHARE_HP,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="percentages", title="Belgian "*flow*": value share (HP, λ=$λ)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig3/HP/" * "fig3_" * flow * "_value_share" * "_HP" * ".png") # export image

    # quantity share
    p = @df subset(df, :FLOW => ByRow(x -> x == flow), :PARTNER_ISO => ByRow(x -> !(x == "EU"))) plot(:DATE, :QUANTITY_SHARE_HP,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="kg", title="Belgian "*flow*": quantity share (HP, λ=$λ)")
        vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
        vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
        vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/comext/fig3/HP/" * "fig3_" * flow * "_quantity_share" * "_HP" * ".png") # export image

end

# BE vis-a-vis GB
p = @df subset(df, :PARTNER_ISO => ByRow(x -> x == "GB")) plot(:DATE, :VALUE_SHARE_HP,
        group=:FLOW, lw=2, legend=:bottomleft, ylabel="percentages", title="Total trade share (values, λ=$λ): \n Belgium vis-a-vis Great Britain")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
savefig(p, dir_dropbox * "results/images/comext/fig3/HP/" * "fig3_" * "BE_GB" * "_value_share" * "_HP" * ".png") # export image


#TODO make test for this
"""
    Example of 'intake_info' = Dict("gt" => Dict("ub" => 100.0, "c" => 10.0))
"""
function apply_bound!(metnet::MetNet, xi, intake_info::Dict = Dict();
        emptyfirst = false, ignore_miss = false)
    xi < 0 && error("xi must be positive")

    emptyfirst && empty!(metnet.intake_info)
    for (exch, exch_info) in intake_info
        ignore_miss && !(exch in metnet.rxns) && continue
        metnet.intake_info[exch] = exch_info
    end

    for (exch, exch_info) in metnet.intake_info

        !haskey(exch_info, "c") && error("'c' not defined for exchange '$(rxns(metnet, exch))'");

        c = exch_info["c"]
        c < 0.0 && error("negative 'c' ($c) in exchange '$(rxns(metnet, exch))'")

        if haskey(exch_info, "ub")
            ub = exch_info["ub"]
            ub!(metnet, exch, stst_bound(c, xi, ub))
        elseif haskey(exch_info, "lb")
            lb = exch_info["lb"]
            lb!(metnet, exch, stst_bound(c, xi, lb))
        else
            error("neither 'lb' or 'ub' defined for exchange '$(rxns(metnet, exch))'")
        end
    end
    return metnet
end


__precompile__()

module Currencies

# Exports
export AbstractMonetary, Monetary
export currency, decimals, majorunit, @usingcurrencies
export currencyinfo, iso4217num, iso4217alpha, shortsymbol, longsymbol
export valuate, ExchangeRateTable, ecbrates
export Basket, StaticBasket, DynamicBasket
export simplefv, compoundfv
export newcurrency!, @usingcustomcurrency
export format

# DeclarativeFormatting (not specific to Currencies; under development)
include("DeclarativeFormatting/DeclarativeFormatting.jl")

# Currency data
include("CurrencyData/CurrencyData.jl")

# Core features
include("CurrenciesBase/CurrenciesBase.jl")
using .CurrenciesBase

include("Baskets/Basket.jl")
using .Baskets

include("Valuation/Valuation.jl")
using .Valuation

# Interface (display/formatting, convenience macro)
include("CurrencyFormatting/CurrencyFormatting.jl")
using .CurrencyFormatting

# Deprecations
include("deprecated.jl")

end # module


#
# meta.jl --
#
# Functions to generate code for interpolation methods.
#
#------------------------------------------------------------------------------
#
# This file is part of the LinearInterpolators package licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2016-2021, Éric Thiébaut.
#

module Meta

const Arg = Union{Number,Symbol,Expr}

#"""
#```julia
#@getcoefs N j w ker lim pos
#```
#
#yields the code for extracting the interpolation coefficients, `N` is the
#interpolation kernel size, `j` and `w` are prefixes for the variables to store
#the indices and weights for interpolation (the actual variable names are given
#by appending a `_i` suffix for `i` in `1:N`), `ker` is the interpolation
#kernel, `lim` are the limits and `pos` is the position of interpolation.
#
#"""
#macro getcoefs(N, j, w, ker, lim, pos)
#    generate_getcoefs(N, j, w, ker, lim, pos)
#end
#
#macro nonsep_interp(alpha, src, beta, dst,
#                    s1, ker1, lim1, pos1,
#                    s2, ker2, lim2, pos2)
#    return generate_nonsep_interp(alpha, src, beta, dst,
#                          s1, ker1, lim1, pos1,
#                          s2, ker2, lim2, pos2)
#end
#
#macro nonsep_interp_adj(src, dst,
#                        s1, ker1, lim1, pos1,
#                        s2, ker2, lim2, pos2)
#    return generate_nonsep_interp_adj(src, dst,
#                              s1, ker1, lim1, pos1,
#                              s2, ker2, lim2, pos2)
#end

"""
    generate_getcoefs(n, j, w, ker, lim, pos)

generates the code for getting the interpolation coefficients.  Here `n` is the
size of `ker` the interpolation kernel, `j` and `w` are symbols or strings used
as prefixes for local variables to store the interpolation indices and weights
respectively, `ker` and `lim` are symbols with the name of the local variables
which store the interpolation kernel and the limits along the dimension of
interpolation, `pos` is a symbol or an expression which gives the position to
interpolate.  For instance:

    generate_getcoefs(2, :i, :c, :kr, :lm, :(x[i]))
        -> :((i_1, i_2, c_1, c_2) = getcoefs(kr, lm, x[i]))

Another possibility is:

    generate_getcoefs(J, W, ker, lim, pos)

where `J` and `W` are vectors of length `n`, the size of `ker` the
interpolation kernel, with the symbolic names of the local variables to store
the interpolation indices and weights respectively.

"""
function generate_getcoefs(n::Integer,
                           j::Union{AbstractString,Symbol},
                           w::Union{AbstractString,Symbol},
                           ker::Symbol, lim::Symbol, pos::Arg)
    return generate_getcoefs([Symbol(j,"_",s) for s in 1:n],
                             [Symbol(w,"_",s) for s in 1:n], ker, lim, pos)
end

function generate_getcoefs(J::AbstractVector{Symbol},
                           W::AbstractVector{Symbol},
                           ker::Symbol, lim::Symbol, pos::Arg)
    @assert length(J) == length(W)
    vars =  Expr(:tuple, J..., W...)
    return :($vars = getcoefs($ker, $lim, $pos))
end

"""
    generate_sum(ex)

generates an expression whose result is the sum of the elements of the vector
`ex` which are symbols (being interpreted as the name of variables) or
expressions.

"""
generate_sum(args::Arg...) = generate_sum(args)
generate_sum(ex::Union{AbstractVector,Tuple{Vararg{Arg}}}) =
    (length(ex) == 1 ? ex[1] : Expr(:call, :+, ex...))

"""
    group_expressions(ex...) -> code

generates a single expression from all expressions given in argument.  The
result may be the same as the input if it is a single expression or a block of
expressions if several expressions are specified.

To insert the result as a block of statements (like a `begin` ... `end` block)
in a quoted expression, write something like:

    quote
        some_pre_code
        \$code
        some_post_code
    end

To strip the surrounding `begin` ... `end` keywords, write instead:

    quote
        some_pre_code
        \$(code.args...)
        some_post_code
    end

"""
group_expressions(args::Expr...) = group_expressions(args)
group_expressions(ex::Union{AbstractVector{Expr},Tuple{Vararg{Expr}}}) =
    (length(ex) == 1 ? ex[1] : Expr(:block, ex...))

function generate_interp_expr(arr::Symbol,
                              J::AbstractVector{Symbol},
                              W::AbstractVector{Symbol})
    @assert length(J) == length(W)
    return generate_sum([:($arr[$(J[i])]*$(W[i])) for i in 1:length(J)])
end

function generate_interp_expr(arr::Symbol,
                              J::AbstractVector{Symbol},
                              W::AbstractVector{Symbol},
                              inds::Arg)
    @assert length(J) == length(W)
    return generate_sum([:($arr[$(J[i]),$inds]*$(W[i])) for i in 1:length(J)])
end

function generate_interp_expr(arr::Symbol,
                              J1::AbstractVector{Symbol},
                              W1::AbstractVector{Symbol},
                              J2::AbstractVector{Symbol},
                              W2::AbstractVector{Symbol})
    @assert length(J2) == length(W2)
    n = length(J2)
    ex = Array{Expr}(undef, n)
    for i in 1:n
        sub = generate_interp_expr(arr, J1, W1, J2[i])
        ex[i] = :($sub*$(W2[i]))
    end
    return generate_sum(ex)
end

end # module


abstract type AbstractLShapedSolver <: AbstractStructuredModel end

nscenarios(lshaped::AbstractLShapedSolver) = lshaped.nscenarios
ncuts(lshaped::AbstractLShapedSolver) = lshaped.data.ncuts
niterations(lshaped::AbstractLShapedSolver) = lshaped.data.iterations
tolerance(lshaped::AbstractLShapedSolver) = lshaped.parameters.τ

# Initialization #
# ======================================================================== #
function initialize!(lshaped::AbstractLShapedSolver, subsolver::SubSolver)
    # Initialize progress meter
    lshaped.progress.thresh = lshaped.parameters.τ
    # Initialize subproblems
    initialize_subproblems!(lshaped, scenarioproblems(lshaped.stochasticprogram), decision(lshaped), subsolver)
    # Prepare the master optimization problem
    prepare_master!(lshaped)
    # Initialize feasibility handler
    initialize_feasibility!(lshaped)
    # Initialize regularization policy
    initialize_regularization!(lshaped)
    # Finish initialization
    finish_initilization!(lshaped)
    return nothing
end
# ======================================================================== #

# Functions #
# ======================================================================== #
function set_params!(lshaped::AbstractLShapedSolver; kwargs...)
    for (k,v) in kwargs
        setfield!(lshaped.parameters,k,v)
    end
    return nothing
end

function prepare_master!(lshaped::AbstractLShapedSolver)
    # θs
    for i = 1:nthetas(lshaped)
        MPB.addvar!(lshaped.mastersolver.lqmodel ,-1e10, Inf, initial_theta_coefficient(lshaped.feasibility))
        push!(lshaped.mastervector, -1e10)
    end
    return nothing
end

function update_solution!(lshaped::AbstractLShapedSolver, solver::LQSolver)
    ncols = decision_length(lshaped.stochasticprogram)
    nb = nthetas(lshaped)
    x = copy(getsolution(solver))
    lshaped.mastervector .= x[1:ncols+nb]
    lshaped.x[1:ncols] .= x[1:ncols]
    set_model_objectives(lshaped, x[ncols+1:ncols+nb])
    return nothing
end

function get_obj(lshaped::AbstractLShapedSolver)
    first_stage = StochasticPrograms.get_stage_one(lshaped.stochasticprogram)
    c = JuMP.prepAffObjective(first_stage)
    return first_stage.objSense == :Min ? c : -c
end

function current_objective_value(lshaped::AbstractLShapedSolver)
    return get_obj(lshaped)⋅current_decision(lshaped) + sum(subobjectives(lshaped))
end

function calculate_estimate(lshaped::AbstractLShapedSolver)
    return get_obj(lshaped)⋅lshaped.x + sum(model_objectives(lshaped))
end

function log!(lshaped::AbstractLShapedSolver; optimal = false)
    @unpack Q,θ,iterations = lshaped.data
    @unpack keep, offset, indent = lshaped.parameters
    push!(lshaped.Q_history, Q)
    push!(lshaped.θ_history, θ)
    lshaped.data.iterations += 1

    log_regularization!(lshaped)
    if lshaped.parameters.log
        current_gap = optimal ? 0.0 : gap(lshaped)
        ProgressMeter.update!(lshaped.progress,current_gap,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", objective(lshaped)),
                                  ("$(indentstr(indent))Gap", current_gap),
                                  ("$(indentstr(indent))Number of cuts", ncuts(lshaped)),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
    end
    return nothing
end

function log!(lshaped::AbstractLShapedSolver, t::Integer; optimal = false)
    @unpack Q,θ,iterations = lshaped.data
    @unpack keep, offset, indent = lshaped.parameters
    lshaped.Q_history[t] = Q
    lshaped.θ_history[t] = θ

    log_regularization!(lshaped,t)

    if lshaped.parameters.log
        current_gap = optimal ? 0.0 : gap(lshaped)
        ProgressMeter.update!(lshaped.progress, current_gap,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", objective(lshaped)),
                                  ("$(indentstr(indent))Gap", current_gap),
                                  ("$(indentstr(indent))Number of cuts", ncuts(lshaped)),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
    end
    return nothing
end

function indentstr(n::Integer)
    return repeat(" ", n)
end

function check_optimality(lshaped::AbstractLShapedSolver)
    @unpack τ = lshaped.parameters
    @unpack θ = lshaped.data
    return θ > -Inf && gap(lshaped) <= τ
end
# ======================================================================== #

# Cut functions #
# ======================================================================== #
active(lshaped::AbstractLShapedSolver, hyperplane::AbstractHyperPlane) = active(hyperplane, decision(lshaped), tolerance(lshaped))
active(lshaped::AbstractLShapedSolver, cut::HyperPlane{OptimalityCut}) = optimal(cut, decision(lshaped), model_objectives(lshaped)[cut.id], tolerance(lshaped))
active(lshaped::AbstractLShapedSolver, cut::AggregatedOptimalityCut) = optimal(cut, decision(lshaped), sum(model_objectives(lshaped)[cut.ids]), tolerance(lshaped))
satisfied(lshaped::AbstractLShapedSolver, hyperplane::AbstractHyperPlane) = satisfied(hyperplane, decision(lshaped), tolerance(lshaped))
satisfied(lshaped::AbstractLShapedSolver, cut::HyperPlane{OptimalityCut}) = satisfied(cut, decision(lshaped), model_objectives(lshaped)[cut.id], tolerance(lshaped))
satisfied(lshaped::AbstractLShapedSolver, cut::AggregatedOptimalityCut) = satisfied(cut, decision(lshaped), sum(model_objectives(lshaped)[cut.ids]), tolerance(lshaped))
violated(lshaped::AbstractLShapedSolver, hyperplane::AbstractHyperPlane) = !satisfied(lshaped, hyperplane)
gap(lshaped::AbstractLShapedSolver, hyperplane::AbstractHyperPlane) = gap(hyperplane, decision(lshaped))
gap(lshaped::AbstractLShapedSolver, cut::HyperPlane{OptimalityCut}) = gap(cut, decision(lshaped), model_objectives(lshaped)[cut.id])
gap(lshaped::AbstractLShapedSolver, cut::AggregatedOptimalityCut) = gap(cut, decision(lshaped), sum(model_objectives(lshaped)[cut.ids]))

function add_cut!(lshaped::AbstractLShapedSolver, cut::AbstractHyperPlane; consider_consolidation = true, check = true)
    added = add_cut!(lshaped, cut, model_objectives(lshaped), subobjectives(lshaped), check = check)
    added && update_objective!(lshaped, cut)
    if consider_consolidation
        added && add_cut!(lshaped, lshaped.consolidation, cut)
    end
    return added
end

function add_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane{OptimalityCut}, θs::AbstractVector, subobjectives::AbstractVector, Q::AbstractFloat; check = true)
    θ = θs[cut.id]
    @unpack τ, cut_scaling = lshaped.parameters
    # Update objective
    subobjectives[cut.id] = Q
    # Check if cut gives new information
    if check && θ > -Inf && (θ + τ >= Q || θ + τ >= cut(lshaped.x))
        # Optimal with respect to this subproblem
        return false
    end
    # Add optimality cut
    process_cut!(lshaped, cut)
    MPB.addconstr!(lshaped.mastersolver.lqmodel, lowlevel(cut, cut_scaling)...)
    lshaped.data.ncuts += 1
    if lshaped.parameters.debug
        push!(lshaped.cuts, cut)
    end
    return true
end
function add_cut!(lshaped::AbstractLShapedSolver, cut::AggregatedOptimalityCut, θs::AbstractVector, subobjectives::AbstractVector, Q::AbstractFloat; check = true)
    θs = θs[cut.ids]
    θ = sum(θs)
    @unpack τ, cut_scaling = lshaped.parameters
    # Update objective
    subobjectives[cut.ids] .= Q/length(cut.ids)
    # Check if cut gives new information
    if check && θ > -Inf && (θ + τ >= Q || θ + τ >= cut(lshaped.x))
        # Optimal with respect to these subproblems
        return false
    end
    # Add optimality cut
    process_cut!(lshaped, cut)
    MPB.addconstr!(lshaped.mastersolver.lqmodel, lowlevel(cut, cut_scaling)...)
    lshaped.data.ncuts += 1
    if lshaped.parameters.debug
        push!(lshaped.cuts, cut)
    end
    return true
end
add_cut!(lshaped::AbstractLShapedSolver, cut::AnyOptimalityCut, θs::AbstractVector, subobjectives::AbstractVector, x::AbstractVector; check = true) = add_cut!(lshaped, cut, θs, subobjectives, cut(x); check = check)
add_cut!(lshaped::AbstractLShapedSolver, cut::AnyOptimalityCut, θs::AbstractVector, subobjectives::AbstractVector; check = true) = add_cut!(lshaped, cut, θs, subobjectives, current_decision(lshaped); check = check)

function add_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane{FeasibilityCut}, ::AbstractVector, subobjectives::AbstractVector, Q::AbstractFloat; check = true)
    # Ensure that there is no false convergence
    subobjectives[cut.id] = Q
    # Add feasibility cut
    process_cut!(lshaped, cut)
    MPB.addconstr!(lshaped.mastersolver.lqmodel, lowlevel(cut)...)
    lshaped.data.ncuts += 1
    if lshaped.parameters.debug
        push!(lshaped.feasibility.cuts, cut)
    end
    return true
end
add_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane{FeasibilityCut}, θs::AbstractVector, subobjectives::AbstractVector; check = true) = add_cut!(lshaped, cut, θs, subobjectives, Inf)

function add_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane{Infeasible}, ::AbstractVector, subobjectives::AbstractVector; check = true)
    subobjectives[cut.id] = Inf
    return true
end

function add_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane{Unbounded}, ::AbstractVector, subobjectives::AbstractVector; check = true)
    subobjectives[cut.id] = -Inf
    return true
end


using Oceananigans.Fields: AbstractField, compute!

fetch_output(output, model, field_slicer) = output(model)

# Needed to support `fetch_output` with `model::Nothing`.
time(model) = model.clock.time
time(::Nothing) = nothing

function fetch_output(field::AbstractField, model, field_slicer)
    compute!(field, time(model))
    return slice_parent(field_slicer, field)
end

convert_output(output, writer) = output
convert_output(output::AbstractArray, writer) = writer.array_type(output)

fetch_and_convert_output(output, model, writer) =
    convert_output(fetch_output(output, model, writer.field_slicer), writer)


# Copyright 2019 Martin Holters
# See accompanying license file.

function wrap_nleq_expr(nn, nq, expr)
    Base.depwarn("nonlinear_eq should be given as a function, not an expression", :Element)

    res_symbs = [gensym("res_$n") for n in 1:nn]
    J_symbs = [gensym("J_$(m)_$n") for m in 1:nn, n in 1:nq]

    function rewrite_refs(expr::Expr)
        if expr.head == :ref
            if expr.args[1] === :res
                return res_symbs[expr.args[2]]
            elseif expr.args[1] === :J
                return J_symbs[expr.args[2], expr.args[3]]
            end
            return expr
        else
            return Expr(expr.head, rewrite_refs.(expr.args)...)
        end
    end

    rewrite_refs(x::Any) = x

    expr = rewrite_refs(expr)

    return eval(quote
        @inline function (q)
            $(nn > 0 || nq > 0 ? Expr(:local, res_symbs..., J_symbs...) : nothing)
            $(expr)
            res = SVector{$nn}($(res_symbs...))
            J = SMatrix{$nn,$nq}($(J_symbs...))
            return (res, J)
        end
    end)
end

function ports_from_old_pins(pins)
    ports = [pins[2i-1] => pins[2i] for i in 1:length(pins)÷2]
    Base.depwarn("`Element(..., pins=$(repr(pins)))` is depreated, use `Element(..., ports=$(repr(ports, context=:typeinfo=>typeof(ports))))` instead", :Element)
    return ports
end

Base.@deprecate(composite_element(circ::Circuit, pins::Vector{<:Pair}),
    composite_element(circ,  pinmap=Dict(pins...)))


#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2019
#
# This is a Julia package for discrete event simulation
#
using DiscreteEvents, .Threads

sleeptime = 0.01
println("... testing real time clock ...")

rtc = createRTClock(0.01, 4711)
stopRTClock(rtc)
rtc.time = 4711.0
thrd = nthreads()
str = "RTClock 4711 on thread $thrd: state=:idle, t=4711.0 s, T=0.01 s, prc:0\n   scheduled ev:0, cev:0, sampl:0\n"
@test repr(rtc) == str

rtc = createRTClock(0.01, 4711)
sleep(1.01)
@test tau(rtc) ≥ 1
t = tau(rtc)
resetClock!(rtc)
sleep(sleeptime)
@test tau(rtc) < t

put!(rtc.cmd, DiscreteEvents.Clear())
sleep(sleeptime)
put!(rtc.cmd, DiscreteEvents.Diag())
err = take!(rtc.back).x
@test err[1] isa ErrorException

a = [0]
b = [0]
incr!(x) = x[1] += 1
event!(rtc, fun(incr!, a), every, 0.1)
periodic!(rtc, fun(incr!, b))
sleep(1)
@test a[1] ≥ 10
@test b[1] ≥ 100

old = DiscreteEvents._handle_exceptions[end]
DiscreteEvents._handle_exceptions[end] = false
sleep(1)
@test a[1] ≥ 20
@test b[1] ≥ 200
DiscreteEvents._handle_exceptions[end] = old

stopRTClock(rtc)
t = tau(rtc)
sleep(0.1)
@test t == tau(rtc)


module PortNumbers

import Base: string, convert

export PortClass, EstablishedPort, RootPort, UserPort, AnyPort, PortNumber
export portmin, portmax
export PortNumber
export string, convert
export make_port_number

const ESTABLISHED_MIN_DEFAULT = 1
const ESTABLISHED_MAX_DEFAULT = 256
const ROOT_MIN_DEFAULT = 1
const ROOT_MAX_DEFAULT = 1023
const USER_MIN_DEFAULT = 1024
const USER_MAX_DEFAULT = 65535
const ANY_MIN_DEFAULT = 1
const ANY_MAX_DEFAULT = 65535

const ESTABLISHED_MIN_DEFAULT = 1
const ESTABLISHED_MAX_DEFAULT = 256
const ROOT_MIN_DEFAULT = 1
const ROOT_MAX_DEFAULT = 1023
const USER_MIN_DEFAULT = 1024
const USER_MAX_DEFAULT = 65535
const ANY_MIM_DEFAULT = 1
const ANY_MAX_DEFAULT = 65535

est_min = ESTABLISHED_MIN_DEFAULT
est_max = ESTABLISHED_MAX_DEFAULT

root_min = ROOT_MIN_DEFAULT
root_max = ROOT_MAX_DEFAULT

user_min = USER_MIN_DEFAULT
user_max = USER_MAX_DEFAULT

any_min = ANY_MIN_DEFAULT
any_max = ANY_MAX_DEFAULT

"""
### PortClass

An abstract type used as the super type
 for the various port number classes.
"""
abstract PortClass

abstract AnyPort <: PortClass
minport( ::Type{AnyPort} ) = any_min
maxport( ::Type{AnyPort} ) = any_max

abstract EstablishedPort <: PortClass
minport( ::Type{EstablishedPort} ) = est_min
maxport( ::Type{EstablishedPort} ) = est_max

abstract RootPort <: PortClass
minport( ::Type{RootPort} ) = root_min
maxport( ::Type{RootPort} ) = root_max

abstract UserPort <: PortClass
minport( ::Type{UserPort} ) = user_min
maxport( ::Type{UserPort} ) = user_max

immutable PortNumber{T<:PortClass}
    value::Integer
    function PortNumber( x::Integer )
        @assert x >= minport( T )
        @assert x <= maxport( T )
        new( x )
    end
end

convert{T<:Integer}( ::Type{T}, x::PortNumber ) = T( x.value )
string( x::PortNumber ) = Int( x ) |> string
convert{T <: AbstractString }( ::Type{T}, x::PortNumber ) = string( x ) |> T

"""
### `make_port_number( x::Integer )`

Return a port number object.

#### Example

`*julia> make_port_number( 22 )`

  --> `PortNumber{EstablishedPort}( 22 )`

"""
function make_port_number( x::Integer )
    @assert x >= minport( AnyPort )
    @assert x <= maxport( AnyPort )
    if x <= maxport( EstablishedPort )
        return PortNumber{EstablishedPort}( x )
    elseif x <= maxport( RootPort )
        return PortNumber{RootPort}( x )
    else
        return PortNumber{UserPort}( x )
    end
end

end # module