"""
Collection of data models used for content parsing
"""
module Models

using Base: @kwdef
using Dates

export AbstractStonksRecord,
  AssetInfo,
  AssetPrice,
  ExchangeRate,
  IncomeStatement,
  BalanceSheet,
  CashflowStatement,
  Earnings

"""
Abstract type inteded to be subtyped by an data model.

All subtypes are constructed from keyword args and optional values default to `missing`.
"""
abstract type AbstractStonksRecord end

"""
Stores general information about a quoted symbol / ticker.

### Constructors
```julia
AssetInfo(;
  symbol::String,
  currency::String,
  name::Union{String,Missing}=missing,
  type::Union{String,Missing}=missing,
  exchange::Union{String,Missing}=missing,
  country::Union{String,Missing}=missing,
  industry::Union{String,Missing}=missing,
  sector::Union{String,Missing}=missing,
  timezone::Union{String,Missing}=missing,
  employees::Union{Int,Missing}=missing,
)
```
"""
@kwdef struct AssetInfo <: AbstractStonksRecord
  symbol::String
  currency::String
  name::Union{String,Missing} = missing
  type::Union{String,Missing} = missing
  exchange::Union{String,Missing} = missing
  country::Union{String,Missing} = missing
  industry::Union{String,Missing} = missing
  sector::Union{String,Missing} = missing
  timezone::Union{String,Missing} = missing
  employees::Union{Int,Missing} = missing
end
"""
Stores a time series datapoint with price information. Lowest frequency is daily.

### Constructors
```julia
AssetPrice(
  symbol::String,
  date::Date,
  close::Float64,
  open::Union{Float64,Missing} = missing,
  high::Union{Float64,Missing} = missing,
  low::Union{Float64,Missing} = missing,
  close_adjusted::Union{Float64,Missing} = missing,
  volume::Union{Integer,Missing} = missing,
)
end
```
"""
@kwdef struct AssetPrice <: AbstractStonksRecord
  symbol::String
  date::Date
  close::Float64
  open::Union{Float64,Missing} = missing
  high::Union{Float64,Missing} = missing
  low::Union{Float64,Missing} = missing
  close_adjusted::Union{Float64,Missing} = missing
  volume::Union{Integer,Missing} = missing
end

"""
Stores an exchange rate datapoint. Lowest frequency is daily.

### Constructors
```julia
ExchangeRate(;
  base::String,
  target::String,
  date::Date,
  rate::Float64,
)
```
"""
@kwdef struct ExchangeRate <: AbstractStonksRecord
  base::String
  target::String
  date::Date
  rate::Float64
end

"""
Stores a datapoint containing income statement information.
Follows normalized fields mapped to GAAP and IFRS taxonomies of the SEC.

### Constructors
```julia
IncomeStatement(;
  symbol::String,
  frequency::String,
  date::Date,
  currency::Union{String,Missing} = missing,
  total_revenue::Int64,
  cost_of_revenue::Int64,
  gross_profit::Int64,
  operating_income::Int64,
  selling_general_and_administrative::Union{Int64,Missing} = missing,
  research_and_development::Union{Int64,Missing} = missing,
  depreciation::Union{Int64,Missing} = missing, # Depreciation, Total
  depreciation_and_amortization::Union{Int64,Missing} = missing,
  income_before_tax::Union{Int64,Missing} = missing,
  income_tax_expense::Union{Int64,Missing} = missing,
  interest_expense::Union{Int64,Missing} = missing,
  interest_and_debt_expense::Union{Int64,Missing} = missing,
  ebit::Union{Int64,Missing} = missing,
  ebitda::Union{Int64,Missing} = missing,
  net_income::Int64,
  net_income_common_shares::Union{Int64,Missing} = missing,
)
```
"""
@kwdef struct IncomeStatement <: AbstractStonksRecord
  symbol::String
  frequency::String
  date::Date
  currency::Union{String,Missing} = missing
  total_revenue::Int64
  cost_of_revenue::Int64
  gross_profit::Int64
  operating_income::Int64
  selling_general_and_administrative::Union{Int64,Missing} = missing
  research_and_development::Union{Int64,Missing} = missing
  depreciation::Union{Int64,Missing} = missing # Depreciation, Total
  depreciation_and_amortization::Union{Int64,Missing} = missing
  income_before_tax::Union{Int64,Missing} = missing
  income_tax_expense::Union{Int64,Missing} = missing
  interest_expense::Union{Int64,Missing} = missing
  interest_and_debt_expense::Union{Int64,Missing} = missing
  #comprehensive_income_net_of_tax::Union{Int64,Missing} = missing
  ebit::Union{Int64,Missing} = missing
  ebitda::Union{Int64,Missing} = missing
  net_income::Int64
  net_income_common_shares::Union{Int64,Missing} = missing
end

"""
Stores a datapoint containing balance sheet information.
Follows normalized fields mapped to GAAP and IFRS taxonomies of the SEC.

### Constructors
```julia
BalanceSheet(;
  symbol::String,
  frequency::String,
  date::Date,
  currency::Union{String,Missing} = missing,
  total_assets::Int64,
  total_liabilities::Int64,
  total_shareholder_equity::Int64,
  cash_and_equivalents::Union{Int64,Missing} = missing,
  current_net_receivables::Union{Int64,Missing} = missing,
  inventory::Union{Int64,Missing} = missing,
  short_term_investments::Union{Int64,Missing} = missing,
  other_current_assets::Union{Int64,Missing} = missing,
  total_current_assets::Union{Int64,Missing} = missing,
  property_plant_equipment::Union{Int64,Missing} = missing,
  goodwill::Union{Int64,Missing} = missing,
  long_term_investments::Union{Int64,Missing} = missing,
  intangible_assets::Union{Int64,Missing} = missing,
  total_noncurrent_assets::Union{Int64,Missing} = missing,
  current_accounts_payable::Union{Int64,Missing} = missing,
  deferred_revenue::Union{Int64,Missing} = missing,
  short_term_debt::Union{Int64,Missing} = missing,
  other_current_liabilities::Union{Int64,Missing} = missing,
  total_current_liabilities::Union{Int64,Missing} = missing,
  current_debt::Union{Int64,Missing} = missing,
  current_long_term_debt::Union{Int64,Missing} = missing,
  long_term_debt::Union{Int64,Missing} = missing,
  long_term_debt_noncurrent::Union{Int64,Missing} = missing,
  capital_lease_obligations::Union{Int64,Missing} = missing,
  other_noncurrent_liabilities::Union{Int64,Missing} = missing,
  total_noncurrent_liabilities::Union{Int64,Missing} = missing,
  treasury_stock::Union{Int64,Missing} = missing,
  retained_earnings::Union{Int64,Missing} = missing,
  common_stock::Union{Int64,Missing} = missing,
  common_stock_shares_outstanding::Union{Int64,Missing} = missing,
)
```
"""
@kwdef struct BalanceSheet <: AbstractStonksRecord
  symbol::String
  frequency::String
  date::Date
  currency::Union{String,Missing} = missing
  total_assets::Int64
  total_liabilities::Int64
  total_shareholder_equity::Int64
  cash_and_equivalents::Union{Int64,Missing} = missing
  current_net_receivables::Union{Int64,Missing} = missing
  inventory::Union{Int64,Missing} = missing
  short_term_investments::Union{Int64,Missing} = missing
  other_current_assets::Union{Int64,Missing} = missing
  total_current_assets::Union{Int64,Missing} = missing
  property_plant_equipment::Union{Int64,Missing} = missing
  goodwill::Union{Int64,Missing} = missing
  long_term_investments::Union{Int64,Missing} = missing
  intangible_assets::Union{Int64,Missing} = missing
  total_noncurrent_assets::Union{Int64,Missing} = missing
  current_accounts_payable::Union{Int64,Missing} = missing
  deferred_revenue::Union{Int64,Missing} = missing
  short_term_debt::Union{Int64,Missing} = missing
  other_current_liabilities::Union{Int64,Missing} = missing
  total_current_liabilities::Union{Int64,Missing} = missing
  current_debt::Union{Int64,Missing} = missing
  current_long_term_debt::Union{Int64,Missing} = missing
  long_term_debt::Union{Int64,Missing} = missing
  long_term_debt_noncurrent::Union{Int64,Missing} = missing
  capital_lease_obligations::Union{Int64,Missing} = missing
  other_noncurrent_liabilities::Union{Int64,Missing} = missing
  total_noncurrent_liabilities::Union{Int64,Missing} = missing
  treasury_stock::Union{Int64,Missing} = missing
  retained_earnings::Union{Int64,Missing} = missing
  common_stock::Union{Int64,Missing} = missing
  common_stock_shares_outstanding::Union{Int64,Missing} = missing
end

"""
Stores a datapoint containing cashflow statement information.

### Constructors
```julia
CashflowStatement(;
  symbol::String,
  frequency::String,
  date::String,
  currency::String,
  operating_cashflow::Union{Int64,Missing} = missing,
  cashflow_investment::Union{Int64,Missing} = missing,
  cashflow_financing::Union{Int64,Missing} = missing,
  change_operating_liabilities::Union{Int64,Missing} = missing,
  change_receivables::Union{Int64,Missing} = missing,
  change_inventory::Union{Int64,Missing} = missing,
  change_cash_and_equivalents::Union{Int64,Missing} = missing,
  depreciation_and_amortization::Union{Int64,Missing} = missing,
  capital_expenditures::Union{Int64,Missing} = missing,
  dividend_payout::Union{Int64,Missing} = missing,
  stock_repurchase::Union{Int64,Missing} = missing,
  net_income::Union{Int64,Missing} = missing,
)
```
"""
@kwdef struct CashflowStatement <: AbstractStonksRecord
  symbol::String
  frequency::String
  date::Date
  currency::String
  operating_cashflow::Union{Int64,Missing} = missing
  cashflow_investment::Union{Int64,Missing} = missing
  cashflow_financing::Union{Int64,Missing} = missing
  change_operating_liabilities::Union{Int64,Missing} = missing
  change_receivables::Union{Int64,Missing} = missing
  change_inventory::Union{Int64,Missing} = missing
  change_cash_and_equivalents::Union{Int64,Missing} = missing
  depreciation_and_amortization::Union{Int64,Missing} = missing
  capital_expenditures::Union{Int64,Missing} = missing
  dividend_payout::Union{Int64,Missing} = missing
  stock_repurchase::Union{Int64,Missing} = missing
  net_income::Union{Int64,Missing} = missing
end

"""
Stores a datapoint containing earnings (per share) information.
### Constructors
```julia
Earnings(;
  symbol::String,
  frequency::String,
  date::Date,
  currency::Union{String,Missing} = missing,
  actual::Float16,
  estimate::Union{Float16,Missing} = missing,
)
```
"""
@kwdef struct Earnings <: AbstractStonksRecord
  symbol::String
  frequency::String
  date::Date
  currency::Union{String,Missing} = missing
  actual::Float16
  estimate::Union{Float16,Missing} = missing
end

# @kwdef struct EconomicIndicator <: AbstractStonksRecord
#   symbol::String
#   name::String
#   date::Date
#   value::Number
# end

end


# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

function _offer(filter::_ViolationFilter, v::_Violation)::Nothing
    if v.monitored_line.offset ∉ keys(filter.queues)
        filter.queues[v.monitored_line.offset] =
            PriorityQueue{_Violation,Float64}()
    end
    q::PriorityQueue{_Violation,Float64} =
        filter.queues[v.monitored_line.offset]
    if length(q) < filter.max_per_line
        enqueue!(q, v => v.amount)
    else
        if v.amount > peek(q)[1].amount
            dequeue!(q)
            enqueue!(q, v => v.amount)
        end
    end
    return nothing
end

function _query(filter::_ViolationFilter)::Array{_Violation,1}
    violations = Array{_Violation,1}()
    time_queue = PriorityQueue{_Violation,Float64}()
    for l in keys(filter.queues)
        line_queue = filter.queues[l]
        while length(line_queue) > 0
            v = dequeue!(line_queue)
            if length(time_queue) < filter.max_total
                enqueue!(time_queue, v => v.amount)
            else
                if v.amount > peek(time_queue)[1].amount
                    dequeue!(time_queue)
                    enqueue!(time_queue, v => v.amount)
                end
            end
        end
    end
    while length(time_queue) > 0
        violations = [violations; dequeue!(time_queue)]
    end
    return violations
end


"""
    jgraph_morph(g::JGraph, layout_to, lims_to; frames=nothing)

This functions allows to animate the change from one layout to another for a graph.

# Arguments
- `g::JGraph` the `JGraph` to animate.
- `layout_to` a `NetworkLayout.AbstractLayout` or a `Vector{NetworkLayout.AbstractLayout}` that will be applied to nodes and edges
to determine the configuration to reach.
- `lims_to` a tuple or a vector of tuples each holding width and height to scale the positions of nodes and edges in the animation.

# Keywords
- `frames` a range of frames during which the morph will take place.
If not specified the animation will last as long aas the JGraph itself.
- `starting_positions` specifies the starting positions of each node, defaults
to [`get_starting_positions(g)`](@ref)
"""
function jgraph_morph(
    g::JGraph,
    layout_to::Vector,
    lims_to::Vector;
    frames = nothing,
    starting_positions = get_starting_positions(g),
    closed=false,
)

    frames = isnothing(frames) ? g.data.frames : frames
    if closed
        push!(layout_to, g.data.layout)
        push!(lims_to[1], g.data.width)
        push!(lims_to[2], g.data.height)
    end

    frames_stops =
        floor.(Int, range(first(frames) - 1, last(frames), length = length(layout_to) + 1))

    for i in 1:length(layout_to)
        starting_positions = jgraph_morph(
            g,
            layout_to[i],
            lims_to[i],
            frames = (frames_stops[i] + 1):frames_stops[i + 1],
            starting_positions = starting_positions,
        )
    end
    return starting_positions
end

function jgraph_morph(
    g::JGraph,
    layout_to,
    lims_to;
    frames = nothing,
    starting_positions = get_starting_positions(g),
)
    Jnodes = jnodes(g)
    frames = isnothing(frames) ? g.data.frames : frames
    positions = _adapt_to_lims(layout_to(g.data.graph), lims_to...)
    for (node, pos_from, pos_to) in zip(Jnodes, starting_positions, positions)
        act!(node, Action(frames, anim_translate(pos_from, pos_to)))
    end
    return positions
end


"""
    check_filters(node; scale = nothing, symbol = nothing, link = nothing)

Check if the filters are consistant with the mtg onto which they are applied

# Examples

```julia
check_filters(mtg, scale = 1)
check_filters(mtg, scale = (1,2))
check_filters(mtg, scale = (1,2), symbol = "Leaf", link = "<")
```
"""
function check_filters(node; scale = nothing, symbol = nothing, link = nothing) where {T}

    root_node = get_root(node)

    nodeMTG_type = typeof(node.MTG)
    if root_node[:scales] !== nothing
        check_filter(nodeMTG_type, :scale, scale, unique(root_node.attributes[:scales]))
    end

    if root_node[:symbols] !== nothing
        check_filter(nodeMTG_type, :symbol, symbol, unique(root_node.attributes[:symbols]))
    end

    if root_node[:link] !== nothing
        check_filter(nodeMTG_type, :link, link, ("/", "<", "+"))
    end

    return nothing
end

function check_filter(nodetype, type::Symbol, filter, filters)
    if !isnothing(filter)
        filter_type = fieldtype(nodetype, type)
        !(typeof(filter) <: filter_type) &&
            @warn "The $type argument should be of type $filter_type"
        if !(filter in filters)
            @warn "The $type argument should be one of: $filters, and you provided $filter."
        end
    end
end

function check_filter(nodetype, type::Symbol, filter::T, filters) where {T<:Union{Tuple,AbstractArray}}
    for i in filter
        check_filter(nodetype, type, i, filters)
    end
end

"""
    is_filtered(node, scale, symbol, link, filter_fun)

Is a node filtered in ? Returns `true` if the node is kept, `false` if it is filtered-out.
"""
function is_filtered(node, scale, symbol, link, filter_fun)

    link_keep = is_filtered(link, node.MTG.link)
    symbol_keep = is_filtered(symbol, node.MTG.symbol)
    scale_keep = is_filtered(scale, node.MTG.scale)
    filter_fun_keep = isnothing(filter_fun) || filter_fun(node)

    scale_keep && symbol_keep && link_keep && filter_fun_keep
end

function is_filtered(filter, value)
    isnothing(filter) || value in filter
end

function is_filtered(filter::String, value)
    isnothing(filter) || value in (filter,)
end

function is_filtered(filter, value::T) where {T<:Union{Tuple,Array}}
    all(map(x -> is_filtered(filter, x), value))
end

"""
    parse_macro_args(args)

Parse filters and arguments given as a collection of expressions. This function is used to
get the filters as keyword arguments in macros.

# Examples

```julia
args = (:(x = length(node.name)), :(y = node.x + 2), :(scale = 2))
MultiScaleTreeGraph.parse_macro_args(args)
```
"""
function parse_macro_args(args)
    filters = Dict{Symbol,Any}(
        :scale => nothing,
        :symbol => nothing,
        :link => nothing,
        :filter_fun => nothing
    )

    kwargs = Dict{Symbol,Any}(
        :all => true,
        :traversal => AbstractTrees.PreOrderDFS,
    )

    args_array = []

    for i in args
        if i.head == :(=) && i.args[1] in keys(filters)
            filters[i.args[1]] = i.args[2]
        elseif i.head == :(=) && i.args[1] in keys(kwargs)
            kwargs[i.args[1]] = i.args[2]
        else
            push!(args_array, i)
        end
    end
    return (; filters...), (; kwargs...), (args_array...,)
end


remaps(mapping, x::Symbol) = get(mapping, x, x)
remaps(mapping, x::Expr) = Expr(x.head, remaps(mapping, x.args[1]), x.args[2:end]...)
remaps(mapping, x::Number) = x

type Op
    name::Symbol
    inputs::Vector
    outputs::Vector
    body::Vector
    info::Expr

    function Op(name, inputs, outputs, body, info, mapping = Dict{Symbol, Symbol}())
        if !isempty(body)
            remap(x) = remaps(mapping, x)
            map!(remap, inputs, inputs)
            uniques = Set{Symbol}()
            for op in body
                map!(remap, op.inputs, op.inputs)
                union!(uniques, filter(o->isa(o, Symbol), op.inputs))
                if isdefined(Symbol("δ$(op.name)_const")) || all(isconst, op.inputs)
                    [mapping[o] = Symbol("$(o)_const") for o in filter(isvar, op.outputs)]
                end
                [mapping[o] = gensym(o) for o in filter(o->o in inputs, op.outputs)]
                map!(remap, op.outputs, op.outputs)
                op.name = name_const(op.name, op.inputs, op.outputs)
            end
            map!(remap, outputs, outputs)
            name = name_const(name, inputs, outputs)
        end
        new(name, inputs, outputs, body, info)
    end
end

function name_const(name, inputs, outputs)
    if isdefined(Symbol("δ$(name)_const")) || all(isconst, inputs) || all(isvar, inputs) || all(isconst, outputs)
        return name
    end
    n = string(name)
    for k = eachindex(inputs)
        if !isvar(inputs[k])
            n *= "_const$k"
        end
    end
    Symbol(n)
end

function parse_function(expr, info, mapping = Dict{Symbol, Symbol}())
    @assert (expr.head == :function || expr.head == :(=)) && length(expr.args) == 2  "Only functions can be differentiated"
    header = expr.args[1]
    @assert header.head == :call "Only functions can be differentiated"
    name = header.args[1]
    inputs = header.args[2:end]
    body = expr.args[2]
    @assert body.head == :block "Body of the function is not found"

    local outputs
    ops = []
    [(info, outputs) = parse_line!(ops, info, line) for line in body.args]
    Op(name, inputs, outputs, ops, info, mapping)
end

parse_line!(ops, info, line::LineNumberNode) = info, [line]
parse_line!(ops, info, line::Symbol) = info, [line]
function parse_line!(ops, info, line::Expr)
    outputs = []
    if line.head == :(=)
        outputs = parse_assign!(ops, info, line.args...)
    elseif line.head == :call || line.head == :(.)
        outputs = [parse_arg!(ops, info, line)]
    elseif line.head == :return && isa(line.args[1], SlotNumber)
        outputs = [Symbol(string(line.args[1]))]
    elseif line.head == :return && isa(line.args[1], Symbol)
        outputs = [line.args[1]]
    elseif line.head == :tuple || line.head == :return && line.args[1].head != :tuple
        outputs = [parse_arg!(ops, info, arg) for arg in line.args]
    elseif line.head == :return && line.args[1].head == :tuple
        outputs = [parse_arg!(ops, info, arg) for arg in line.args[1].args]
    elseif line.head == :line
        info = line
    else error("In function $expr do not know how to handle $(line.head) on $line")
    end
    info, outputs
end

function parse_assign!(ops, info, vals, expr::Symbol)
    @assert vals.head == :tuple "In assignment $vals = $expr do not know how to handle $(vals.head)"
    outputs = [vals.args...]
    push!(ops, Op(:fanout, [expr], outputs, [], info))
    outputs
end

function parse_assign!(ops, info, vals, expr::Expr)
    func, inputs = parse_expr!(ops, info, expr)
    outputs = outs(vals)
    push!(ops, Op(func, inputs, outputs, [], info))
    outputs
end

outs(vals::Symbol) = [vals]
outs(vals::SlotNumber) = [Symbol(string(vals))]
outs(vals::Expr) = [vals.args...]

function parse_expr!(ops, info, expr::Expr)
    if expr.head == :tuple
        args = [parse_arg!(ops, info, arg) for arg in expr.args]
        :tuple, args
    elseif expr.head == :call && expr.args[1] == GlobalRef(Base, :broadcast)
        if isa(expr.args[2], GlobalRef)
            expr.args[2] = expr.args[2].name
        end
        ".$(expr.args[2])", [parse_arg!(ops, info, arg) for arg in expr.args[3:end]]
    elseif expr.head == :call
        if isa(expr.args[1], GlobalRef)
            expr.args[1] = expr.args[1].name
        end
        args = [parse_arg!(ops, info, arg) for arg in expr.args[2:end]]
        while length(args) > 2 && (expr.args[1] == :(+) || expr.args[1] == :(*))
            a = shift!(args)
            arg = Symbol("tmp$(length(ops)+1)")
            push!(ops, Op(opname(expr.args[1]), [a, args[1]], [arg], [], info))
            args[1] = arg
        end
        opname(expr.args[1]), args
    elseif expr.head == :(.)
        @assert expr.args[2].head == :tuple
        ".$(expr.args[1])", [parse_arg!(ops, info, arg) for arg in expr.args[2].args]
    elseif expr.head == :ref
        args = [parse_arg!(ops, info, arg) for arg in expr.args]
        :ref, args
    elseif expr.head == :(:)
        args = [parse_arg!(ops, info, arg) for arg in expr.args]
        :colon, args
    else error("In expr $expr do not know how to handle $(expr.head)")
    end
end

opname(name) = get(opnames, name, name)
const opnames = Dict(:(.*) => :dot_times, :(*) => :times, :(.+) => :dot_plus, :(+) => :plus,
                     :(./) => :dot_divide, :(/) => :divide, :(.-) => :dot_minus, :(-) => :minus,
                     :(.^) => :dot_power, :(^) => :power, :getindex => :ref)

function parse_arg!(ops, info, arg::Expr)
    func, inputs = parse_expr!(ops, info, arg)
    arg = Symbol("tmp$(length(ops)+1)")
    push!(ops, Op(func, inputs, [arg], [], info))
    arg
end

parse_arg!(ops, info, arg::Symbol) = arg
parse_arg!(ops, info, arg::Number) = arg
parse_arg!(ops, info, arg::SlotNumber) = Symbol(string(arg))


#!/usr/bin/julia

using Random

function passwd(name :: String, site :: String, seed :: Int64, n :: Int64)
    Random.seed!(seed)
    nums = map((x)->Int64(x), collect(name))
    seed = atan(Int64(sum(nums)) * 131);
    for i=1:n
        res = 213*rand()*sin(seed * (Int64(name[(i % length(name)) + 1]) / 13 + Int64(site[(i % length(site)) + 1]) / 31))
        seed *= 3*sin(res/5)+5
        nextch = Int64(round(abs(res))) % 26 + 'a'
	print(rand() < 0.5 ? nextch : nextch - 32)
    end

end

passwd(ARGS[1], ARGS[2], parse(Int64, ARGS[3]), parse(Int64, ARGS[4]))
println()


##### Multi-dimensional interpolations

# TODO: Add implementations that perform interps with a single pass and no alloc

#####
##### CellCenter <-> CellEdge
#####

function interp!(cent::CellCenter, edge::CellEdge{dim}, grid::Grid) where {dim}
    face_orth = CellFace(grid, first(orthogs(dim)))
    interp!(face_orth, edge, grid)
    interp!(cent, face_orth, grid)
end

function interp!(edge::CellEdge{dim}, cent::CellCenter, grid::Grid) where {dim}
    face_orth = CellFace(grid, first(orthogs(dim)))
    interp!(face_orth, cent, grid)
    interp!(edge, face_orth, grid)
end

#####
##### CellCorner <-> CellFace
#####

function interp!(corn::CellCorner, face::CellFace{dim}, grid::Grid) where {dim}
    edge_orth = CellEdge(grid, first(orthogs(dim)))
    interp!(edge_orth, face, grid)
    interp!(corn, edge_orth, grid)
end

function interp!(face::CellFace{dim}, corn::CellCorner, grid::Grid) where {dim}
    edge_orth = CellEdge(grid, first(orthogs(dim)))
    interp!(edge_orth, corn, grid)
    interp!(face, edge_orth, grid)
end

#####
##### CellFace <-> CellEdge
#####

function interp!(face::CellFace{dim}, edge::CellEdge{dim}, grid::Grid) where {dim}
    face_orth = CellFace(grid, first(orthogs(dim)))
    cent = CellCenter(grid)
    interp!(face_orth, edge, grid)
    interp!(cent, face_orth, grid)
    interp!(face, cent, grid)
end

function interp!(edge::CellEdge{dim}, face::CellFace{dim}, grid::Grid) where {dim}
    cent = CellCenter(grid)
    face_orth = CellFace(grid, first(orthogs(dim)))
    interp!(cent, face, grid)
    interp!(face_orth, cent, grid)
    interp!(edge, face_orth, grid)
end

#####
##### CellCenter <-> CellCorner
#####

function interp!(cent::CellCenter, corn::CellCorner, grid::Grid)
    edge = CellEdge(grid, 1)
    interp!(edge, corn, grid)
    interp!(cent, edge, grid)
end

function interp!(corn::CellCorner, cent::CellCenter, grid::Grid)
    edge = CellEdge(grid, 1)
    interp!(edge, cent, grid)
    interp!(corn, edge, grid)
end


# Copyright 2017 Stephen L. Smith and Frank Imeson
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Select a removal and an insertion method using powers, and then perform
removal followed by insertion on tour.  Operation done in place.
"""
function remove_insert(current::Tour, best::Tour, dist::Array{Int64,2}, member::Array{Int64,1},
						setdist::Distsv, sets::Array{Any,1},
						powers, param::Dict{Symbol,Any}, phase::Symbol)
	# make a new tour to perform the insertion and deletion on
    trial = Tour(copy(current.tour), current.cost)
	pivot_tour!(trial.tour)
	num_removals = rand(param[:min_removals]:param[:max_removals])

	removal = power_select(powers["removals"], powers["removal_total"], phase)
	if removal.name == "distance"
		sets_to_insert = distance_removal!(trial.tour, dist, num_removals,
													member, removal.value)
    elseif removal.name == "worst"
		sets_to_insert = worst_removal!(trial.tour, dist, num_removals,
													member, removal.value)
	else
		sets_to_insert = segment_removal!(trial.tour, num_removals, member)
	end

	randomize_sets!(sets, sets_to_insert)
	# then perform insertion
	insertion = power_select(powers["insertions"], powers["insertion_total"], phase)
	noise = power_select(powers["noise"], powers["noise_total"], phase)
	if insertion.name == "cheapest"
		cheapest_insertion!(trial.tour, sets_to_insert, dist, setdist, sets)
	else
		randpdf_insertion!(trial.tour, sets_to_insert, dist, setdist, sets,
							insertion.value, noise)
	end
	rand() < param[:prob_reopt] && opt_cycle!(trial, dist, sets, member, param, setdist, "partial")

	# update power scores for remove and insert
	score = 100 * max(current.cost - trial.cost, 0)/current.cost
	insertion.scores[phase] += score
	insertion.count[phase] += 1
	removal.scores[phase] += score
	removal.count[phase] += 1
	noise.scores[phase] += score
	noise.count[phase] += 1
	return trial
end

"""
Select an integer between 1 and num according to
and exponential distribution with lambda = power
# goes from left of array if power is positive
# and right of array if it is negative
"""
function select_k(num::Int64, power::Float64)
	base = (1/2)^abs(power)
	# (1 - base^num)/(1 - base) is sum of geometric series
	rand_select = (1 - base^num)/(1 - base) * rand()
	bin = 1
	@inbounds for k = 1:num
		if rand_select < bin
			return (power >= 0 ? (num - k + 1) : k)
		end
		rand_select -= bin
		bin *= base
	end
	return (power >=0 ? num : 1)
end

"""
selecting a random k in 1 to length(weights) according to power
and then selecting the kth smallest element in weights
"""
function pdf_select(weights::Array{Int64,1}, power::Float64)
    power == 0.0 && return rand(1:length(weights))
    power > 9.0 && return rand_select(weights, maximum(weights))
    power < - 9.0 && return rand_select(weights, minimum(weights))

	# select kth smallest.  If 1 or length(weights), simply return
	k = select_k(length(weights), power)
	k == 1 && return rand_select(weights, minimum(weights))
	k == length(weights) && return rand_select(weights, maximum(weights))
	val = partialsort(weights, k)

	return rand_select(weights, val)
end

"""  choose set with pdf_select, and then insert in best place with noise  """
function randpdf_insertion!(tour::Array{Int64,1}, sets_to_insert::Array{Int64,1},
							dist::Array{Int64, 2}, setdist::Distsv,
							sets::Array{Any, 1}, power::Float64, noise::Power)

    mindist = [typemax(Int64) for i=1:length(sets_to_insert)]
    @inbounds for i = 1:length(sets_to_insert)
        set = sets_to_insert[i]
        for vertex in tour
            if setdist.min_sv[set, vertex] < mindist[i]
                mindist[i] = setdist.min_sv[set, vertex]
            end
        end
    end
    new_vertex_in_tour = 0

    @inbounds while length(sets_to_insert) > 0
        if new_vertex_in_tour != 0
            for i = 1:length(sets_to_insert)
                set = sets_to_insert[i]
                if setdist.min_sv[set, new_vertex_in_tour] < mindist[i]
                    mindist[i] = setdist.min_sv[set, new_vertex_in_tour]
                end
            end
        end
        set_index = pdf_select(mindist, power) # select set to insert from pdf
        # find the closest vertex and the best insertion in that vertex
        nearest_set = sets_to_insert[set_index]
		if noise.name == "subset"
			bestv, bestpos = insert_subset_lb(tour, dist, sets[nearest_set], nearest_set,
											  setdist, noise.value)
		else
			bestv, bestpos =
					insert_lb(tour, dist, sets[nearest_set], nearest_set, setdist, noise.value)
		end
        insert!(tour, bestpos, bestv)  # perform the insertion
        new_vertex_in_tour = bestv
        # remove the inserted set from data structures
        splice!(sets_to_insert, set_index)
        splice!(mindist, set_index)
    end
end

function cheapest_insertion!(tour::Array{Int64,1}, sets_to_insert::Array{Int64,1},
	dist::Array{Int64, 2}, setdist::Distsv, sets::Array{Any, 1})
    """
	choose vertex that can be inserted most cheaply, and insert it in that position
	"""
	while length(sets_to_insert) > 0
        best_cost = typemax(Int64)
        best_v = 0
        best_pos = 0
        best_set = 0
        for i = 1:length(sets_to_insert)
            set_ind = sets_to_insert[i]
            # find the best place to insert the vertex
            best_v, best_pos, cost = insert_cost_lb(tour, dist, sets[set_ind], set_ind, setdist,
										  best_v, best_pos, best_cost)
			if cost < best_cost
				best_set = i
				best_cost = cost
			end
        end

        # now, perform the insertion
        insert!(tour, best_pos, best_v)
        # remove the inserted set from data structures
        splice!(sets_to_insert, best_set)
    end
end

"""
Given a tour and a set, this function finds the vertex in the set with minimum
insertion cost, along with the position of this insertion in the tour.  If
best_position is i, then vertex should be inserted between tour[i-1] and tour[i].
"""
@inline function insert_lb(tour::Array{Int64,1}, dist::Array{Int64,2}, set::Array{Int64, 1},
							setind::Int, setdist::Distsv, noise::Float64)
	best_cost = typemax(Int64)
	bestv = 0
	bestpos = 0

	@inbounds for i = 1:length(tour)
		v1 = prev_tour(tour, i)
		lb = setdist.vert_set[v1, setind] + setdist.set_vert[setind, tour[i]] - dist[v1, tour[i]]
		lb > best_cost && continue

		for v in set
	        insert_cost = dist[v1, v] + dist[v, tour[i]] - dist[v1, tour[i]]
			noise > 0.0 && (insert_cost += round(Int64, noise * rand() * abs(insert_cost)))
			if insert_cost < best_cost
				best_cost = insert_cost
				bestv = v
				bestpos = i
	        end
		end
    end
    return bestv, bestpos
end

@inline function insert_subset_lb(tour::Array{Int64,1}, dist::Array{Int64,2}, set::Array{Int64, 1},
							setind::Int, setdist::Distsv, noise::Float64)
	best_cost = typemax(Int64)
	bestv = 0
	bestpos = 0
	tour_inds = collect(1:length(tour))

	@inbounds for i = 1:ceil(Int64, length(tour) * noise)
		i = incremental_shuffle!(tour_inds, i)
		v1 = prev_tour(tour, i)
		lb = setdist.vert_set[v1, setind] + setdist.set_vert[setind, tour[i]] - dist[v1, tour[i]]
		lb > best_cost && continue

		for v in set
	        insert_cost = dist[v1, v] + dist[v, tour[i]] - dist[v1, tour[i]]
	        if insert_cost < best_cost
				best_cost = insert_cost
				bestv = v
				bestpos = i
	        end
		end
    end
    return bestv, bestpos
end

############ Initial Tour Construction ##########################

"""build tour from scratch on a cold restart"""
function initial_tour!(lowest::Tour, dist::Array{Int64, 2}, sets::Array{Any, 1},
						setdist::Distsv, trial_num::Int64, param::Dict{Symbol,Any})
	sets_to_insert = collect(1:param[:num_sets])
	best = Tour(Int64[], typemax(Int64))

	# compute random initial tour only if past first trial
	# in this case, randomly choose between random and insertion tour.
	if param[:init_tour] == "warm" && param[:warm_tour] != nothing
 		best.tour = param[:warm_tour]
	elseif param[:init_tour] == "rand" && (trial_num > 1) && (rand() < 0.5)
		random_initial_tour!(best.tour, sets_to_insert, dist, sets)
	else
		random_insertion!(best.tour, sets_to_insert, dist, sets, setdist)
	end
	best.cost = tour_cost(best.tour, dist)
	lowest.cost > best.cost && (lowest = best)
	return best
end

"""
Randomly shuffle the sets, and then insert the best vertex from each set back into
the tour where sets are considered in shuffled order.
"""
function random_insertion!(tour::Array{Int64,1}, sets_to_insert::Array{Int64,1},
						   dist::Array{Int64, 2}, sets::Array{Any, 1}, setdist::Distsv)
    shuffle!(sets_to_insert)  # randomly permute the sets
    for set in sets_to_insert
        # only have to compute the insert cost for the changed portion of the tour
        if isempty(tour)
            best_vertex = rand(sets[set])
            best_position = 1
        else
            best_vertex, best_position = insert_lb(tour, dist, sets[set], set, setdist, 0.75)
        end
        # now, perform the insertion
        insert!(tour, best_position, best_vertex)
    end
end

"""
Randomly shuffle the sets, and then insert the best vertex from each set back into
the tour where sets are considered in shuffled order.
"""
function random_initial_tour!(tour::Array{Int64,1}, sets_to_insert::Array{Int64,1},
							  dist::Array{Int64, 2}, sets::Array{Any, 1})
    shuffle!(sets_to_insert)
    for set in sets_to_insert
		push!(tour, rand(sets[set]))
    end
end

######################### Removals ################################
"""
Remove the vertices randomly, but biased towards those that add the most length to the
tour.  Bias is based on the power input.  Vertices are then selected via pdf select.
"""
function worst_removal!(tour::Array{Int64,1}, dist::Array{Int64, 2},
							num_to_remove::Int64, member::Array{Int64,1}, power::Float64)
    deleted_sets = Array{Int}(undef, 0)
	while length(deleted_sets) < num_to_remove
		removal_costs = worst_vertices(tour, dist)
		ind = pdf_select(removal_costs, power)
		set_to_delete = member[tour[ind]]

        # perform the deletion
        push!(deleted_sets, set_to_delete)
		splice!(tour, ind)
	end
    return deleted_sets
end

""" removing a single continuos segment of the tour of size num_remove """
function segment_removal!(tour::Array{Int64, 1}, num_to_remove::Int64, member::Array{Int64,1})
	i = rand(1:length(tour))
	deleted_sets = Array{Int}(undef, 0)
	while length(deleted_sets) < num_to_remove
		i > length(tour) && (i = 1)
		push!(deleted_sets, member[tour[i]])
		splice!(tour, i)
	end
	return deleted_sets
end

"""  pick a random vertex, and delete its closest neighbors  """
function distance_removal!(tour::Array{Int64,1}, dist::Array{Int64, 2},
							   num_to_remove::Int64, member::Array{Int64,1}, power::Float64)
    deleted_sets = Array{Int}(undef, 0)
    deleted_vertices = Array{Int}(undef, 0)

    seed_index = rand(1:length(tour))
    push!(deleted_sets, member[tour[seed_index]])
    push!(deleted_vertices, tour[seed_index])
    splice!(tour, seed_index)

    while length(deleted_sets) < num_to_remove
        # pick a random vertex from the set of deleted vertices
        seed_vertex = rand(deleted_vertices)
        # find closest vertex to the seed vertex that's still in the tour
        mindist = zeros(Int64, length(tour))
        for i = 1:length(tour)
			mindist[i] = min(dist[seed_vertex, tour[i]], dist[tour[i], seed_vertex])
        end
        del_index = pdf_select(mindist, power)
        push!(deleted_sets, member[tour[del_index]])
        push!(deleted_vertices, tour[del_index])
        splice!(tour, del_index)
    end

    return deleted_sets
end

"""
determine the cost of removing each vertex from the tour, given that all others remain.
"""
function worst_vertices(tour::Array{Int64, 1}, dist::Array{Int64, 2})
    removal_cost = zeros(Int64, length(tour))
    @inbounds for i = 1:length(tour)
        if i == 1
            removal_cost[i] = dist[tour[end], tour[i]] +
								dist[tour[i], tour[i+1]] - dist[tour[end], tour[i+1]]
        elseif i == length(tour)
            removal_cost[i] = dist[tour[i-1], tour[i]] +
								dist[tour[i], tour[1]] - dist[tour[i-1], tour[1]]
        else
            removal_cost[i] = dist[tour[i-1], tour[i]] +
								dist[tour[i], tour[i+1]] - dist[tour[i-1], tour[i+1]]
		end
    end
    return removal_cost
end


# Rcpp derivatives (RcppSundials)
function create_derivatives_rcpp(language, sorted_model, name)
  code = ""
  for level in 1:length(sorted_model.SortedEquations)
    for (lhs, rhs) in sorted_model.SortedEquations[level]
        if level == 1
          if(ismatch(r"params|states|forcs", string(rhs.Expr)))
           number = match(r"(?<=\[)[\d]+(?=\])",string(rhs.Expr))
           new_expr = replace(string(rhs.Expr), "[$(number.match)]", "[$(int(number.match) - 1)]")
           code *= "const double $(new_expr);\n"
         else
           code *= "const double $(rhs.Expr);\n"
         end
        else
          mod_expr = copy(rhs.Expr)
          mod_expr = sub_product(mod_expr)
          mod_expr = substitute_power(mod_expr)
          mod_expr = sub_minmax(mod_expr)
          code *= "const double $lhs" * " = " * replace(replace(string(mod_expr), ".*", "*"), ":", "") * ";\n"
        end
    end
  end
  # Return line
  return_line = create_return_line_Rcpp(sorted_model.NamesDerivatives,sorted_model.NamesObserved)
  # Return the output
  code = replace(code, "1.0 *", "")
  up_boiler_plate =
"""
  array<vector<double>, 2> $(name)(const double& t, const vector<double>& states,
            const vector<double>& params, const vector<double>& forcs) { \n
"""
  low_boiler_plate =
"""
}
"""
  return paste("\n", up_boiler_plate, code, return_line,low_boiler_plate)
end

# Rcpp jacobian (RcppSundials)
function create_jacobian_rcpp(language, compressed_model, jacobian_matrix, names_derivatives, name)
  # Substitute * in AST to avoid pretty printing of scalar product
  for i = 1:size(jacobian_matrix, 1), j = 1:size(jacobian_matrix, 2)
      if isa(jacobian_matrix[i,j], Expr)
        jacobian_matrix[i,j] = sub_product(jacobian_matrix[i,j])
      end
  end
  code = ""
  for (lhs, rhs) in compressed_model.SortedEquations[1]
      if (ismatch(r"params|states|forcs", string(rhs.Expr)))
       number = match(r"(?<=\[)[\d]+(?=\])",string(rhs.Expr))
       new_expr = replace(string(rhs.Expr), "[$(number.match)]", "[$(int(number.match) - 1)]")
       code *= "const double $(new_expr);\n"
     else
       code *= "const double $(rhs.Expr);\n"
     end
  end
  for i = 1:size(jacobian_matrix, 1), j = 1:size(jacobian_matrix, 2)
    if isa(jacobian_matrix[i,j], Number) && eval(jacobian_matrix[i,j]) == 0
      continue
    else
      rhs = substitute_power(jacobian_matrix[i,j])
      rhs = sub_minmax(rhs)
      rhs = replace(replace(string(rhs), ".*", "*"), ":", "")
      code *= "output.at($(i-1),$(j - 1)) = $rhs;\n"
    end
  end

  up_boiler_plate =
"""
arma::mat $(name)_jacobian(const double& t, const vector<double>& states,
          const vector<double>& params, const vector<double>& forcs) {
  arma::mat output = arma::eye(states.size(), states.size());\n
"""
  low_boiler_plate =
"""
    return output;
  }
"""

  return paste("\n", up_boiler_plate, code,low_boiler_plate)

end

# Write the file with all the model functions
function write_code_rcpp!(dynamic_type, model_function, name, file)

  up_boiler_plate =
"""
#include <RcppArmadillo.h>
#include <array>
#include <vector>
#include <math.h>
using namespace std;

double ifelse(bool condition, const double& result1, const double& result2) {
  if(condition) {
    return result1;
  } else {
    return result2;
  }
}

inline double heaviside(const double& arg) {
  return arg <= 0.0 ? 0.0 : 1.0;
}

inline double dirac(const double& arg) {
  return abs(arg) <= numeric_limits<double>::epsilon()  ? 0 : numeric_limits<double>::infinity();
}

inline double Min(const double& arg1) {
  return arg1;
}

inline double Max(const double& arg1) {
  return arg1;
}

inline double Min(const double& arg1, const double& arg2) {
  return arg1 <= arg2 ? arg1 : arg2;
}

inline double Max(const double& arg1, const double& arg2) {
  return arg1 >= arg2 ? arg1 : arg2;
}

extern "C" {
"""
  low_boiler_plate =
"""
};
"""

  f = open("$(file).cpp","w")
  println(f, up_boiler_plate)
  println(f, model_function)
  println(f, low_boiler_plate)
  close(f)

  nothing

end

############################ AUXILLIARY FUNCTIONS ##############################

# Because Julia will use "pretty printing" for scalar product, we need to fool the parser
# by temporarily substituying by another binary operator and then doing string replacement
function sub_product(ex::Expr)
  for i in 1:length(ex.args)
  if ex.args[i] == :*
    ex.args[i] = :.*
  elseif isa(ex.args[i], Expr)
     ex.args[i] = sub_product(ex.args[i])
  end
end
return ex
end

sub_product(ex::Symbol) = ex

# Substitute the annoying x^y for pow(x,y)
function substitute_power(ex::Expr)
  new_ex = deepcopy(ex)
  for i in 1:length(new_ex.args)
    if isa(new_ex.args[i], Expr)
      new_ex.args[i] = substitute_power(new_ex.args[i])
    elseif new_ex.args[i] == :^
      new_ex.args[i] = :pow
    end
  end
  return new_ex
end

substitute_power(ex::Any) = ex

function sub_minmax(ex::Expr)
  new_ex = deepcopy(ex)
  for i in 1:length(new_ex.args)
    if new_ex.args[i] == :min || new_ex.args[i] == :max
        new_ex.args[i] == :min ?  (new_ex.args[i] = :Min) : (new_ex.args[i] = :Max)
    elseif isa(new_ex.args[i], Expr)
        new_ex.args[i] = sub_minmax(new_ex.args[i])
  end
end
return new_ex
end

sub_minmax(ex::Symbol) = ex

# Create a return line with an array containing the derivatives and observed functions (for the STL version of the model)
function create_return_line_Rcpp(states, observed)
    return_line = "vector<double> derivatives{"
    for i in states
      if i != states[end]
        return_line = return_line  * i * ","
      else
        return_line = return_line  * i
      end
    end
    return_line = return_line * "};\n vector<double> observed{"
    for i in observed
      if i != observed[end]
        return_line = return_line * i * ", "
      else
        return_line = return_line  * i
      end
    end
    return_line = return_line * "};\n array<vector<double>,2> output{derivatives, observed};\n return output;"
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

struct SMPSModel{N, T <: AbstractFloat, M <: AbstractMatrix}
    raw::RawSMPS{T}
    stages::NTuple{N, SMPSStage{T,M}}

    function SMPSModel(raw::RawSMPS{T}, stages::NTuple{N, SMPSStage{T,M}}) where {N, T <: AbstractFloat, M <: AbstractMatrix}
        return new{N,T,M}(raw, stages)
    end
end

function SMPSModel(raw::RawSMPS)
    N = length(raw.tim.stages)
    constructor = sparsity(raw.cor) >= 0.5 ? SparseLPData : LPData
    stages = ntuple(Val(N)) do i
        parse_stage(raw, i, N, constructor)
    end
    return SMPSModel(raw, stages)
end

function parse_stage(raw::RawSMPS, stage::Integer, N::Integer, constructor)
    1 <= stage <= N || error("$(stage.id) not in range 1 to $N.")
    # Get required data
    tim = raw.tim
    cor = raw.cor
    sto = raw.sto
    if stage == 1
        # Get row and column of the given stage
        row = tim.row_delims[stage]
        col = tim.col_delims[stage]
        # Compute the rowrange up to the next stage
        next_row = tim.row_delims[stage + 1]
        rowrange = collect(keys(cor.rows))[cor.rows[row][1]:cor.rows[next_row][1]-1]
        # Compute the colrange up to the next stage
        next_col = tim.col_delims[stage + 1]
        colrange = collect(keys(cor.vars))[cor.vars[col]:cor.vars[next_col]-1]
        # Extract first-stage model
        model = constructor(cor; colrange = colrange, rowrange = rowrange, include_constant = true)
        # Return stage
        return SMPSStage(model)
    else
        # Get row and column of the current stage
        row = tim.row_delims[stage]
        col = tim.col_delims[stage]
        # and the column of the previous stage
        prev_col = tim.col_delims[stage - 1]
        # Compute the rowrange up to the next stage
        rowrange = if stage < N
            next_row = tim.row_delims[stage + 1]
            rowrange = collect(keys(cor.rows))[cor.rows[row][1]:cor.rows[next_row][1]-1]
        else
            rowrange = collect(keys(cor.rows))[cor.rows[row][1]:end]
        end
        # Compute the colrange from the previos stage up to the next stage
        colrange = if stage < N
            next_col = tim.col_delims[stage + 1]
            colrange = collect(keys(cor.vars))[cor.vars[prev_col]:cor.vars[next_col]-1]
        else
            colrange = collect(keys(cor.vars))[cor.vars[prev_col]:end]
        end
        # Extract second-stage model
        model = constructor(cor; colrange = colrange, rowrange = rowrange)
        # Compute technology and recourse ranges
        delim      = findfirst(x -> x == col, colrange)
        technology = 1:delim-1
        recourse   = delim:length(colrange)
        # Extract uncertainty perturbation
        uncertainty = uncertainty_template(sto, cor, model, stage - 1)
        # Return stage
        return SMPSStage(stage, model, uncertainty, technology, recourse)
    end
end

function stochastic_model(smps::SMPSModel{2})
    # Create stochastic model
    sm = @stochastic_model begin
        @stage 1 begin
            @parameters begin
                data = smps.stages[1]
            end
            # Unpack first-stage SMPS data
            n    = data.model.indexmap.n
            m₂   = data.model.indexmap.m₂
            c₁   = data.model.c₁
            c₂   = data.model.c₂
            A    = data.model.A
            b    = data.model.b
            C, d = canonical(data.model.C,
                             data.model.d₁,
                             data.model.d₂)
            lb   = data.model.lb
            ub   = data.model.ub
            bin  = data.model.is_binary
            int  = data.model.is_integer
            # Define all first-stage variables as decisions
            if any(isfinite.(lb))
                if any(isfinite.(ub))
                    # Double bounds
                    @decision(model, lb[i] <= x[i in 1:n] <= ub[i])
                else
                    # Lower bounds
                    @decision(model, x[i in 1:n] >= lb[i])
                end
            elseif any(isfinite.(ub))
                # Upper bounds
                @decision(model, x[i in 1:n] <= ub[i])
            else
                # Free
                @decision(model, x[i in 1:n])
            end
            # Add any binary or integer restrictions
            for i in 1:n
                if bin[i]
                    @constraint(model, x[i] in MOI.ZeroOne())
                elseif int[i]
                    @constraint(model, x[i] in MOI.Integer())
                end
            end
            # Define objective and constraints
            if any(abs.(c₁) .> sqrt(eps())) || abs(c₂) > sqrt(eps())
                @objective(model, Min, dot(c₁, x) + c₂)
            end
            if length(A) > 0 && length(b) > 0
                @constraint(model, A * x .== b)
            end
            if length(C) > 0 && length(d) > 0
                @constraint(model, C * x .<= d)
            end
        end
        @stage 2 begin
            @parameters begin
                data = smps.stages[2]
            end
            @uncertain Δq ΔT ΔW Δh ΔC Δd₁ Δd₂ from SMPSScenario
            # Unpack first-stage SMPS data
            n    = length(data.recourse)
            m₂   = data.model.indexmap.m₂
            q    = data.model.c₁[data.recourse]
            Δq   = Δq[data.recourse]
            T    = data.model.A[:, data.technology]
            W    = data.model.A[:, data.recourse]
            h    = data.model.b
            C̃, h̃ = canonical(data.model.C + ΔC,
                             data.model.d₁ + Δd₁,
                             data.model.d₂ + Δd₂)
            T̃    = C̃[:, data.technology]
            W̃    = C̃[:, data.recourse]
            lb   = data.model.lb[data.recourse]
            ub   = data.model.ub[data.recourse]
            bin  = data.model.is_binary[data.recourse]
            int  = data.model.is_integer[data.recourse]
            # Define all recourse variables
            if any(isfinite.(lb))
                if any(isfinite.(ub))
                    # Double bounds
                    @recourse(model, lb[i] <= y[i in 1:n] <= ub[i])
                else
                    # Lower bounds
                    @recourse(model, y[i in 1:n] >= lb[i])
                end
            elseif any(isfinite.(ub))
                # Upper bounds
                @recourse(model, y[i in 1:n] <= ub[i])
            else
                # Free
                @recourse(model, y[i in 1:n])
            end
            # Add any binary or integer restrictions
            for i in 1:n
                if bin[i]
                    @constraint(model, y[i] in MOI.ZeroOne())
                elseif int[i]
                    @constraint(model, y[i] in MOI.Integer())
                end
            end
            # Define objective and constraints
            @objective(model, Min, dot((q + Δq), y))
            if (length(T) > 0 || length(W) > 0) && length(h) > 0
                @constraint(model, (T + ΔT) * x + (W + ΔW) * y .== (h + Δh))
            end
            if (length(T̃) > 0 || length(W̃) > 0) && length(h̃) > 0
                @constraint(model, T̃ * x + W̃ * y .<= h̃)
            end
        end
    end
    return sm
end


using IntcodeVM;
using DelimitedFiles;

# In day 9 we complete the Intcode VM and validate its correctness.
# Many Advent of Code participants (including myself) would have found
# on this day that their Intcode interpreters were under-engineered.
# The big thing people needed was encapsulation. Intcode VMs should be
# easy to start, reasonably self-contained, and easy to get output from.

print("Day 9 Part 1: ");
IntcodeVM.run(vec(readdlm(ARGS[1], ',', Int, '\n')), inputs=[1], out=stdout)

print("Day 9 Part 2: ");
IntcodeVM.run(vec(readdlm(ARGS[1], ',', Int, '\n')), inputs=[2], out=stdout)


using CSV

df = CSV.File(joinpath(@__DIR__, "benchmarks.csv"))

# Display formatted table

function pretty_print(io, x::Number)
    d, r = fldmod(x, 1)
    r ≈ 0 ? print(io, Int(d)) : print(io, x)
end

function print_benchmarks(io, df)
    print(io, raw"""
    \begin{tabular}{llllll}
        \hline
        \textbf{machine} & \textbf{size} & \textbf{device} & \textbf{forward (\si{\milli\second})} & \textbf{backward (\si{\milli\second})} & \textbf{ratio} \\ \hline
    """)
    for row in df
        print(io, "    ")
        print(io, row.machine)
        print(io, " & ")
        print(io, row.problemsize)
        print(io, " & ")
        print(io, row.device)
        print(io, " & ")
        pretty_print(io, round(row.forward / 1000, sigdigits=5))
        print(io, " & ")
        pretty_print(io, round(row.backward / 1000, sigdigits=5))
        print(io, " & ")
        pretty_print(io, round(row.backward / row.forward, digits=3))
        print(io, raw" \\ ")
        print(io, '\n')
    end
    print(io, raw"""
        \hline
    \end{tabular}
    """)
end

print_benchmarks(stdout, df)

# Display figure

using AlgebraOfGraphics, CairoMakie
set_aog_theme!()

passes = [:forward, :backward]

plt = data(df) * mapping(
    :machine => sorter("dense", "convolution", "recurrent"),
    passes .=> (t -> (t / 1000)) .=> "runtime (μs)",
    color = dims(1) => renamer(passes) => "pass",
    dodge = dims(1) => renamer(passes) => "pass",
    col = :device,
    row = :problemsize => renamer("small" => "small size", "medium" => "medium size")
)

fig = Figure(resolution=(600, 800))

ag = draw!(
    fig, plt * visual(BarPlot, fillto = 1, dodge_gap = 0.),
    axis=(xticklabelrotation=π/6, yscale=log10, yticks= LogTicks(LinearTicks(6)))
)
legend!(fig[0, 1:2], ag, orientation=:horizontal)

display(fig)


"""
Run two coarse rising thermal bubble simulations and make sure that when
restarting from a checkpoint, the restarted simulation matches the non-restarted
simulation numerically.
"""
function run_thermal_bubble_checkpointer_tests()
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 100, 100, 100
    Δt = 6

    true_model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), ν=4e-2, κ=4e-2)

    # Add a cube-shaped warm temperature anomaly that takes up the middle 50%
    # of the domain volume.
    i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4)
    j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4)
    k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
    true_model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01

    checkpointed_model = deepcopy(true_model)

    time_step!(true_model, 10, Δt)

    checkpointer = Checkpointer(dir=".", prefix="test_", frequency=5, padding=1)
    push!(checkpointed_model.output_writers, checkpointer)

    # Checkpoint should be saved as "test_model_checkpoint_5.jld" after the
    # 5th iteration.
    time_step!(checkpointed_model, 5, Δt)

    # Remove all knowledge of the checkpointed model.
    checkpointed_model = nothing

    restored_model = restore_from_checkpoint("test_model_checkpoint_5.jld")

    time_step!(restored_model, 5, Δt)

    # Now the true_model and restored_model should be identical.
    @test all(restored_model.velocities.u.data .≈ true_model.velocities.u.data)
    @test all(restored_model.velocities.v.data .≈ true_model.velocities.v.data)
    @test all(restored_model.velocities.w.data .≈ true_model.velocities.w.data)
    @test all(restored_model.tracers.T.data    .≈ true_model.tracers.T.data)
    @test all(restored_model.tracers.S.data    .≈ true_model.tracers.S.data)
    @test all(restored_model.G.Gu.data         .≈ true_model.G.Gu.data)
    @test all(restored_model.G.Gv.data         .≈ true_model.G.Gv.data)
    @test all(restored_model.G.Gw.data         .≈ true_model.G.Gw.data)
    @test all(restored_model.G.GT.data         .≈ true_model.G.GT.data)
    @test all(restored_model.G.GS.data         .≈ true_model.G.GS.data)
end

"""
Run a coarse thermal bubble simulation and save the output to NetCDF at the
10th time step. Then read back the output and test that it matches the model's
state.
"""
function run_thermal_bubble_netcdf_tests()
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 100, 100, 100
    Δt = 6

    model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), ν=4e-2, κ=4e-2)

    # Add a cube-shaped warm temperature anomaly that takes up the middle 50%
    # of the domain volume.
    i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4)
    j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4)
    k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
    model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01

    nc_writer = NetCDFOutputWriter(dir=".", prefix="test_", frequency=10, padding=1)
    push!(model.output_writers, nc_writer)

    time_step!(model, 10, Δt)

    u = read_output(nc_writer, "u", 10)
    v = read_output(nc_writer, "v", 10)
    w = read_output(nc_writer, "w", 10)
    T = read_output(nc_writer, "T", 10)
    S = read_output(nc_writer, "S", 10)

    @test all(u .≈ model.velocities.u.data)
    @test all(v .≈ model.velocities.v.data)
    @test all(w .≈ model.velocities.w.data)
    @test all(T .≈ model.tracers.T.data)
    @test all(S .≈ model.tracers.S.data)
end


using DAWG
using BenchmarkTools

#download csv dict from https://github.com/neologd/mecab-unidic-neologd/tree/master/seed
#decompress to any directory such as c:/tmp
#use only word with 3 to 11 characters (about 1.5 million words)

cd("C:/tmp")
sampledict = Dict()
f = open("mecab-unidic-user-dict-seed.20161215.csv")
for ln in eachline(f)
  word = (split(ln,","))[1]
  wlen = length(word)
  if wlen < 3 || wlen > 11
    continue
  end

  if !haskey(sampledict,word)
    sampledict[word]=true
  end
end
close(f)

l = collect(keys(sampledict))
println(length(l)," words in dictionary.")

@time smdawg = makeSparseMatrixDAWG(l)
#534.940754 seconds (673.87 M allocations: 30.877 GB, 18.57% gc time)

@time dadawg = makeDoubleArrayDAWG(l)

teststring = """
公開当日に生放送された。
出演者
木村拓哉
北川景子
陣内智則のプレゼンは時間の関係で省略になった（エンディングの浜辺のシーンで松たか子とのアドリブ）。
"""

smdawglist = DictionaryMatch(teststring,smdawg)
dadawglist = DictionaryMatch(teststring,dadawg)

@benchmark smdawglist = DictionaryMatch(teststring,smdawg,extract=false)

@benchmark dadawglist = DictionaryMatch(teststring,dadawg,extract=false)


# perturb an array m by percentage of m0, when mask is true
function perturb!(m, m0, perc, mask::Function)
    for i in CartesianIndices(m)
        ii = Tuple(i)
        if (mask(ii))
            m[i] = m[i] + perc * 1e-2 * m0
        end
    end
end
function perturb!(m, m0, perc, idx::CartesianIndices)
    for i in idx
        m[i] = m[i] + perc * 1e-2 * m0
    end
end

"""
In-place method to add a rectangular box or cuboid in the medium. Its position is specified by its corners
and perc denotes the percentage of the perturbation.

```julia
update!(medium, field; rectangle=[[0,10],[10,10]], perc=10) # for 2D medium
update!(medium, field; rectangle=[[0,10,15],[10,10,15]], perc=10) # for 3D medium
```

In-place method to add random noise to medium.

```julia
update!(medium, fields; randn_perc=10)
```
"""
function update!(
    medium::Medium,
    fields::Vector{Symbol};
    rectangle = nothing,
    perc = 0.0,
    randn_perc = 0.0,
)

    mgrid = medium.mgrid
    if (!(rectangle === nothing))
        @assert length(rectangle)==2
        @assert length.(rectangle)==fill(ndims(medium),2)
        rectangle_indices = CartesianIndices(
            Tuple([
                argmin(abs.(mg .- rectangle[1][im])):argmin(abs.(mg .- rectangle[2][im])) for (im, mg) in enumerate(mgrid)
            ]),
        )
        if (!iszero(perc))
            for field in fields
                m = medium.m[field]
                m0 = medium.ref[field]
                perturb!(m, m0, perc, rectangle_indices)
            end
        end
    end
    if (!iszero(randn_perc))
        for field in fields
            m = medium.m[field]
            m0 = medium.ref[field]
            for i in eachindex(m)
                m[i] = χ(Float64(χ(m[i], m0, 1) + randn() * randn_perc * 1e-2), m0, -1)
            end

        end
    end

    return medium
end

# # Keyword Arguments

# * `point_loc::Vector{Float64}=[0., 0.,]` : approx location of point pert
# * `point_pert::Float64=0.0` : perturbation at the point scatterer
# * `ellip_loc::Vector{Float64}=nothing` : location of center of perturbation, [z, x]
# * `ellip_rad::Float64=0.0` : size of elliptic perturbation
# * `ellip_pert::Float64=0.1` : perturbation inside the ellipse
# * `ellip_α=0.0` : rotate the ellipse
# * `rect_loc::Array{Float64}=nothing` : rectangle location, [zmin, xmin, zmax, xmax]
# * `rect_pert::Float64=0.1` : perturbation in a rectangle
# * `constant_pert::Float64=0.0` : constant perturbation
# * `randn_pert::Float64=0.0` : percentage of reference values for additive random noise
# * `fields::Vector{Symbol}=[:χvp,:χrho,:χvs]` : which fields are to be modified?
# * `onlyin` : `medium` is modified only when field values are in these ranges
# """
# function update!(
#     medium::Medium,
#     fields::Vector{Symbol};
#     point_loc = fill(0.0, ndims(medium)),
#     point_pert::Float64 = 0.0,
#     ellip_loc = fill(0.0, ndims(medium)),
#     ellip_rad = 0.0,
#     ellip_pert::Real = 0.0,
#     ellip_α = 0.0,
#     rect_loc = fill(0.0, ndims(medium)),
#     rect_pert::Float64 = 0.0,
#     constant_pert::Float64 = 0.0,
#     randn_perc::Real = 0.0,
#     layerlocations = nothing,
# )

#     # only editing basic fields is allowed
#     for field in fields
#         @assert field ∈ [:vp, :vs, :rho, :Q]
#     end
#     rect_loc = convert.(Float64, rect_loc)
#     ellip_loc = convert.(Float64, ellip_loc)

#     temp = zeros(length.(medium.mgrid)...)

#     ipointloc = [
#         Interpolation.indminn(medium.mgrid[i], Float64(point_loc[i]), 1)[1] for
#         i = 1:length(medium.mgrid)
#     ]
#     temp[ipointloc...] += point_pert

#     if (!(ellip_pert == 0.0))
#         α = ellip_α * pi / 180.0
#         # circle or ellipse
#         rads =
#             (length(ellip_rad) == 1) ? [ellip_rad[1], ellip_rad[1]] :
#             [ellip_rad[1], ellip_rad[2]]

#         temp += [
#             (
#                 (
#                     (
#                         (
#                             (medium.mgrid[2][ix] - ellip_loc[2]) * cos(α) +
#                             (medium.mgrid[1][iz] - ellip_loc[1]) * sin(α)
#                         )^2 * inv(rads[1]^2) +
#                         (
#                             (-medium.mgrid[1][iz] + ellip_loc[1]) * cos(α) +
#                             (medium.mgrid[2][ix] - ellip_loc[2]) * sin(α)
#                         )^2 * inv(rads[2]^2)
#                     ) <= 1.0
#                 ) ? Float64(ellip_pert) : 0.0
#             ) for iz = 1:length(medium.mgrid[1]), ix = 1:length(medium.mgrid[2])
#         ]
#     end
#     if (!(rect_pert == 0.0))
#         temp += [
#             (
#                 (
#                     (medium.mgrid[2][ix] - rect_loc[4]) * (medium.mgrid[2][ix] - rect_loc[2]) < 0.0
#                 ) & (
#                     (medium.mgrid[1][iz] - rect_loc[3]) * (medium.mgrid[1][iz] - rect_loc[1]) < 0.0
#                 )
#             ) ? rect_pert : 0.0 for iz = 1:length(medium.mgrid[1]),
#             ix = 1:length(medium.mgrid[2])
#         ]
#     end
#     if (!(constant_pert == 0.0))
#         temp .+= constant_pert
#     end

#     for (iff, field) in enumerate(fields)
#         m = medium.m[field]
#         if (!(layerlocations === nothing))
#             onlyatvalues = []
#             for ipos in layerlocations
#                 ipx = Interpolation.indminn(medium.mgrid[2], Float64(ipos[2]), 1)[1]
#                 ipz = Interpolation.indminn(medium.mgrid[1], Float64(ipos[1]), 1)[1]
#                 push!(onlyatvalues, m[ipz, ipx])
#             end
#         end
#         for i in eachindex(m)
#             # if(((m[i]-onlyin[iff][1])*(m[i]-onlyin[iff][2]))<0.0)
#             if (layerlocations === nothing)
#                 m[i] += temp[i]
#             else
#                 if (m[i] ∈ onlyatvalues)
#                     m[i] += temp[i]
#                 end
#             end
#         end
#     end

#     return nothing
# end


export plotsDir
export dataDir
export Beamsplitter
export classicalFreqData
export classicalSinglePlot
export classicalSumPlot

"""
    plotsDir(name::String,dirname::String = "plots")

Concatenates string with name of plots directory
"""
function plotsDir(name::String,dirname::String = "plots")
    return joinpath(dirname,name)
end

"""
    dataDir(name::String,dirname::String = "data")

Concatenates string with name of data directory
"""
function dataDir(name::String,dirname::String = "data")
    return joinpath(dirname,name)
end

"""
    γIntensityPlot(times::Array,γint::Array,prefix::String)

Plots average photon counts and saves to file. Returns relative path to resulting plot.
"""
function γIntensityPlot(times::Array,γint::Array,prefix::String)
    γplotName = string(prefix,"avg-photons-vs-time.svg")
    γplot = plot(times,γint,label=false)
    xlabel!("time (ns)")
    ylabel!("Avg photon counts")
    title!("Average Photon Counts vs Time")
    savefig(γplot,γplotName)
    @info "Saved $γplotName"
    return γplotName
end

export γIntensityPlot

"""
    γCountPlot(times::Array,γint::Array,prefix::String)

Plots 'measured' photon counts and saves to file. Returns relative path to resulting plot.
"""
function γCountPlot(times::Array,γcount::Array, γint::Array,prefix::String)
    γplotName = string(prefix,"photon-counts-vs-time.svg")
    γplot = bar(times,γcount,label="Simulated Counts")
    sig1 = sqrt.(γint)
    sig1low = map(x -> x < 1 ? x : sqrt(x),γint)
    plot!(γplot,times,γint,label = "Average Counts"; ribbon = (sig1low,sig1))
    xlabel!("time (ns)")
    ylabel!("photon counts")
    title!("Photon Counts vs Time")
    savefig(γplot,γplotName)
    @info "Saved $γplotName"
    return γplotName
end

export γCountPlot
"""
    γCorrTimePlot(timeVec::Vector, g2Vec::Vector, params::Dict, prefix::String)

Saves plot of photon correlation time series
"""
function γCorrTimePlot(timeDF::DataFrame, params::Dict, prefix::String)
    out = []
    γplotName = string(prefix,"time-domain-photon-correlation.svg")
    inzcorr  = timeDF[!,:corr1] .> 0
    nzcorr = timeDF[!,:corr1][inzcorr]
    nztime = timeDF[!,:time][inzcorr]
	γplot = plot(nztime,nzcorr,label = false)
    xlabel!(L"\tau \textrm{ (ns)}")
	ylabel!("\$g^{(2)}(\\tau)\$")
    title!("Photon Correlations vs Time Offset")
    savefig(γplot,γplotName)
    push!(out,γplotName)
    @info "Saved $(γplotName)"
    if params[:repeat] > 1
        γplotName = string(prefix,"time-domain-photon-correlation-sum.svg")
        nzcorr = timeDF[!,:sum][inzcorr]
        γplot = plot(nztime,nzcorr,label = false)
        xlabel!(L"\tau \textrm{ (ns)}")
        ylabel!("\$\\sum_{i=1}^{$(params[:repeat])}g_i^{(2)}(\\tau)\$")
        title!("Sum of Photon Correlations vs Time Offset")
        savefig(γplot,γplotName)
        push!(out,γplotName)
        @info "Saved $(γplotName)"
    end
    return out
end

export γCorrTimePlot

"""
    γCorrFreqPlot(freqVec::Vector, fftVec::Vector, params::Dict, prefix::String)

Saves plot of photon correlation Fourier transform
"""
function γCorrFreqPlot(freqDF::DataFrame, params::Dict, prefix::String)
    out = []
    γplotName = string(prefix,"frequency-domain-photon-correlation.svg")
	γplot = plot(freqDF[!,:freq],freqDF[!,:corr1],label = false)
	xlabel!("frequency (GHz)")
	ylabel!("\$\\hat{g}^{(2)}(\\nu)\$")
	title!("Photon Correlations vs Frequency")
    if length(params[:νm]) > 1
        shifts = map(x->abs(x[2]-x[1]),subsets(params[:νm],2))
        shiftname = length(params[:νm]) > 2 ? "frequency shifts" : "frequency shift"
        vline!(shifts,label=shiftname,ls=:dash)
    end
    savefig(γplot,γplotName)
    push!(out,γplotName)
    @info "Saved $(γplotName)"
    if params[:repeat] > 1
        γplotName = string(prefix,"frequency-domain-photon-correlation-sum.svg")
        γplot = plot(freqDF[!,:freq],freqDF[!,:sum],label = false)
        xlabel!("frequency (GHz)")
        ylabel!("\$\\sum_{i=1}^{$(params[:repeat])}\\hat{g}^{(2)}(\\nu)\$")
        title!("Sum of Photon Correlations vs Frequency")
        if length(params[:νm]) > 1
            shifts = map(x->abs(x[2]-x[1]),subsets(params[:νm],2))
            shiftname = length(params[:νm]) > 2 ? "frequency shifts" : "frequency shift"
            vline!(shifts,label=shiftname,ls=:dash)
        end
        savefig(γplot,γplotName)
        push!(out,γplotName)
        @info "Saved $(γplotName)"
    end
    return out
end

export γCorrFreqPlot


##################################################################################
# This file is part of ModelBaseEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

abstract type AbstractEquation end

# equations must have these fields: expr, vinds, vsyms, eval_resid, eval_RJ
for fn in (:expr, :vinds, :vsyms, :eval_resid, :eval_RJ)
    local qnfn = QuoteNode(fn)
    quote
        @inline $fn(eqn::AbstractEquation) = getfield(eqn, $qnfn)
    end |> eval
end

const default_eqn_type = :any

# equations might have these fields. If not, we provide defaults
@inline type(eqn::AbstractEquation) = :type in fieldnames(typeof(eqn)) ? getfield(eqn, :type) : default_eqn_type
@inline doc(eqn::AbstractEquation) = :doc in fieldnames(typeof(eqn)) ? getfield(eqn, :doc) : ""

@inline eval_resid(eqn::AbstractEquation, x) = eval_resid(eqn)(x)
@inline eval_RJ(eqn::AbstractEquation, x) = eval_RJ(eqn)(x)

#
function Base.show(io::IO, eqn::AbstractEquation)
    if get(io, :compact, false)
        return print(io, expr(eqn))
    end
    docstr = isempty(doc(eqn)) ? "" : "\"$(doc(eqn))\" "
    typestr = type(eqn) === default_eqn_type ? "" : "$(type(eqn)) "
    print(io, docstr, typestr, expr(eqn))
end

Base.:(==)(e1::AbstractEquation, e2::AbstractEquation) = type(e1) == type(e1) && expr(e1) == expr(e2)
Base.hash(e::AbstractEquation, h::UInt) = hash((type(e), expr(e)), h)

abstract type AbstractModel end

# a subtype of AbstractModel is expected to have a number of fields.
# If it doesn't, the creater of the new model type must define the
# access methods that follow.

@inline variables(m::AM) where AM <: AbstractModel = getfield(m, :variables)
@inline nvariables(m::AM) where AM <: AbstractModel = length(variables(m))

@inline shocks(m::AM) where AM <: AbstractModel = getfield(m, :shocks)
@inline nshocks(m::AM) where AM <: AbstractModel = length(shocks(m))

@inline allvars(m::AM) where AM <: AbstractModel = vcat(variables(m), shocks(m))
@inline nallvars(m::AM) where AM <: AbstractModel = length(variables(m)) + length(shocks(m))

@inline sstate(m::AM) where AM <: AbstractModel = getfield(m, :sstate)

@inline parameters(m::AM) where AM <: AbstractModel = getfield(m, :parameters)

@inline equations(m::AM) where AM <: AbstractModel = getfield(m, :equations)
@inline nequations(m::AM) where AM <: AbstractModel = length(equations(m))

@inline alleqns(m::AM) where AM <: AbstractModel = getfield(m, :equations)
@inline nalleqns(m::AM) where AM <: AbstractModel = length(equations(m))

export parameters
export variables, nvariables
export shocks, nshocks
export equations, nequations
export sstate
#######

# @inline moduleof(f::Function) = parentmodule(f)
@inline moduleof(e::AbstractEquation) = parentmodule(eval_resid(e))
function moduleof(m::AbstractModel)
    eqns = equations(m)
    if isempty(eqns)
        error("Unable to determine the module containing the given model. Try adding equations to it and call `@initialize`.")
    end
    return moduleof(first(eqns))
end


"""
```
Species(id, [value=0])
```

Define a `Species` with a `name` and initial population `value`.
"""
struct Species
  identifier :: Symbol
  population :: Int

  function Species(id, value::Int=0)
    if value < 0
      error("Species population must be nonnegative.")
    end
    new(Symbol(id), value)
  end
end

function Base.show(io::IO, x::Species)
  println(io, x.population)
end


"""
`DPGopts(m;σβ=1.,αΘ=0.0001,αw=0.001,αv=0.001,αu=0.001,γ=0.99,τ=0.001,iters=20_000, critic_update=:gradient,λrls=0.999,stepreduce_interval=1000,stepreduce_factor=0.995,hold_actor=1000)`

Structure with options to the DMP

# Fields

`σβ` Exploration noise covariance\n
`αΘ::Float64` Actor step size\n
`αw::Float64` Q-function step size 1\n
`αv::Float64` Q-function step size 2\n
`αu::Float64` Currently not used\n
`γ::Float64` Discount factor\n
`τ::Float64` Tracking factor between target and training networks\n
`λ::Float64` Regularization parameter for leastsquares Q-function fitting\n
`iters::Int64` Number of iterations to run\n
`m::Int64` Action dimension\n
`critic_update::Symbol` How to update the critic, can be chosen as `:gradient`, `:rls`, `:kalman`, `:leastsquares`\n
`λrls::Float64` If rls is used for critic update, use this forgetting factor\n
`stepreduce_interval::Int` The stepsize is reduced with this interval\n
`stepreduce_factor::Float64` The stepsize is reduced with this factor\n
`hold_actor::Int` Keep the actor from being updated for a few iterations in the beginning to allow the critic to obtain reasonable values\n
`experience_replay::Int` How many transitions to store in the replay buffer\n
`experience_ratio::Int` How many stored transitions to train with for every new experience\n
`momentum::Float64` Momentum term in actor gradient ascent (0.9)\n
`rmsprop::Bool` Use rmsprop for gradient training\n
`rmspropfactor::Float64` How much rmsprop, (0.9)\n
`critic_update_interval::Int` If critic is trained using leastsquares, only relearn the critic between this many mini-batches\n
`eval_interval::Int` With this interval the policy is evaluated without noise, parameters saved and stepsizes are reduced if divergence is detected.\n
`divergence_threshold::Float64` If the cost is higher than `divergence_threshold*bestcost`, the parameters are reset to their previously best values and stepsizes are reduced.\n
See example file or the paper Deterministic Policy Gradient Algorithms, Silver et al. 2014
"""
immutable DPGopts
    σβ
    αΘ::Float64
    αw::Float64
    αv::Float64
    αu::Float64
    γ::Float64
    τ::Float64
    λ::Float64
    iters::Int64
    m::Int64
    critic_update::Symbol
    λrls::Float64
    stepreduce_interval::Int
    stepreduce_factor::Float64
    hold_actor::Int
    experience_replay::Int
    experience_ratio::Int
    momentum::Float64
    rmsprop::Bool
    rmspropfactor::Float64
    critic_update_interval::Int
    eval_interval::Int
    divergence_threshold::Float64
end

DPGopts(m;
σβ=1.,
αΘ=0.001,
αw=0.01,
αv=0.01,
αu=0.001,
γ=0.99,
τ=0.001,
λ=1e-4,
iters=20_000,
critic_update=:gradient,
λrls=0.999,
stepreduce_interval=1000,
stepreduce_factor=0.995,
hold_actor=1000,
experience_replay=0,
experience_ratio=10,
momentum=0.9,
rmsprop=true,
rmspropfactor=0.9,
critic_update_interval=1,
eval_interval=10,
divergence_threshold=1.5) =
DPGopts(σβ,αΘ,αw,αv,αu,γ,τ,λ,iters,m,critic_update,λrls,stepreduce_interval,stepreduce_factor,hold_actor,experience_replay,experience_ratio,momentum,rmsprop,rmspropfactor,critic_update_interval,eval_interval,divergence_threshold)

"""
Structure with functions to pass to the DMP

# Fields

`μ,Q,gradients,simulate,exploration,reward`

See example file or the paper by Silver et al. 2014
"""
immutable DPGfuns
    μ::Function
    Q::Function
    gradients::Function
    simulate::Function
    exploration::Function
    reward::Function
end

"""
Structure which contains the parameters of the DPG optimization problem\n
`Θ` parameters in the actor\n
`w` parameters in the Q-function\n
`v` parameters in the Q-function\n
All parameters should be a subtype of AbstractVector
A typical Q-function looks like `Q = (∇μ(s)*(a-μ(s)))'w + V(s,v)`

"""
type DPGstate{T1<:AbstractVector,T2<:AbstractVector,T3<:AbstractVector}
    Θ::T1
    w::T2
    v::T3
end

function J(x,a,r)
    cost = @parallel (+) for t = 1:size(x,1)
    r(x[t,:][:],a[t,:][:],t)
end
-cost
end

# """
# `cost, Θ, w, v = dpg(opts, funs, state0, x0)`
#
# Main function.
#
# # Arguments
# `opts::DPGopts` structure with options and parameters\n
# `funs::DPGfuns` structure with functions\n
# `state0::DPGstate` initial parameters
# `x0` initial system state
# """
function dpg(opts, funs, state0, x0,C, progressfun = (Θ,w,v,C,i,x,uout,cost)->0)
    println("=== Deterministic Policy Gradient ===")
    # Expand input structs
    σβ          = opts.σβ
    αΘ          = opts.αΘ
    αw          = opts.αw
    αv          = opts.αv
    αu          = opts.αu
    γ           = opts.γ
    τ           = opts.τ
    iters       = opts.iters
    m           = opts.m
    n           = length(x0)
    critic_update= opts.critic_update
    λrls        = opts.λrls
    mem         = SequentialReplayMemory(opts.experience_replay)
    μ           = funs.μ
    Q           = funs.Q
    gradients   = funs.gradients
    simulate    = funs.simulate
    exploration = funs.exploration
    r           = funs.reward
    println("Training using $critic_update")

    # Initialize parameters
    Θ           = deepcopy(state0.Θ) # Weights
    w           = deepcopy(state0.w)
    v           = deepcopy(state0.v)
    Θt          = deepcopy(Θ) # Tracking weights
    wt          = deepcopy(w)
    vt          = deepcopy(v)
    Pw          = size(Θ,1)
    Pv          = size(v,1)
    Θb          = deepcopy(Θ) # Best weights
    wb          = deepcopy(w)
    vb          = deepcopy(v)
    dΘs         = 1000ones(Pw) # Weight gradient states
    dΘs2        = 0*ones(Pw)
    dws         = 100ones(Pw)
    dvs         = 100ones(Pv)
    cost        = zeros(iters)
    bestcost    = Inf
    x,uout      = simulate(Θ, x0)
    T           = size(x,1)

    # TODO: Make the parameters below part of the options
    if critic_update == :rls
        Pvw = 10eye(Pw+Pv)
    elseif critic_update == :kalman
        Pk = 1000eye(Pw+Pv)
        R2 = 1
        R12 = zeros(Pw+Pv)
    end

    s = zeros(n)

    function train!(tvec, update_actor,update_critic=true)
        dΘ,dw,dv = zeros(Θ),zeros(w),zeros(v)
        if critic_update ==:leastsquares && update_critic
            basis = sample_uniform!(mem,Pv)
            for (it,t) in enumerate(basis)
                C[:,it] = t.s
            end
            push!(mem,basis)
            A = Matrix{Float64}(length(mem),Pv+Pw)
            targets = Vector{Float64}(length(mem))
            for (it,t) in enumerate(mem)
                a1          = μ(t.s1,Θ,t.t)
                _, ∇wQ, ∇vQ, _ = gradients(t.s1,t.s,a1,t.a,Θ,w,v,t.t,C)
                y           = t.r + γ * Q(t.s1,a1,vt,wt,Θt,t.t,C)
                targets[it] = y
                # t.δ         = (y - Q(t.s,t.a,v,w,Θ,t.t))[1] # Not needed for batch learning
                A[it,1:Pv]  = ∇vQ
                A[it,Pv+1:end]  = ∇wQ
            end
            vw = [A; opts.λ*eye(Pv+Pw)]\[targets+100;zeros(Pv+Pw)] #TODO magic numer 100?
            v[:],w[:] = vw[1:Pv],vw[Pv+1:end]
        end
        for t in tvec
            a1          = μ(t.s1,Θt,t.t)
            ∇aQ, ∇wQ,∇vQ, ∇μ = gradients(t.s1,t.s,a1,t.a,Θ,w,v,t.t,C)
            y           = t.r + γ * Q(t.s1,a1,vt,wt,Θt,t.t,C)
            t.δ         = (y - Q(t.s,t.a,v,w,Θ,t.t,C))[1]
            if critic_update == :rls
                vw,Pvw      = RLS([v;w], y, [∇vQ;∇wQ], Pvw, λrls)
                v[:],w[:]   = vw[1:Pv],vw[Pv+1:end]
            elseif critic_update == :kalman

                Φ = [∇vQ;∇wQ]
                # R1 = ΦΦ', to only update covariance in the direction of incoming data
                vw,Pk = kalman(Φ*Φ',R2,R12,[v;w], y, Φ, Pk)
                v[:],w[:] = vw[1:Pv],vw[Pv+1:end]
            elseif critic_update ==:gradient
                dw += t.δ * ∇wQ  #- γ * ϕ(s1,a1) * ϕu
                dv += t.δ * ∇vQ   #- γ * ϕ(s1) * ϕu
            end
            dΘ += ∇μ*∇aQ
        end
        dΘ./= length(tvec)
        # RMS prop update parameters (gradient divided by running average of RMS gradient, see. http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
        dΘs[:] = opts.rmspropfactor*dΘs + (1-opts.rmspropfactor)*dΘ.^2
        if update_actor
            ΔΘ = αΘ * dΘ
            if opts.rmsprop
                ΔΘ ./= sqrt(dΘs+1e-10) # RMSprop
            end
            # ΔΘ = dΘ./sqrt(dΘs+1e-10).*sqrt(dΘs2) # ADAdelta
            # dΘs2[:] = 0.9dΘs2 + 0.1ΔΘ.^2 # ADAdelta
            dΘs2 = opts.momentum*dΘs2 + ΔΘ # Momentum + RMSProp
            Θ[:]  += dΘs2
        end

        if critic_update == :gradient
            Δw = αw/T * dw
            Δv = αv/T * dv
            dws[:] = opts.rmspropfactor*dws + (1-opts.rmspropfactor)*dw.^2
            dvs[:] = opts.rmspropfactor*dvs + (1-opts.rmspropfactor)*dv.^2
            if opts.rmsprop
                Δw ./= sqrt(dws)+1e-10
                Δv ./= sqrt(dvs)+1e-10
            end
            w += Δw
            v += Δv
        end

        # Update tracking networks
        Θt[:], wt[:], vt[:] = τ*Θ + (1-τ)*Θt, τ*w + (1-τ)*wt, τ*v + (1-τ)*vt
        nothing
    end

    # Main loop ================================================================
    for i = 1:iters
        x0i         = x0 #+ 2randn(n) # TODO: this should not be hard coded
        noise       = exploration(σβ)
        x,uout      = simulate(Θ, x0i, noise)
        batch = Vector{Transition}(T-1)
        for ti = 1:T-1
            s1          = x[ti+1,:][:]
            s           = x[ti,:][:]
            a           = uout[ti,:][:]
            ri          = r(s1,a,ti)
            cost[i]    -= ri
            trans       = Transition(s,s1,a,ri,0.,ti,i)
            batch[ti]   = trans
            opts.experience_replay > 0 && push!(mem, trans)
        end
        train!(batch, i > opts.hold_actor)

        if i % opts.stepreduce_interval == 0
            αΘ  *= opts.stepreduce_factor
            αw  *= opts.stepreduce_factor
            αv  *= opts.stepreduce_factor
        end

        if opts.experience_replay > 0 && i > 10
            for ei = 1:opts.experience_ratio
                trans = sample_uniform!(mem,T-1)
                train!(trans, i > opts.hold_actor, ei-1 % opts.critic_update_interval == 0)
                push!(mem, trans)
            end
            ((i % 50) == 0) && sort!(mem)
        end

        if (i-1) % opts.eval_interval == 0 # Simulate without noise and evaluate cost
            x,uout = simulate(Θ, x0) # Evaluate using latest parameters and possibly revert back to tracking parameters
            cost[i] = J(x,uout,r)
            progressfun(Θ,w,v,C,i,x,uout, cost)
            if critic_update == :gradient
                println(i, ", cost: ", cost[i] |> r5, " norm ∇Θ: ", Σ½(dΘs) |> r5, " norm ∇w: ", Σ½(dws) |> r5, " norm ∇v: ", Σ½(dvs) |> r5)#, " trace(P): ", trace(Pvw) |> r5)
            else
                println(i, ", cost: ", cost[i] |> r5, " norm ∇Θ: ", Σ½(dΘs) |> r5)
            end
            if cost[i] < bestcost
                bestcost = cost[i]
                # TODO: changed to saveing tracking networks, to be more likely to escape local minima
                Θb = deepcopy(Θ)
                wb = deepcopy(w)
                vb = deepcopy(v)

            elseif cost[i] > opts.divergence_threshold*bestcost
                print_with_color(:orange,"Reducing stepsizes due to divergence\n")
                αΘ  /= 5
                αw  /= 5
                αv  /= 5
                αu  /= 5
                σβ  /= 2
                Θ, w, v = deepcopy(Θb), deepcopy(wb), deepcopy(vb) # reset parameters
            end
        end

    end
    println("Done. Minimum cost: $(minimum(cost[1:opts.eval_interval:end])), ($(minimum(cost)))")
    return cost, Θb, wb, vb, mem # Select the parameters with lowest cost
end