__precompile__()
module Reduce
using ForceImport, SyntaxTree, LinearAlgebra
!isdefined(Reduce,:linefilter!) && (linefilter! = SyntaxTree.linefilter)

#   This file is part of Reduce.jl. It is licensed under the MIT license
#   Copyright (C) 2017 Michael Reed

include(joinpath(@__DIR__,"../deps/svn.jl"))

struct PSL <: Base.AbstractPipe
    input::Pipe
    output::Pipe
    process::Base.Process
    function PSL()
        # Setup pipes and reduce process
        input = Pipe()
        output = Pipe()
        rsl = red()
        dirf = joinpath(@__DIR__,"..","deps")
        if !Sys.iswindows()
            try
                process = _spawn(red(), input, output)
            catch
                process = _spawn(redsys(dirf), input, output)
            end
        else
            process = _spawn(redsys(dirf), input, output)
        end
        # Close the unneeded ends of Pipes
        close(input.out)
        close(output.in)
        return new(input, output, process)
    end
end

Base.kill(rs::PSL) = kill(rs.process)
Base.process_exited(rs::PSL) = process_exited(rs.process)

export error, ReduceError
import Base: error

struct ReduceError <: Exception
    errstr::String
end

Base.showerror(io::IO, err::ReduceError) = print(io,"Reduce: "*chomp(err.errstr))

function ReduceCheck(output) # check for REDUCE errors
    occursin(r"(([*]{5})|([+]{3}) )|( ?  \(Y or N\))",output) && throw(ReduceError(output))
end

function ReduceWarn(output) # check for REDUCE warnings
    if occursin(r"[*]{3}",output)
        @info "REDUCE: $(chomp(chomp(output)))"
        join(split(output,r"[*]{3}.*\n"))
    else
        output
    end
end

function PipeClogged(tf::Bool,c::Int,info::String)
    @warn "Reduce pipe clogged by $info, $(tf ? "success" : "failure") after $c tries"
end

clear(rs::PSL) = (write(rs.input,";\n"); readavailable(rs.output))
clears = (()->(c=true; return (tf=c)->(c≠tf && (c=tf); return c)))()

const EOT = Char(4) # end of transmission character
EOTstr = "symbolic write(int2id $(Int(EOT)))"

function Base.write(rs::PSL, input::String)
    clears() && clear(rs)
    write(rs.input,"$input; $EOTstr;\n")
end

const SOS = "[0-9]+: " # REDUCE terminal prompt
const RES = Regex("\n($EOT\n$SOS)|(\n$SOS\n$EOT)|(\n$SOS$EOT\n)|($EOT\n)")

function Base.read(rs::PSL) # get result and strip prompts/EOT char
    out = String(readuntil(rs.output,EOT))*String(readavailable(rs.output))
    Sys.iswindows() && (out = replace(out,r"\r" => ""))
    out = replace(replace(out,r"\$\n\n" => "\n\n"),RES=>"")
    out = replace(out,Regex(SOS) => "") |> chomp |> chomp |> join
    ReduceCheck(out)
    return ReduceWarn(out)
end

readsp(rs::PSL) = split(read(rs),"\n\n\n")

include("rexpr.jl") # load RExpr features
include("parser.jl") # load parser generator
include("repl.jl") # load repl features
include("switch.jl") # load switch operators

module Algebra
using Reduce, LinearAlgebra
import AbstractTensors
import AbstractTensors: conj, inv, dot, PROD, SUM, -, /, √
import AbstractTensors: sqrt, abs, exp, expm1, log, log1p, sin, cos, sinh, cosh, ^
const *,+ = PROD,SUM

export expm1, log1p, √
for T ∈ (RExpr,Expr,Symbol,String)
    @eval begin
        Base.expm1(z::$T) = exp(z)-1
        Base.log1p(z::$T) = 1+log(z)
        dot(a::$T,b) = conj(a)*b
        dot(a,b::$T) = conj(a)*b
        dot(a::$T,b::$T) = conj(a)*b
    end
end
dot(a::Expr,b::Symbol) = conj(a)*b
dot(a::Symbol,b::Expr) = conj(a)*b

include("unary.jl") # load unary operators
include("args.jl") # load calculus operators

function init_subtype(name)
    Expr(:block,[:(Base.$i(r::$name...)=$i(RExpr.(r)...)|>$name) for i ∈ [alg;iops]]...) |> eval
    Expr(:block,[:($i(r::$name...)=$i(RExpr.(r)...)|>$name) for i ∈ [calculus;cnan;cmat]]...) |> eval
    Expr(:block,[:(Base.$i(r::$name)=$i(RExpr(r))|>$name) for i ∈ [sbas;[:length]]]...) |> eval
    Expr(:block,[:($i(r::$name)=$i(RExpr(r))|>$name) for i ∈ [sfun;snan;snum;scom;sint;sran;smat]]...) |> eval
    :(Base.promote(r::$name,x) = (r,$name(x))) |> eval
    :(Base.promote(x,r::$name) = (r,$name(x))) |> eval
    :(Base.signbit(r::$name) = false) |> eval
end

const variables = [
    :root_multiplicities,
    :requirements,
    :assumptions,
    :low_pow,
    :high_pow,
    :euler_gamma,
    :khinchin,
]

for var ∈ [variables;[:ws]]
    :($var() = rcall(RExpr($(string(var)))) |> parse) |> eval
end

function bypass(op::Symbol,T::ExprSymbol,args)
    1 ∈ args && (@eval $op(x::X) where X<:$T = Base.$op(x))
    if 2 ∈ args
        @eval begin
            $op(a::A,b::B) where {A<:ExprSymbol,B<:$T} = Base.$op(a,b)
            $op(a::A,b::B) where {A<:$T,B<:ExprSymbol} = Base.$op(a,b)
        end
    end
end

import AbstractTensors: TensorAlgebra, value
for op ∈ (:exp,:reverse)
    bypass(op,:TensorAlgebra,(1,))
end
for op ∈ (:+,:*,:/)
    bypass(op,:TensorAlgebra,(2,))
end
@eval bypass(:-,:TensorAlgebra,(1,2))
rank(x::X) where X<:TensorAlgebra = LinearAlgebra.rank(x)

@doc """
    ws()

Several mechanisms are available for saving and retrieving previously evaluated expressions. The simplest of these refers to the last algebraic expression simplified. When an assignment of an algebraic expression is made, or an expression is evaluated at the top level, (i.e., not inside a compound statement or procedure) the results of the evaluation are automatically saved in a variable `ws` that we shall refer to as the workspace. (More precisely, the expression is assigned to the variable `ws` that is then available for further manipulation.)

*Example:* If we evaluate the expression `(x+y)^2` at the top level and next wish to differentiate it with respect to `y`, we can simply say
``
julia> Algebra.df(Algebra.ws(),:y)
```
to get the desired answer.
""" Reduce.Algebra.ws

@doc """
    root_multiplicities()

Solution multiplicities are stored in the global variable `root_multiplicities` rather than the solution list. The value of this variable is a list of the multiplicities of the solutions for the last call of `solve`. For example,
```Julia
julia> Algebra.solve(:(x^2==2x-1),:x); Algebra.root_multiplicities()
```
gives the results
```Julia
(:(x = 1),)

(2,)
```
""" Reduce.Algebra.root_multiplicities

@doc """
    requirements()

The proper design of a variable sequence supplied as a second argument to `solve` is important for the structure of the solution of an equation system. Any unknown in the system not in this list is considered totally free. E.g. the call
```Julia
Algebra.solve((:(x==2z),:(z==2y)),(:z,))
```
produces an empty list as a result because there is no function `z = z(x,y)` which fulfills both equations for arbitrary `x` and `y` values. In such a case the share variable `requirements` displays a set of restrictions for the parameters of the system:

```Julia
julia> Algebra.requirements()
(:(x - 4y),)
```

The non-existence of a formal solution is caused by a contradiction which disappears only if the parameters of the initial system are set such that all members of the requirements list take the value zero. For a linear system the set is complete: a solution of the requirements list makes the initial system solvable. E.g. in the above case a substitution ``x = 4y`` makes the equation set consistent. For a non-linear system only one inconsistency is detected. If such a system has more than one inconsistency, you must reduce them one after the other. 1 The set shows you also the dependency among the parameters: here one of ``x`` and ``y`` is free and a formal solution of the system can be computed by adding it to the variable list of `solve`. The requirement set is not unique – there may be other such sets.
""" Reduce.Algebra.requirements

@doc """
    assumptions()

A system with parameters may have a formal solution, e.g.
```Julia
julia> Algebra.solve((:(x==a*z+1),:(0==b*z-y)),(:z,:x))
(:(z = y // b), :(x = (a * y + b) // b))
```
which is not valid for all possible values of the parameters. The variable `assumptions` contains then a list of restrictions: the solutions are valid only as long as none of these expressions vanishes. Any zero of one of them represents a special case that is not covered by the formal solution. In the above case the value is
```Julia
julia> Algebra.assumptions()
(:b,)
```
which excludes formally the case ``b = 0``; obviously this special parameter value makes the system singular. The set of assumptions is complete for both, linear and non–linear systems.
""" Reduce.Algebra.assumptions

@doc """
    euler_gamma()

Euler's constant, also available as -\$\\psi(1)\$.
""" Reduce.Algebra.euler_gamma

@doc """
    khinchin()

Khinchin's constant, defined as

\$ \\prod_{n=1}^\\infty \\left( 1 + \\frac{1}{n(n+2)} \\right)^{\\log_2 n}. \$
""" Reduce.Algebra.khinchin

end

export Algebra, @force, @vars

macro vars(syms...)
    Expr(:block,[Expr(:(=),esc(sym),QuoteNode(sym)) for sym ∈ syms]...)
end

Base.write(rs::PSL,r::RExpr) = write(rs,convert(String,r))

import Base: zero, one

for T ∈ [:Any,:Expr,:Symbol]
    @eval begin
        zero(::Type{$T}) = 0
        zero(::$T) = 0
        one(::Type{$T}) = 1
        one(::$T) = 1
    end
end

zero(::Type{RExpr}) = R"0"
zero(::RExpr) = R"0"
one(::Type{RExpr}) = R"1"
one(::RExpr) = R"1"

import LinearAlgebra: transpose, adjoint

transpose(r::ExprSymbol) = r
transpose(r::RExpr) = r
adjoint(r::ExprSymbol) = Algebra.conj(r)
adjoint(r::RExpr) = Algebra.conj(r)

## Setup

const offlist = [:nat,:latex,:exp]

export load_package, @load_package

"""
    load_package(::Symbol)

Loads the specified package into REDUCE

## Examples
```julia-repl
julia> load_package(:rlfi)
```
"""
function load_package(pkg::Union{String,Symbol},pkgs...)
    "load_package $pkg" |> rcall
    for extra in pkgs
        load_package(extra)
    end
    return nothing
end
function load_package(pkgs::Union{Array{String,1},Array{Symbol,1}})
    for pkg in pkgs
        load_package(pkg)
    end
end

macro load_package(pkg...)
    load_package(pkg...)
end

"""
    Reduce.Reset()

Kills the REDUCE process and starts a new instance.

## Examples
```julia-repl
julia> Reduce.Reset()
Reduce (Free PSL version, revision 4015),  5-May-2017 ...
```
"""
Reset() = (kill(rs); Load())
__init__() = (Load(); repl_init(); atexit(() -> kill(rs)))

# Server setup

function Load()
    global rs = PSL()

    offs = ""
    for o in offlist
        o != :nat && (offs = offs*"off $o; ")
    end
    write(rs.input,"off nat; $EOTstr;\n")
    banner = readuntil(rs.output,EOT) |> String
    readavailable(rs.output)
    rcsl = occursin(" CSL ",banner)
    if Sys.iswindows()
        banner = replace(banner,r"\r" => "")
    else
        ReduceCheck(banner)
    end
    global banner = split(String(banner),'\n')[rcsl ? 1 : end-3]
    load_package(:rlfi)
    offs |> RExpr |> rcall
    rcall(R"on savestructr")
    return nothing
end

global preload = false
try
    global preload
    (ENV["REDPRE"] ≠ "0") && (preload = true)
catch
end
preload && include("precomp.jl")

end # module


# This file is part of Kpax3. License is MIT.

function test_state_list_constructor()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile)

  x = Kpax3.AminoAcidData(settings)

  (m, n) = size(x.data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(x.data, settings.γ, settings.r)

  slist = Kpax3.AminoAcidStateList(x.data, [3; 3; 3; 1; 1; 2], priorR, priorC, settings)

  @test isa(slist.state, Vector{Kpax3.AminoAcidState})
  @test isa(slist.logpp, Vector{Float64})
  @test isa(slist.rank, Vector{Int})
  @test all(slist.logpp .< 0.0)
  @test length(unique(slist.logpp)) > 1
  @test all(0 .< slist.rank .< settings.popsize + 1)
  @test length(unique(slist.rank)) == settings.popsize
  @test all(diff(slist.logpp[slist.rank]) .<= 0)

  @test slist.state[1].R == [3; 3; 3; 1; 1; 2]

  for i in 1:settings.popsize
    t = Kpax3.AminoAcidState(x.data, slist.state[i].R, priorR, priorC, settings)

    l = t.cl[1:t.k]

    @test slist.state[i].R == t.R
    @test slist.state[i].C == t.C
    @test slist.state[i].emptycluster == t.emptycluster
    @test slist.state[i].cl == t.cl
    @test slist.state[i].k == t.k
    @test slist.state[i].v == t.v
    @test slist.state[i].n1s == t.n1s
    for g in l
      @test slist.state[i].unit[g][1:slist.state[i].v[g]]==t.unit[g][1:t.v[g]]
    end
    @test slist.state[i].logpR == t.logpR
    @test slist.state[i].logpC == t.logpC
    @test slist.state[i].loglik == t.loglik
    @test slist.state[i].logpp == t.logpp

    @test isapprox(slist.logpp[i], slist.state[i].logpp, atol=ε)
  end

  state = Kpax3.AminoAcidState(x.data, [3; 3; 3; 1; 1; 2], priorR, priorC, settings)
  slist = Kpax3.AminoAcidStateList(settings.popsize, state)

  @test isa(slist.state, Vector{Kpax3.AminoAcidState})
  @test isa(slist.logpp, Vector{Float64})
  @test isa(slist.rank, Vector{Int})
  @test all(slist.logpp .< 0.0)
  @test length(unique(slist.logpp)) == 1
  @test all(0 .< slist.rank .< settings.popsize + 1)
  @test length(unique(slist.rank)) == settings.popsize
  @test all(diff(slist.logpp[slist.rank]) .<= 0)
  @test slist.rank == Int[i for i in 1:settings.popsize]

  for i in 1:settings.popsize
    @test slist.state[i].R == state.R
    @test slist.state[i].C == state.C
    @test slist.state[i].emptycluster == state.emptycluster
    @test slist.state[i].cl == state.cl
    @test slist.state[i].k == state.k
    @test slist.state[i].v == state.v
    @test slist.state[i].n1s == state.n1s
    for g in 1:3
      @test (slist.state[i].unit[g][1:slist.state[i].v[g]] ==
             state.unit[g][1:state.v[g]])
    end
    @test slist.state[i].logpR == state.logpR
    @test slist.state[i].logpC == state.logpC
    @test slist.state[i].loglik == state.loglik
    @test slist.state[i].logpp == state.logpp
    @test slist.logpp[i] == state.logpp
  end

  nothing
end

test_state_list_constructor()

function test_state_list_copy_basic()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile)

  x = Kpax3.AminoAcidData(settings)

  (m, n) = size(x.data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(x.data, settings.γ, settings.r)

  s1 = Kpax3.AminoAcidStateList(x.data, [3; 3; 3; 1; 1; 2], priorR, priorC, settings)
  s2 = Kpax3.AminoAcidStateList(x.data, [1; 1; 1; 1; 1; 1], priorR, priorC, settings)

  Kpax3.copystatelist!(s2, s1, settings.popsize)

  for i in 1:settings.popsize
    j = s1.rank[i]

    r = 1:s1.state[j].k
    l = s1.state[j].cl[r]

    @test s2.state[i].R == s1.state[j].R
    @test s2.state[i].C[l, :] == s1.state[j].C[l, :]
    @test s2.state[i].emptycluster[r] == s1.state[j].emptycluster[r]
    @test s2.state[i].cl[r] == s1.state[j].cl[r]
    @test s2.state[i].k == s1.state[j].k
    @test s2.state[i].v[l] == s1.state[j].v[l]
    @test s2.state[i].n1s[l, :] == s1.state[j].n1s[l, :]
    for g in l
      @test (s2.state[i].unit[g][1:s2.state[i].v[g]] ==
             s1.state[j].unit[g][1:s1.state[j].v[g]])
    end
    @test s2.state[i].logpR == s1.state[j].logpR
    @test s2.state[i].logpC == s1.state[j].logpC
    @test s2.state[i].loglik == s1.state[j].loglik
    @test s2.state[i].logpp == s1.state[j].logpp

    @test s2.logpp[i] == s1.logpp[j]

    @test s2.rank[i] == i
  end

  nothing
end

test_state_list_copy_basic()


mutable struct GeneralDataType <: MyXMLElement
    el::XMLOrNothing
    states::Vector{Int}
    id::String

    function GeneralDataType(states::Vector{Int}, id::String)
        return new(nothing, states, id::String)
    end
end

function GeneralDataType(n::Int; id::String = "discreteStates")
    return GeneralDataType(collect(0:(n - 1)), id)
end

function name(::GeneralDataType)
    return bn.GENERAL_DATA_TYPE
end

function make_xml(gdt::GeneralDataType)
    el = new_element(name(gdt))
    set_attribute(el, bn.ID, gdt.id)

    for i = 1:length(gdt.states)
        state_el = new_child(el, bn.STATE)
        set_attribute(state_el, bn.CODE, gdt.states[i])
    end

    gdt.el = el
    return el
end

################################################################################
## aligments
################################################################################

const MISSING_SEQUENCE = 9

mutable struct Alignment <: MyXMLElement
    el::XMLOrNothing
    data_type::MyXMLElement
    taxa::Vector{<:AbstractString}
    sequences::Matrix{Int}
    id::String

    function Alignment(data_type::MyXMLElement,
                        taxa::Vector{<:AbstractString},
                        sequences::Matrix{Int};
                        id::String = bn.ALIGNMENT)
        @assert length(taxa) == size(sequences, 1)
        return new(nothing, data_type, taxa, sequences, id)
    end
end

function name(::Alignment)
    return bn.ALIGNMENT
end

function make_xml(al::Alignment)
    el = new_element(name(al))
    set_attribute(el, bn.ID, al.id)
    add_ref_el(el, al.data_type)
    for i = 1:size(al.sequences, 1)
        seq_el = new_child(el, bn.SEQUENCE)
        taxon_el = new_child(seq_el, bn.TAXON)
        set_attribute(taxon_el, bn.IDREF, al.taxa[i])
        seq_text = join(al.sequences[i, :], ' ')
        add_text(seq_el, seq_text)
    end

    al.el = el
    return el
end

function get_num_classes(alignment::Alignment)
    @unpack sequences = alignment
    p = size(alignment.sequences, 2)
    return length.([setdiff(unique(sequences[:, i]), MISSING_SEQUENCE) for i = 1:p])
end

################################################################################
## patterns
################################################################################

mutable struct Patterns <: MyXMLElement
    el::XMLOrNothing
    alignment::Alignment
    id::String

    function Patterns(alignment::Alignment; id::String = bn.PATTERNS)
        return new(nothing, alignment, id)
    end
end

function name(::Patterns)
    return bn.PATTERNS
end

function make_xml(patterns::Patterns)
    @unpack alignment, id = patterns

    el = new_element(name(patterns))
    set_attribute(el, bn.ID, id)
    set_attribute(el, bn.FROM, "1")
    set_attribute(el, bn.UNIQUE, bn.FALSE)
    add_ref_el(el, alignment)
    patterns.el = el

    return el
end

function get_num_classes(patterns::Patterns)
    return get_num_classes(patterns.alignment)
end

################################################################################
## likelihood
################################################################################

mutable struct OrderedLatentLiability <: MyXMLElement
    el::XMLOrNothing
    patterns::Patterns
    treeModel::TreeModelXMLElement
    threshold::CompoundParameter
    id::String
    trait_name::String

    function OrderedLatentLiability(patterns::Patterns,
                                    treeModel::TreeModelXMLElement,
                                    threshold::CompoundParameter;
                                    id::String = bn.ORDERED_LATENT_LIABILITY,
                                    trait_name::String = treeModel.param_names[1])

        return new(nothing, patterns, treeModel, threshold, id, trait_name)
    end
end

function OrderedLatentLiability(patterns::Patterns,
                                treeModel::TreeModelXMLElement;
                                id::String = bn.ORDERED_LATENT_LIABILITY,
                                trait_name::String = treeModel.param_names[1])
    n_states = get_num_classes(patterns)
    need_thresh = findall(x -> x > 2, n_states)

    params = Parameter[]
    if length(need_thresh) == 0
        push!(params, Parameter([0.0]))
    else
        for ind in need_thresh
            p_id = "$trait_name$ind.threshold"
            thresholds = default_thresholds(n_states[ind])
            push!(params, Parameter(thresholds, p_id, lower=0.0))
        end
    end

    cp = CompoundParameter(params, "$trait_name.threshold")
    return OrderedLatentLiability(patterns, treeModel, cp, id = id, trait_name = trait_name)

end

function name(::OrderedLatentLiability)
    return bn.ORDERED_LATENT_LIABILITY
end

function make_xml(oll::OrderedLatentLiability)
    @unpack patterns, treeModel, threshold, id = oll
    el = new_element(name(oll))
    add_ref_el(el, patterns)
    add_ref_el(el, treeModel)
    tip_el = new_child(el, bn.TIP_TRAIT)
    add_ref_el(tip_el, Parameter(Float64[], oll.trait_name))

    thresh_el = new_child(el, bn.THRESHOLD)
    add_child(thresh_el, make_xml(threshold))

    classes_el = new_child(el, bn.NUM_CLASSES)
    classes = get_num_classes(patterns)
    class_param = Parameter(Float64.(classes))
    add_child(classes_el, make_xml(class_param))

    set_attribute(el, bn.N_TRAITS, string(length(classes)))
    set_attribute(el, bn.N_DATA, "1")
    set_attribute(el, bn.ID, oll.id)

    oll.el = el
    return el
end

################################################################################
## operator
################################################################################

mutable struct ExtendedLatendLiabilityOperator <: OperatorXMLElement
    el::XMLOrNothing
    trait_likelihood::TraitLikelihoodXMLElement
    latent_liability::OrderedLatentLiability
    weight::Float64

    function ExtendedLatendLiabilityOperator(
            trait_likelihood::TraitLikelihoodXMLElement,
            latent_liability::OrderedLatentLiability;
            weight::Float64 = 1.0)

        return new(nothing, trait_likelihood, latent_liability, weight)
    end

end

function name(::ExtendedLatendLiabilityOperator)
    return bn.EXTENDED_GIBBS_OPERATOR
end

function make_xml(elo::ExtendedLatendLiabilityOperator)
    el = new_element(name(elo))

    set_attribute(el, bn.WEIGHT, elo.weight)
    add_ref_el(el, elo.trait_likelihood)
    add_ref_el(el, elo.latent_liability)
    elo.el = el
    return el
end

################################################################################
## helpers
################################################################################

function default_thresholds(states::Int)
    return convert.(Float64, collect(1:(states - 2)))
end

function find_state(x::Float64, thresholds::Vector{Float64})
    n = length(thresholds)

    if isnan(x)
        if n >= MISSING_SEQUENCE
            error("Cannot accomodate more than $MISSING_SEQUENCE patterns current missing code.")
        end
        return MISSING_SEQUENCE
    end

    state = 0
    for i = 1:n
        if x < thresholds[i]
            return i - 1
        end
    end
    return n
end

function data_to_alignment(data::Matrix{Float64}, states::Vector{Int};
            thresholds::Vector{Vector{Float64}} = default_thresholds.(states),
            bin_discrete::Bool = false)

    n, p = size(data)
    sequences = zeros(Int, n, p)

    for j = 1:p
        n_states = states[j]
        if n_states > 1
            if bin_discrete
                for i = 1:n
                    sequences[i, j] = find_state(data[i, j], thresholds[j])
                end
            else
                for i = 1:n
                    if isnan(data[i, j])
                        sequences[i, j] = MISSING_SEQUENCE
                    else
                        sequences[i, j] = Int(data[i, j])
                    end
                end
            end
        end
    end

    return sequences

end

################################################################################
## latent liability
################################################################################

function add_latent_liability(bx::BEASTXMLElement, trait_names::Vector{String}, states::Vector{Int})
    m = length(trait_names)
    gdt = GeneralDataType(maximum(states))
    ops_ind = find_element_ind(bx, OperatorsXMLElement)

    add_child(bx, gdt, ops_ind)

    for nm in trait_names
        latent_liablity = latent_liablity_elements(bx, nm, states, gdt)
    end
end

function add_latent_liability(bx::BEASTXMLElement, discrete_inds::Vector{Int})

    if length(discrete_inds) == 0
        return nothing
    end

    data_el = find_element(bx, DataXMLElement)
    @unpack data_mats = data_el
    ps = [size(x, 2) for x in data_mats]
    @assert length(unique(ps)) == 1
    p = ps[1]
    states = ones(Int, p)

    for data in data_el.data_mats
        for ind in discrete_inds
            ind_states = setdiff(unique(@view data[:, ind]), NaN)
            states[ind] = max(Int(maximum(ind_states)) + 1, states[ind])
        end
    end

    set_factor_precision_indices(bx, setdiff(1:p, discrete_inds))

    add_latent_liability(bx, data_el.trait_names, states)
end

function latent_liablity_elements(bx::BEASTXMLElement, trait_name::String, states::Vector{Int}, gdt::GeneralDataType)

    data_el = find_element(bx, DataXMLElement)
    treeModel = find_element(bx, TreeModelXMLElement)
    trait_likelihood = find_element(bx, TraitLikelihoodXMLElement)
    int_fac = trait_likelihood.extension_el

    if BEASTXMLConstructor.get_trait_name(int_fac) != trait_name
        flpd = find_element(bx, FactorLogPredictiveDensity)
        trait_likelihood = flpd.like
    end

    @assert BEASTXMLConstructor.get_trait_name(trait_likelihood.extension_el) == trait_name

    precs = get_precision(int_fac)
    set_value(precs, fill(2.0, length(precs)))

    taxa = data_el.taxa
    trait_ind = findfirst(isequal(trait_name), data_el.trait_names)
    data = data_el.data_mats[trait_ind]

    sequence = data_to_alignment(data, states)
    alignment = Alignment(gdt, taxa, sequence, id=trait_name * ".alignment")
    pattern = Patterns(alignment, id=trait_name * ".patterns")

    param_name = treeModel.param_names[findfirst(isequal(trait_name), treeModel.node_traits)]

    likelihood = OrderedLatentLiability(pattern, treeModel, id = trait_name * ".latentLiability", trait_name = param_name)

    operator = ExtendedLatendLiabilityOperator(trait_likelihood, likelihood)
    thresh_ops = MyXMLElement[]
    if (maximum(get_num_classes(likelihood.patterns)) > 2)
        for param in likelihood.threshold.parameters
            push!(thresh_ops, ScaleOperator(param))
        end
    end
    # threshold_operator = ScaleOperator(likelihood.threshold)

    ops_ind = find_element_ind(bx, OperatorsXMLElement)
    operators = bx.components[ops_ind]
    for el in [alignment, pattern, likelihood]
        add_child(bx, el, ops_ind)
        ops_ind += 1
    end

    push!(operators.els, operator)
    for op in thresh_ops
        push!(operators.els, op)
    end

    mcmc_el = find_element(bx, MCMCXMLElement)
    BEASTXMLConstructor.add_likelihood!(mcmc_el, likelihood)
    BEASTXMLConstructor.add_loggable(bx, likelihood.threshold, already_made = true)
    BEASTXMLConstructor.add_loggable(bx, Parameter(Float64[], param_name), already_made = true)

    return likelihood
end


function gather_bodies_initial_coordinates(simulation::NBodySimulation)
    system = simulation.system
    bodies = system.bodies;
    len = n = length(bodies)
    if simulation.thermostat isa NoseHooverThermostat
        len += 1
    end
    u0 = zeros(3, len)
    v0 = zeros(3, len)
    for i = 1:n
        u0[:, i] = bodies[i].r
        v0[:, i] = bodies[i].v
    end
    (u0, v0, n)
end

function gather_bodies_initial_coordinates(simulation::NBodySimulation{<:WaterSPCFw})
    system = simulation.system
    molecules = system.bodies
    n = length(molecules)
    len = 3*n
    if simulation.thermostat isa NoseHooverThermostat
        len = 3*n+1
    end
    (u0,v0) = gather_atom_coordinates(simulation.system, len)
    (u0, v0, n)
end

function gather_atom_coordinates(system::WaterSPCFw{<:MassBody}, len)
    molecules = system.bodies
    n = length(molecules)
    T = eltype(first(molecules).r)
    u0 = zeros(T,3,len)
    v0 = zeros(T,3,len)
    for i = 1:n
        p = system.scpfw_parameters
        indO, indH1, indH2 = 3*(i-1)+1, 3*(i-1)+2, 3*(i-1)+3
        u0[:, indO]  = molecules[i].r
        u0[:, indH1] = molecules[i].r .+ [p.rOH, 0.0, 0.0]
        u0[:, indH2] = molecules[i].r .+ [cos(p.aHOH) * p.rOH, 0.0, sin(p.aHOH) * p.rOH]
        v0[:, indO]  = molecules[i].v
        v0[:, indH1] = molecules[i].v
        v0[:, indH2] = molecules[i].v
    end
    (u0, v0)
end

function gather_atom_coordinates(system::WaterSPCFw{<:WaterMolecule}, len)
    molecules = system.bodies
    n = length(molecules)
    T = eltype(first(molecules).r)
    u0 = zeros(T,3,len)
    v0 = zeros(T,3,len)
    for i = 1:n
        p = system.scpfw_parameters
        indO, indH1, indH2 = 3*(i-1)+1, 3*(i-1)+2, 3*(i-1)+3
        u0[:, indO]  = molecules[i].O.r
        u0[:, indH1] = molecules[i].H1.r
        u0[:, indH2] = molecules[i].H2.r
        v0[:, indO]  = molecules[i].O.v
        v0[:, indH1] = molecules[i].H1.v
        v0[:, indH2] = molecules[i].H2.v
    end
    (u0, v0)
end

function gather_accelerations_for_potentials(simulation::NBodySimulation{CustomAccelerationSystem})
    acceleration_functions = []
    push!(acceleration_functions, simulation.system.acceleration)
    acceleration_functions
end

function gather_accelerations_for_potentials(simulation::NBodySimulation{<:PotentialNBodySystem})
    acceleration_functions = []
    for (potential, parameters) in simulation.system.potentials
        push!(acceleration_functions, get_accelerating_function(parameters, simulation))
    end
    acceleration_functions
end

function gather_group_accelerations(simulation::NBodySimulation{<:WaterSPCFw})
    accelerations = []
    push!(accelerations, get_group_accelerating_function(simulation.system.scpfw_parameters, simulation))
    accelerations
end

function get_accelerating_function(parameters::LennardJonesParameters, simulation::NBodySimulation)
    (ms, indxs) = obtain_data_for_lennard_jones_interaction(simulation.system)
    (dv, u, v, t, i) -> pairwise_lennard_jones_acceleration!(dv, u, i, indxs, ms, parameters, simulation.boundary_conditions)
end

function get_accelerating_function(parameters::ElectrostaticParameters, simulation::NBodySimulation)
    (qs, ms, indxs, exclude) = obtain_data_for_electrostatic_interaction(simulation.system)
    (dv, u, v, t, i) -> pairwise_electrostatic_acceleration!(dv, u, i, length(indxs), qs, ms, exclude, parameters, simulation.boundary_conditions)
end

function get_accelerating_function(parameters::SPCFwParameters, simulation::NBodySimulation)
    (ms, neighbouhood) = obtain_data_for_harmonic_bond_interaction(simulation.system, parameters)
    (dv, u, v, t, i) -> harmonic_bond_potential_acceleration!(dv, u, i, ms, neighbouhood, parameters)
end

function get_accelerating_function(parameters::GravitationalParameters, simulation::NBodySimulation)
    (dv, u, v, t, i) -> gravitational_acceleration!(dv, u, i, length(simulation.system.bodies), simulation.system.bodies, parameters)
end

function get_accelerating_function(parameters::MagnetostaticParameters, simulation::NBodySimulation)
    (dv, u, v, t, i) -> magnetostatic_dipdip_acceleration!(dv, u, i, length(simulation.system.bodies), simulation.system.bodies, parameters)
end

function get_group_accelerating_function(parameters::PotentialParameters, simulation::NBodySimulation{<:WaterSPCFw})
    (ms, indxs) = obtain_data_for_lennard_jones_interaction(simulation.system)
    (dv, u, v, t, i) -> valence_angle_potential_acceleration!(dv, u, 3*(i-1)+2, 3*(i-1)+1, 3*(i-1)+3, ms, parameters)
end

function obtain_data_for_harmonic_bond_interaction(system::WaterSPCFw, p::SPCFwParameters)
    neighbouhoods = Dict{Int, Vector{Tuple{Int,Float64}}}()
    n = length(system.bodies)
    ms = zeros(3 * n)
    for i in 1:n
        indO, indH1, indH2 = 3 * (i - 1) + 1, 3 * (i - 1) + 2, 3 * (i - 1) + 3
        ms[indO] = system.mO
        ms[indH1] = system.mH
        ms[indH2] = system.mH
        neighbours_o = Vector{Tuple{Int,Float64}}()
        push!(neighbours_o, (indH1, p.kb))
        push!(neighbours_o, (indH2, p.kb))
        neighbours_h1 = Vector{Tuple{Int,Float64}}()
        push!(neighbours_h1, (indO, p.kb))
        neighbours_h2 = Vector{Tuple{Int,Float64}}()
        push!(neighbours_h2, (indO, p.kb))
        neighbouhoods[indO] = neighbours_o
        neighbouhoods[indH1] = neighbours_h1
        neighbouhoods[indH2] = neighbours_h2
    end
    (ms, neighbouhoods)
end

function obtain_data_for_lennard_jones_interaction(system::PotentialNBodySystem)
    bodies = system.bodies
    n = length(bodies)
    ms = zeros(typeof(first(bodies).m), n)
    indxs = zeros(Int, n)
    for i = 1:n
        ms[i] = bodies[i].m
        indxs[i] = i
    end
    return (ms, indxs)
end

function obtain_data_for_lennard_jones_interaction(system::WaterSPCFw)
    bodies = system.bodies
    n = length(bodies)
    ms = zeros(3 * n)
    indxs = zeros(Int, n)
    for i = 1:n
        indxs[i] = 3 * (i - 1) + 1
        ms[3 * (i - 1) + 1] = system.mO
        ms[3 * (i - 1) + 2] = system.mH
        ms[3 * (i - 1) + 3] = system.mH
    end
    return (ms, indxs)
end

function obtain_data_for_electrostatic_interaction(system::PotentialNBodySystem)
    bodies = system.bodies
    n = length(bodies)
    qs = zeros(typeof(first(bodies).q), n)
    ms = zeros(typeof(first(bodies).m), n)
    indxs = collect(1:n)
    exclude = Dict{Int, Vector{Int}}()
    for i = 1:n
        qs[i] = bodies[i].q
        ms[i] = bodies[i].m
        exclude[i] = [i]
    end
    return (qs, ms, indxs, exclude)
end

function obtain_data_for_electrostatic_interaction(system::WaterSPCFw)
    bodies = system.bodies
    n = length(bodies)
    qs = zeros(3 * n)
    ms = zeros(3 * n)
    indxs = collect(1:3*n)
    exclude = Dict{Int, Vector{Int}}()
    for i = 1:n
        Oind = 3 * (i - 1) + 1
        qs[Oind] = system.qO
        qs[Oind + 1] = system.qH
        qs[Oind + 2] = system.qH
        ms[Oind] = system.mO
        ms[Oind + 1] = system.mH
        ms[Oind + 2] = system.mH
        exclude[Oind] = [Oind, Oind+1, Oind+2,3*n+1]
        exclude[Oind+1] = [Oind, Oind+1, Oind+2,3*n+1]
        exclude[Oind+2] = [Oind, Oind+1, Oind+2,3*n+1]
    end
    return (qs, ms, indxs, exclude)
end

function obtain_data_for_electrostatic_interaction(system::NBodySystem)
    bodies = system.bodies
    n = length(bodies)
    qs = zeros(typeof(first(bodies).m), n)
    ms = zeros(typeof(first(bodies).m), n)
    return (qs, ms)
end

function obtain_data_for_valence_angle_harmonic_interaction(system::WaterSPCFw)
    p = system.scpfw_parameters
    bonds = Vector{Tuple{Int, Int, Int, Float64, Float64}}()
    n = length(system.bodies)
    ms = zeros(3 * n)
    for i = 1:n
        indO, indH1, indH2 = 3 * (i - 1) + 1, 3 * (i - 1) + 2, 3 * (i - 1) + 3
        ms[indO] = system.mO
        ms[indH1] = system.mH
        ms[indH2] = system.mH
        push!(bonds, (indH1, indO, indH2, p.aHOH, p.ka))
    end
    return (ms, bonds)
end

function gather_simultaneous_acceleration(s::NBodySimulation)
    acelerations = []
    if s.thermostat isa BerendsenThermostat
        push!(acelerations, get_berendsen_thermostating_acceleration(s))
    elseif s.thermostat isa NoseHooverThermostat
        push!(acelerations, get_nosehoover_thermostating_acceleration(s))
    end
    acelerations
end

function get_berendsen_thermostating_acceleration(simulation::NBodySimulation)
    (ms, kb, n, nc, p) = obtain_data_for_berendsen_thermostating(simulation)
    (dv, u, v, t) -> berendsen_acceleration!(dv, v, ms, kb, n, nc, p)
end

function obtain_data_for_berendsen_thermostating(simulation::NBodySimulation)
    ms = get_masses(simulation.system)
    kb = simulation.kb
    n = length(simulation.system.bodies)
    nc = 0
    p = simulation.thermostat
    (ms, kb, n, nc, p)
end

function obtain_data_for_berendsen_thermostating(simulation::NBodySimulation{<:WaterSPCFw})
    ms = get_masses(simulation.system)
    kb = simulation.kb
    n = 3*length(simulation.system.bodies)
    nc = 2*length(simulation.system.bodies)
    p = simulation.thermostat
    (ms, kb, n, nc, p)
end

function get_nosehoover_thermostating_acceleration(simulation::NBodySimulation)
    (ms, kb, n, nc, γind, p) = obtain_data_for_nosehoover_thermostating(simulation)
    (dv, u, v, t) -> nosehoover_acceleration!(dv, u, v, ms, kb, n, nc, γind, p)
end

function obtain_data_for_nosehoover_thermostating(simulation::NBodySimulation)
    ms = get_masses(simulation.system)
    kb = simulation.kb
    n = length(simulation.system.bodies)
    γind = 3*n+1
    nc = 0
    p = simulation.thermostat
    (ms, kb, n, nc, γind, p)
end

function obtain_data_for_nosehoover_thermostating(simulation::NBodySimulation{<:WaterSPCFw})
    ms = get_masses(simulation.system)
    kb = simulation.kb
    n = 3*length(simulation.system.bodies)
    γind = 3*n+1
    nc = 2*length(simulation.system.bodies)
    p = simulation.thermostat
    (ms, kb, n, nc, γind, p)
end

function DiffEqBase.ODEProblem(simulation::NBodySimulation{<:PotentialNBodySystem})
    (u0, v0, n) = gather_bodies_initial_coordinates(simulation)

    acceleration_functions = gather_accelerations_for_potentials(simulation)

    function ode_system!(du, u, p, t)
        du[:, 1:n] = @view u[:, n + 1:2n];

        @inbounds for i = 1:n
            a = MVector(0.0, 0.0, 0.0)
            for acceleration! in acceleration_functions
                acceleration!(a, u[:, 1:n], u[:, n + 1:end], t, i);
            end
            du[:, n + i] .= a
        end
    end

    return ODEProblem(ode_system!, hcat(u0, v0), simulation.tspan)
end

function DiffEqBase.SecondOrderODEProblem(simulation::NBodySimulation{<:PotentialNBodySystem})
    (u0, v0, n) = gather_bodies_initial_coordinates(simulation)
    acceleration_functions = gather_accelerations_for_potentials(simulation)
    simultaneous_acceleration = gather_simultaneous_acceleration(simulation)

    function soode_system!(dv, v, u, p, t)
        @inbounds for i = 1:n
            a = MVector(0.0, 0.0, 0.0)
            for acceleration! in acceleration_functions
                acceleration!(a, u, v, t, i);
            end
            dv[:, i] .= a
        end
        for acceleration! in simultaneous_acceleration
            acceleration!(dv, u, v, t);
        end
    end

    SecondOrderODEProblem(soode_system!, v0, u0, simulation.tspan)
end

function DiffEqBase.SecondOrderODEProblem(simulation::NBodySimulation{<:WaterSPCFw})
    (u0, v0, n) = gather_bodies_initial_coordinates(simulation)
    (o_acelerations, h_acelerations) = gather_accelerations_for_potentials(simulation)
    group_accelerations = gather_group_accelerations(simulation)
    simultaneous_acceleration = gather_simultaneous_acceleration(simulation)
    #(ms_lj, indxs_lj) = obtain_data_for_lennard_jones_interaction(simulation.system)
    #(qs, ms_em, indxs_em, exclude) = obtain_data_for_electrostatic_interaction(simulation.system)
    #lj_parameters = simulation.system.lj_parameters
    #e_parameters  = simulation.system.e_parameters
    #pbc = simulation.boundary_conditions

    function soode_system!(dv, v, u, p, t)
        T = eltype(v)
        #=
        @inbounds for i=1:n
            aO  = MVector{3,T}(zero(T), zero(T), zero(T))
            aH1 = MVector{3,T}(zero(T), zero(T), zero(T))
            aH2 = MVector{3,T}(zero(T), zero(T), zero(T))
            pairwise_lennard_jones_acceleration!(aO,  u, 3*i-2, indxs_lj, ms_lj, lj_parameters, pbc)
            pairwise_electrostatic_acceleration!(aO,  u, 3*i-2, length(indxs_em), qs, ms_em, exclude, e_parameters, pbc)
            pairwise_electrostatic_acceleration!(aH1, u, 3*i-1, length(indxs_em), qs, ms_em, exclude, e_parameters, pbc)
            pairwise_electrostatic_acceleration!(aH2, u, 3*i-0, length(indxs_em), qs, ms_em, exclude, e_parameters, pbc)
            dv[:, 3*i-2] .= aO
            dv[:, 3*i-1] .= aH1
            dv[:, 3*i-0] .= aH2
        end
        =#
        @inbounds for i = 1:n
            a = MVector{3,T}(zero(T), zero(T), zero(T))
            for acceleration! in o_acelerations
                acceleration!(a, u, v, t, 3*i-2);
            end
            dv[:, 3*i-2] .= a
        end
        @inbounds for i ∈ 1:n, j ∈ (1,2)
            a = MVector{3,T}(zero(T), zero(T), zero(T))
            for acceleration! in h_acelerations
                acceleration!(a, u, v, t, 3*i-2+j);
            end
            dv[:, 3*i-2+j] .= a
        end
        @inbounds for i = 1:n
            for acceleration! in group_accelerations
                acceleration!(dv, u, v, t, i)
            end
        end
        for acceleration! in simultaneous_acceleration
            acceleration!(dv, u, v, t)
        end
    end
    SecondOrderODEProblem(soode_system!, v0, u0, simulation.tspan)
end

function gather_accelerations_for_potentials(simulation::NBodySimulation{<:WaterSPCFw})
    o_acelerations = []
    h_acelerations = []
    push!(o_acelerations, get_accelerating_function(simulation.system.e_parameters, simulation))
    push!(o_acelerations, get_accelerating_function(simulation.system.scpfw_parameters, simulation))
    push!(o_acelerations, get_accelerating_function(simulation.system.lj_parameters, simulation))
    push!(h_acelerations, get_accelerating_function(simulation.system.e_parameters, simulation))
    push!(h_acelerations, get_accelerating_function(simulation.system.scpfw_parameters, simulation))
    (o_acelerations, h_acelerations)
end

function DiffEqBase.SDEProblem(simulation::NBodySimulation{<:PotentialNBodySystem})
    (u0, v0, n) = gather_bodies_initial_coordinates(simulation)
    acceleration_functions = gather_accelerations_for_potentials(simulation)
    therm = simulation.thermostat

    function deterministic_acceleration!(du, u, p, t)
        du[:, 1:n] = @view u[:, n + 1:2n]
        @inbounds for i = 1:n
            a = MVector(0.0, 0.0, 0.0)
            for acceleration! in acceleration_functions
                acceleration!(a, (@view u[:, 1:n]), (@view u[:, n + 1:end]), t, i);
            end
            du[:, n + i] .= a
        end
        @. du[:, n + 1:end] -= therm.γ*u[:, n + 1:end]
    end

    σ = sqrt(2*therm.γ*simulation.kb*therm.T/simulation.system.bodies[1].m)

    function noise!(du, u, p, t)
        @. du[:, n + 1:end] += σ
    end

    return SDEProblem(deterministic_acceleration!, noise!, hcat(u0, v0), simulation.tspan)
end

function DiffEqBase.SDEProblem(simulation::NBodySimulation{<:WaterSPCFw})
    (u0, v0, n) = gather_bodies_initial_coordinates(simulation)
    (o_acelerations, h_acelerations) = gather_accelerations_for_potentials(simulation)
    group_accelerations = gather_group_accelerations(simulation)
    simultaneous_acceleration = gather_simultaneous_acceleration(simulation)
    therm = simulation.thermostat
    mH = simulation.system.mH
    mO = simulation.system.mO

    function deterministic_acceleration!(du, u, p, t)
        du[:, 1:3*n] = @view u[:, 3*n+1 : 2*3*n]
        @inbounds for i = 1:n
            a = MVector(0.0, 0.0, 0.0)
            for acceleration! in o_acelerations
                acceleration!(a, (@view u[:, 1:3*n]), (@view u[:, 3*n+1 : 2*3*n]), t, 3 * (i - 1) + 1);
            end
            du[:, 3*n+3 * (i - 1) + 1]  .= a
            @. du[:, 3*n+3 * (i - 1) + 1] -= therm.γ*u[:, 3*n+3 * (i - 1) + 1]/mO
            @. du[:, 3*n+3 * (i - 1) + 2] -= therm.γ*u[:, 3*n+3 * (i - 1) + 2]/mH
            @. du[:, 3*n+3 * (i - 1) + 3] -= therm.γ*u[:, 3*n+3 * (i - 1) + 3]/mH
        end
        @inbounds for i in 1:n, j in (2, 3)
            a = MVector(0.0, 0.0, 0.0)
            for acceleration! in h_acelerations
                acceleration!(a, (@view u[:, 1:3*n]), (@view u[:, 3*n+1 : 2*3*n]), t, 3 * (i - 1) + j);
            end
            du[:, 3*n+3 * (i - 1) + j]   .= a
        end
        @inbounds for i = 1:n
            for acceleration! in group_accelerations
                acceleration!((@view du[ :, 3*n+1 : 2*3*n]), (@view u[:, 1:3*n]), (@view u[:, 3*n+1 : 2*3*n]), t, i);
            end
        end
        for acceleration! in simultaneous_acceleration
            acceleration!((@view du[:, 3*n+1 : 2*3*n]), (@view u[:, 1:3*n]), (@view u[:, 3*n+1 : 2*3*n]), t);
        end
        @. du[:, 3*n+1 : 2*3*n] -= therm.γ*u[:, 3*n+1 : 2*3*n]
    end
    σO = sqrt(2*therm.γ*simulation.kb*therm.T)/mO
    σH = sqrt(2*therm.γ*simulation.kb*therm.T)/mH

    function noise!(du, u, p, t)
        @inbounds for i = 1:n
            @. du[:, 3*n+3 * (i - 1) + 1] += σO
            @. du[:, 3*n+3 * (i - 1) + 2] += σH
            @. du[:, 3*n+3 * (i - 1) + 3] += σH
        end
    end

    return SDEProblem(deterministic_acceleration!, noise!, hcat(u0, v0), simulation.tspan)
end


const FLOATING_POINT_REGEX = r"[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?"
const METERS_PER_FOOT = 0.3048

type NGSIMRoadway
    name::Symbol
    boundaries::Vector{Vector{VecE2}}
    centerlines::Vector{Vector{CurvePt}}
end
immutable RoadwayInputParams
    filepath_boundaries::String
    filepath_centerlines::String
end

function read_boundaries(io::IO)

    lines = readlines(io)
    for (i,line) in enumerate(lines)
        lines[i] = strip(line)
    end
    @assert(lines[1] == "BOUNDARIES")

    n_boundaries = parse(Int, lines[2])
    @assert(n_boundaries ≥ 0)

    retval = Array(Vector{VecE2}, n_boundaries)
    line_index = 2
    for i in 1 : n_boundaries

        @assert(lines[line_index+=1] == @sprintf("BOUNDARY %d", i))
        npts = parse(Int, lines[line_index+=1])
        line = Array(VecE2, npts)
        for j in 1 : npts
            matches = matchall(FLOATING_POINT_REGEX, lines[line_index+=1])
            x = parse(Float64, matches[1]) * METERS_PER_FOOT # convert to meters
            y = parse(Float64, matches[2]) * METERS_PER_FOOT
            line[j] = VecE2(x,y)
        end
        retval[i] = line
    end

    retval
end
function read_centerlines(io::IO)

    lines = readlines(io)
    for (i,line) in enumerate(lines)
        lines[i] = strip(line)
    end
    @assert(lines[1] == "CENTERLINES")

    n_centerlines = parse(Int, lines[2])
    @assert(n_centerlines ≥ 0)

    line_index = 2
    retval = Dict{AbstractString, Vector{CurvePt}}()
    for i in 1 : n_centerlines
        @assert(lines[line_index+=1] == "CENTERLINE")
        name = lines[line_index+=1]
        npts = parse(Int, lines[line_index+=1])
        line = Array(VecE2, npts)
        for j in 1 : npts
            matches = matchall(FLOATING_POINT_REGEX, lines[line_index+=1])
            x = parse(Float64, matches[1]) * METERS_PER_FOOT # convert to meters
            y = parse(Float64, matches[2]) * METERS_PER_FOOT # convert to meters
            line[j] = VecE2(x,y)
        end

        # post-process to extract heading and distance along lane
        centerline = Array(CurvePt, npts)
        let
            θ = atan2(line[2]-line[1])
            centerline[1] = CurvePt(VecSE2(line[1],θ), 0.0)
            for i in 2 : npts-1
                s = centerline[i-1].s + hypot(line[i] - line[i-1])
                θ = (atan2(line[i]-line[i-1]) + atan2(line[i+1]-line[i]))/2 # mean angle
                centerline[i] = CurvePt(VecSE2(line[i],θ), s)
            end
            s = centerline[npts-1].s + hypot(line[npts] - line[npts-1])
            θ = atan2(line[npts]-line[npts-1])
            centerline[npts] = CurvePt(VecSE2(line[npts],θ), s)
        end

        retval[name] = centerline

    end

    retval
end
function read_roadway(input_params::RoadwayInputParams)
    boundaries = open(read_boundaries, input_params.filepath_boundaries)
    centerlines = open(read_centerlines, input_params.filepath_centerlines)

    name = symbol(splitext(splitdir(input_params.filepath_boundaries)[2])[1])
    NGSIMRoadway(name, boundaries, collect(values(centerlines)))
end

function write_lwpolyline(io::IO, pts::Vector{CurvePt}, handle_int::Int, is_closed::Bool=false)
    N = length(pts)
    println(io, "  0")
    println(io, "LWPOLYLINE")
    println(io, "  5") # handle (increases)
    @printf(io, "B%d\n", handle_int)
    println(io, "100") # subclass marker
    println(io, "AcDbEntity")
    println(io, "  8") # layer name
    println(io, "150")
    println(io, "  6") # linetype name
    println(io, "ByLayer")
    println(io, " 62") # color number
    println(io, "  256")
    println(io, "370") # lineweight enum
    println(io, "   -1")
    println(io, "100") # subclass marker
    println(io, "AcDbPolyline")
    println(io, " 90") # number of vertices
    @printf(io, "   %d\n", N)
    println(io, " 70") # 0 is default, 1 is closed
    @printf(io, "    %d\n", is_closed ? 1 : 0)
    println(io, " 43") # 0 is constant width
    println(io, "0")

    for i in 1 : N
        println(io, " 10")
        @printf(io, "%.3f\n", pts[i].pos.x)
        println(io, " 20")
        @printf(io, "%.3f\n", pts[i].pos.y)
    end
end

function convert_curves_feet_to_meters!(roadway::Roadway)
    for seg in roadway.segments
        for lane in seg.lanes
            for (i,P) in enumerate(lane.curve)
                lane.curve[i] = CurvePt(
                        VecSE2(P.pos.x*METERS_PER_FOOT, P.pos.y*METERS_PER_FOOT, P.pos.θ),
                        P.s*METERS_PER_FOOT, P.k/METERS_PER_FOOT, P.kd/METERS_PER_FOOT)
            end
        end
    end
    roadway
end

function write_dxf(io::IO, roadway::NGSIMRoadway)

    lines = open(readlines, Pkg.dir("NGSIM", "data", "template.dxf"))
    i = findfirst(lines, "ENTITIES\n")
    i != 0 || error("ENTITIES section not found")

    # write out header
    for j in 1 : i
        print(io, lines[j])
    end

    # write out the lanes
    for (handle_int, lane) in enumerate(roadway.centerlines)
        write_lwpolyline(io, lane, handle_int)
    end

    # write out tail
    for j in i+1 : length(lines)
        print(io, lines[j])
    end
end
function write_roadways_to_dxf()

    roadway_input_80 = RoadwayInputParams(Pkg.dir("NGSIM", "data", "boundaries80.txt"),
                                          Pkg.dir("NGSIM", "data", "centerlines80.txt"))
    roadway_input_101 = RoadwayInputParams(Pkg.dir("NGSIM", "data", "boundaries101.txt"),
                                           Pkg.dir("NGSIM", "data", "centerlines101.txt"))

    ngsimroadway_80 = read_roadway(roadway_input_80)
    ngsimroadway_101 = read_roadway(roadway_input_101)

    open(io->write_dxf(io, ROADWAY_80), Pkg.dir("NGSIM", "data", "ngsim_80.dxf"), "w")
    open(io->write_dxf(io, ROADWAY_101), Pkg.dir("NGSIM", "data", "ngsim_101.dxf"), "w")
end
function write_roadways_from_dxf()

    roadway_80 = open(io->read_dxf(io, Roadway, dist_threshold_lane_connect=2.0), Pkg.dir("NGSIM", "data", "ngsim_80.dxf"), "r")
    roadway_101 = open(io->read_dxf(io, Roadway, dist_threshold_lane_connect=2.0), Pkg.dir("NGSIM", "data", "ngsim_101.dxf"), "r")

    # also converts to meters
    convert_curves_feet_to_meters!(roadway_80)
    convert_curves_feet_to_meters!(roadway_101)

    open(io->write(io, roadway_80), Pkg.dir("NGSIM", "data", "ngsim_80.txt"), "w")
    open(io->write(io, roadway_101), Pkg.dir("NGSIM", "data", "ngsim_101.txt"), "w")
end

const ROADWAY_80 = open(io->read(io, Roadway), Pkg.dir("NGSIM", "data", "ngsim_80.txt"), "r")
const ROADWAY_101 = open(io->read(io, Roadway), Pkg.dir("NGSIM", "data", "ngsim_101.txt"), "r")


# Prepare data for IL
# 1. run MILP to get the optimal solutions for each scenario
# 2. extract state-action dataset of each load for DNN training in behavior cloning
# In the returned dataset (X, y), each column of X is a state, and each value of y is the
# corresponding scalar action.

"""
    run_expert!(ss::AbstractVector{DataFrame}, h::Dict;
        out_bson::Union{String, Nothing} = nothing) -> Vector{DataFrame}

Run an IL expert to get optimal demonstrations for each scenario in `ss`. The optimal solutions are
written into `ss` in place, while the updated `ss` is also returned.

Besides, the updated `ss` is stored into the BSON file `out_bson` with key `:ss`.
"""
function run_expert!(ss::AbstractVector{DataFrame}, h::Dict;
    out_bson::Union{String,Nothing} = nothing)
    train_ss = DataFrame[]
    for s in ss
        # the solutions have been injected into `s` in place
        run_MILP!(s, h)
        push!(train_ss, s)
    end
    if out_bson !== nothing
        @save out_bson ss = train_ss
    end
    return train_ss
end

"""
    build_expert_dataset(ss::AbstractVector{DataFrame}, h::Dict;
        out_bson::Union{String,Nothing} = nothing) -> Dict

Build expert demonstrations dataset for each load in `h` based on optimized scenarios `ss`.
The dataset (X, y) of a load with ID `id` can be retrieved with key `id` from the returned `Dict`.

If the file path `out_bson` is specified, then the returned `Dict` is also stored into the BSON
file with key `:d`.
"""
function build_expert_dataset(ss::AbstractVector{DataFrame}, h::Dict; H=1,
    out_bson::Union{String,Nothing} = nothing)::Dict
    d = Dict(l["id"] => build_AD_dataset(ss, l) for l in h["AD"])
    for l in h["SU"]
        d[l["id"]] = build_SU_dataset(ss, l; H)
    end
    for l in h["SI"]
        d[l["id"]] = build_SI_dataset(ss, l)
    end
    if !isnothing(out_bson)
        @save out_bson d
    end
    return d
end

"""
    build_AD_dataset(ss::AbstractVector{DataFrame}, id::String) -> (Matrix, Vector)

Create a dataset for an AD load with ID `id` from scenarios `ss`.
State: ``[t, ρ, P^{PV}]``. Action: ``P^{AD}``.
"""
function build_AD_dataset(ss::AbstractVector{DataFrame}, id::String)
    n = length(ss)
    T = nrow(ss[1])
    X = zeros(3, n * T)
    y = zeros(n * T)
    i = 1
    for s in ss
        for t = 1:T
            X[:, i] .= [t, s.ρ[t], s.PPV[t]]
            y[i] = s[t, "P_$id"]
            i += 1
        end
    end
    @assert i == length(y) + 1
    return X, y
end

"""
    build_AD_dataset(ss::AbstractVector{DataFrame}, l::Dict) -> (Matrix, Vector)

Create a dataset for an AD load `l`.
See [`build_AD_dataset(ss::AbstractVector{DataFrame}, id::String)`](@ref).
"""
build_AD_dataset(ss::AbstractVector{DataFrame}, l::Dict) = build_AD_dataset(ss, l["id"])

"""
    build_SU_dataset(ss::AbstractVector{DataFrame}, l::Dict) -> (Matrix, Vector)

Create a dataset for a SU load `l` from scenarios `ss`.
State: ``[t, ρ, P^{PV}]``. Action: ``z^{su}``.
Note that we only consider ``t \\in [ts, tf-L+1]``, in which the action may be 1.
"""
function build_SU_dataset(ss::AbstractVector{DataFrame}, l::Dict, ls_AD=nothing; H::Integer = 1)
    @assert H >= 1
    n = length(ss)
    ts, tf, L = l["ts"], l["tf"], l["L"]
    m = tf - L + 1 - ts + 1 # number of state-action pairs in each scenario
    n_AD = 0
    if ls_AD !== nothing
        n_AD = length(ls_AD)
    end
    X = zeros(1 + 2 * H + n_AD, n * m)
    y = zeros(n * m)
    i = 1
    id = l["id"]
    for s in ss
        count_one = 0    # each scenario has exactly one 1
        for t = ts:tf-L+1
            # state: t and historical price & PV generation
            X[1, i] = t
            X[2:2+H-1, i] .= (s.ρ[h] for h = t-H+1:t)
            X[2+H:2+H+H-1, i] .= (s.PPV[h] for h = t-H+1:t)
            # binary variables may not be exactly 0 or 1 in JuMP
            if s[t, "z_$id"] > 0.5
                y[i] = 1
                count_one += 1
            else
                y[i] = 0
            end
            if ls_AD !== nothing
                X[2+2*H:end, i] .= (s[t, "P_$(lad["id"])"] for lad in ls_AD)
            end
            i += 1
        end
        @assert count_one == 1
    end
    @assert i == length(y) + 1
    return X, y
end

"""
    build_SI_dataset(ss::AbstractVector{DataFrame}, l::Dict) -> (Matrix, Vector)

Create a dataset for a SI load `l` from scenarios `ss`.
State: ``[t, ρ, P^{PV}, O^{si}]``. Action: ``z^{si}``.
Note that we only consider ``t \\in [ts, tf]``, in which the action may be 1.
"""
function build_SI_dataset(ss::AbstractVector{DataFrame}, l::Dict)
    n = length(ss)
    ts, tf, L = l["ts"], l["tf"], l["L"]
    m = tf - ts + 1 # number of state-action pairs in each scenario
    X = zeros(4, n * m)
    y = zeros(n * m)
    i = 1
    id = l["id"]
    i_s = 0
    for s in ss
        i_s += 1
        Osi = 0
        for t = ts:tf
            X[:, i] .= [t, s.ρ[t], s.PPV[t], Osi]
            y[i] = s[t, "z_$id"] > 0.5 ? 1 : 0 # JuMP may not yield exactly 1 though it is a binary
            if y[i] == 1
                Osi += 1
            end
            i += 1
        end
        @assert Osi == L
    end
    @assert i == length(y) + 1
    return X, y
end


## Metadata For graph creation
@enum VertexLocation LOC_PMEM LOC_DRAM LOC_DRAM_PRE LOC_SOURCE LOC_SINK

# Helpers - generally pretty simple but helpful for abstracting the `isdram` attribute.
ispmem(loc::VertexLocation) = loc == LOC_PMEM
isdram(loc::VertexLocation) = loc == LOC_DRAM || loc == LOC_DRAM_PRE
issource(loc::VertexLocation) = loc == LOC_SOURCE
issink(loc::VertexLocation) = loc == LOC_SINK

@enum EdgeType begin
    EDGE_NONE
    EDGE_SYNC_READ
    EDGE_SYNC_WRITE
    EDGE_ASYNC_READ
    EDGE_ASYNC_WRITE
end

isasync(et::EdgeType) = in(et, (EDGE_ASYNC_READ, EDGE_ASYNC_WRITE))

@enum MoveType MOVE_NONE MOVE_SYNC MOVE_ASYNC

# Metadata to assign to each node in the liveness graph for tensors.
"""
    VertexMetadata

Metadata associated with the vertex of the tensor graph.

$(FIELDS)
"""
struct VertexMetadata
    "The gadget that this vertex belongs to. Used for edge generation."
    gadget::Int
    "The op index that this gadget refers to."
    op::XNode
    "Where the vertex lives."
    location::VertexLocation
    "What type of moves this vertex allows."
    move_type::MoveType
    "`true` is this vertex is a user of the tensor in whose graph this vertex appear."
    isuser::Bool
    vertex_number::Int
end

struct EdgeMetadata
    edgetype::EdgeType
end
isasync(em::EdgeMetadata) = isasync(em.edgetype)

#####
##### Preprocessing
#####

"""
$(SIGNATURES) -> UnitRange{Int64}

Return a range of indices corresponding to nodes during which `xtensor` is live.
"""
liverange(xtensor::XTensor) = (producer(xtensor).index):(consumer(xtensor).index)

function _getgadgets(A::ILPHolder{Asynchronous}, data::FunctionData, t::XTensor)
    range = liverange(t)
    livenodes = (nodes(data)[x] for x in range)
    refs = Vector{NamedTuple{(:node, :move_type),Tuple{XNode,MoveType}}}()

    # Build the referece map
    reference_map = Dict{XNode, XNode}()
    ref = producer(t)

    # To decide if a node should be considered as an aynchronous move point, we check to see
    # if the node is with `bound` distance of a user of the tensor.
    #
    # The intuition here is that moves will probably be located closer to their producers or
    # consumers rather than further.
    #
    # Making `bound` larger increased the search space of the formulation, which may lead to
    # better results at the cost of a larger mode.
    bound = 15
    move_time = sizeof(t) / A.write_bandwidth
    for ind in range
        node = nodes(data)[ind]
        if in(node, users(t))
            push!(refs, (node = node, move_type = MOVE_SYNC))
            ref = node

        # Check if there is a user node within `bound`. If so, make this an async move node.
        elseif hasprofile(node) &&
            any(
                in(users(t)),
                (nodes(data)[i] for i in max(ind-bound, 1):min(ind+bound, length(nodes(data))))
            )

            push!(refs, (node = node, move_type = MOVE_ASYNC))
            ref = node
        end
        reference_map[node] = ref
    end

    return refs, reference_map
end

function _getgadgets(::ILPHolder{Synchronous}, data::FunctionData, t::XTensor)
    range = liverange(t)
    livenodes = (nodes(data)[x] for x in range)

    # Build the referece map
    reference_map = Dict{XNode, XNode}()
    ref = producer(t)
    for ind in range
        node = nodes(data)[ind]
        if in(node, users(t))
            ref = node
        end
        reference_map[node] = ref
    end

    nt = [
        (node = u, move_type = isone(i) ? MOVE_NONE : MOVE_SYNC) for (i,u) in enumerate(users(t))
    ]

    return nt, reference_map
end

function _getgadgets(::ILPHolder{Static}, data::FunctionData, t::XTensor)
    range = liverange(t)
    producer = nodes(data)[first(range)]

    reference_map = Dict{XNode, XNode}()
    for ind in range
        reference_map[nodes(data)[ind]] = producer
    end

    return [(node = producer, move_type = MOVE_NONE)], reference_map
end

function getgadgets(x, data::FunctionData, t::XTensor)
    refs, reference_map = _getgadgets(x, data, t)
    # Update with argument

    # If this tensor `isarg`, then since it is live for the entire function, we need to
    # insert node references for all nodes pointing to the producer.
    if isarg(t)
        ref = producer(t)
        for node in nodes(data)
            # Shortcut here, `get!` will lookup and automatically insert the mapping if it
            # doesn't exist.
            get!(reference_map, node, ref)
        end
    end

    return refs, reference_map
end

# TODO: When can't do anything because exhausted, make ASCII art for what's going on with
# these tensor graphs.
#=
=#

function edge_metadata(src, dst, s, d, src_move_type, tensor::XTensor)
    # Setup the correct move annotations.
    if src_move_type == MOVE_ASYNC
        edge_read_type = EDGE_ASYNC_READ
        edge_write_type = EDGE_ASYNC_WRITE
    else
        edge_read_type = EDGE_SYNC_READ
        edge_write_type = EDGE_SYNC_WRITE
    end

    # Determine if an edge should be added and what kind of edge it is.

    ### LOC_SOURCE node is source
    if (src, dst) == (LOC_SOURCE, LOC_DRAM_PRE)
        isone(d) && return EdgeMetadata(EDGE_NONE)
    elseif (src, dst) == (LOC_SOURCE, LOC_PMEM)
        isone(d) && return EdgeMetadata(EDGE_NONE)
    # If the tensor is an argument, we need to create an edge directly from source to sink.
    # This indicates that the tensor begins in DRAM - so will always be in DRAM
    elseif (src, dst) == (LOC_SOURCE, LOC_SINK)
        isarg(tensor) && return EdgeMetadata(EDGE_NONE)

    ### LOC_DRAM as source

    # If we're in LOC_DRAM, data already exists in PMEM. Thus all edges originating in
    # LOC_DRAM have no metadata. I.E., no read or write type eges.
    elseif (src, dst) == (LOC_DRAM, LOC_DRAM)
        s == d-1 && return EdgeMetadata(EDGE_NONE)
    elseif (src, dst) == (LOC_DRAM, LOC_PMEM)
        s == d-1 && return EdgeMetadata(EDGE_NONE)
    elseif (src, dst) == (LOC_DRAM, LOC_SINK)
        s == d-1 && return EdgeMetadata(EDGE_NONE)

    ### LOC_DRAM_PRE as source

    # Moving from LOC_DRAM_PRE to LOC_PMEM indicates a movement of data from DRAM to PMEM
    # where the data did not exist in PMEM before. Thus, we must record a write.
    elseif (src, dst) == (LOC_DRAM_PRE, LOC_PMEM)
        s == d-1 && return EdgeMetadata(edge_write_type)
    # Movement between DRAM or to the SINK requires no metadata
    elseif (src, dst) == (LOC_DRAM_PRE, LOC_DRAM_PRE)
        s == d-1 && return EdgeMetadata(EDGE_NONE)
    elseif (src, dst) == (LOC_DRAM_PRE, LOC_SINK)
        s == d-1 && return EdgeMetadata(EDGE_NONE)

    ### LOC_PMEM as source

    # Movement edges from PMEM to DRAM. Must not happen on the first `component` because
    # prefetching from PMEM to DRAM at the time of creation makes no sense.
    elseif (src, dst) == (LOC_PMEM, LOC_DRAM)
        (s == d) && !isone(s) && return EdgeMetadata(edge_read_type)
    elseif (src, dst) == (LOC_PMEM, LOC_PMEM)
        (s == d-1) && return EdgeMetadata(EDGE_NONE)
    # Otherwise, no metadata needed
    elseif (src, dst) == (LOC_PMEM, LOC_SINK)
        (s == d-1) && return EdgeMetadata(EDGE_NONE)
    end
    return nothing
end

function preprocess!(S::ILPHolder, data::FunctionData)
     for tensor in tensors(data)

        # Get the users of this node
        # Get two things from getgadgets:
        #
        # 1. A named tuple (node::XNode, move_type::MoveType)
        # 2. A dictionary implementing the `ref` function.
        gadgets, reference_map = getgadgets(S, data, tensor)

        @assert !isempty(gadgets)

        # Graph building time :D
        g = MetaGraph(DiGraph(), EdgeMetadata, VertexMetadata)

        # If this tensor can only be assigned to a single location - don't generate a graph
        # for it - otherwise, populate the nodes in a tensor graph
        fixed_tensor = length(locations(tensor)) == 1
        if !fixed_tensor
            for (count, nt) in enumerate(gadgets)
                # Unpack the XNode
                node = nt.node
                move_type = nt.move_type
                islast = (count == length(gadgets))

                isuser = in(node, users(tensor))

                if count == 1
                    add_vertex!(g, VertexMetadata(0, node, LOC_SOURCE, move_type, isuser, nv(g)+1))
                end
                # Enumerate over locations that this tensor can live.
                #
                # Do it this way because some nodes can only live in DRAM, so iterating
                # then filtering takes care of that
                for location in locations(tensor)
                    if location == DRAM
                        # Add DRAM nodes
                        #
                        # Do not add these nodes if this tensor is an argument because it will
                        # have already originated in PMM
                        if !isarg(tensor)
                            add_vertex!(g,
                                VertexMetadata(count, node, LOC_DRAM_PRE, move_type, isuser, nv(g)+1)
                            )
                        end

                        # only add a DRAM node if there could have been a write to PMEM
                        if count > 1
                            add_vertex!(g,
                                VertexMetadata(count, node, LOC_DRAM, move_type, isuser, nv(g)+1)
                            )
                        end
                    end

                    if location == PMEM
                        @assert !startswith(nGraph.name(tensor), "Constant")
                        # PMEM node
                        add_vertex!(g,
                            VertexMetadata(count, node, LOC_PMEM, move_type, isuser, nv(g)+1)
                        )
                    end
                end
                if islast
                    # Set the gadget number for the sink to one higher than the last count.
                    add_vertex!(g,
                        VertexMetadata(count + 1, node, LOC_SINK, move_type, isuser, nv(g)+1)
                    )
                end
            end
        end

        # Use a quadratic complexity algorithm for doing edge assignment. It's not
        # perfect but it's simple, and as long as the graphs don't get too big should
        # run quickly enough for our purposes.
        for src in vertices(g), dst in vertices(g)
            src == dst && continue

            src_meta = getmeta(g, src)
            dst_meta = getmeta(g, dst)

            metadata = edge_metadata(
                src_meta.location,
                dst_meta.location,
                src_meta.gadget,
                dst_meta.gadget,
                src_meta.move_type,
                tensor,
            )

            isnothing(metadata) && continue

            add_edge!(g, src, dst, metadata)
        end

        # Create the descriptor
        S.descriptors[tensor] = TensorMeta(
            g,
            [g.node for g in gadgets],
            reference_map,
            fixed_tensor
        )
    end
end


using Pkg
using PackageCompiler

function build_julia_image(project_path::String)

    Pkg.activate(joinpath(project_path, "julia_envs", "Xpress"))
    Pkg.instantiate(verbose=true)

    # must point pycall to the python path we want it to use
    ENV["PYTHON"] = joinpath(project_path, "env", "bin", "python")
    ENV["PROJECT_PATH"] = project_path

    # TODO create REopt module to attach to Base for system image
    if Sys.islinux()
        ext = ".so"
    elseif Sys.isapple()
        ext = ".dylib"
    elseif Sys.iswindows()
        ext = ".dll"
    else
        error("Unsupported operating system")
    end

    PackageCompiler.create_sysimage(
        [:AxisArrays, :JuMP, :MathOptInterface, :PyCall, :Xpress, :MutableArithmetics],
        sysimage_path=joinpath(project_path, "julia_envs", "Xpress", "JuliaXpressSysimage" * ext),
        precompile_execution_file=joinpath(project_path, "julia_envs", "Xpress", "precompile.jl")
    )
end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
"""
    struct ElectricTariff

- data for electric tariff in reopt model
- can be defined using custom rates or URDB rate
- very similar to the URDB struct but includes export rates and bins
"""
struct ElectricTariff
    energy_rates::AbstractArray{Float64, 2} # gets a second dim with tiers
    energy_tier_limits::AbstractArray{Float64,1}
    n_energy_tiers::Int

    monthly_demand_rates::AbstractArray{Float64, 2} # gets a second dim with tiers
    time_steps_monthly::AbstractArray{AbstractArray{Int64,1},1}  # length = 0 or 12
    monthly_demand_tier_limits::AbstractArray{Float64,1}
    n_monthly_demand_tiers::Int

    tou_demand_rates::AbstractArray{Float64, 2} # gets a second dim with tiers
    tou_demand_ratchet_time_steps::AbstractArray{AbstractArray{Int64,1},1}  # length = n_tou_demand_ratchets
    tou_demand_tier_limits::AbstractArray{Float64,1}
    n_tou_demand_tiers::Int

    demand_lookback_months::AbstractArray{Int,1}
    demand_lookback_percent::Float64
    demand_lookback_range::Int

    fixed_monthly_charge::Float64
    annual_min_charge::Float64
    min_monthly_charge::Float64

    export_rates::Dict{Symbol, AbstractArray}
    export_bins::AbstractArray{Symbol,1}

    coincident_peak_load_active_time_steps::AbstractVector{AbstractVector{Int64}}
    coincident_peak_load_charge_per_kw::AbstractVector{Float64}
    coincpeak_periods::AbstractVector{Int64}
end

"""
    ElectricTariff

ElectricTariff constructor
```julia
function ElectricTariff(;
    urdb_label::String="",
    urdb_response::Dict=Dict(),
    urdb_utility_name::String="",
    urdb_rate_name::String="",
    year::Int=2020,
    time_steps_per_hour::Int=1,
    NEM::Bool=false,
    wholesale_rate::T1=nothing,
    export_rate_beyond_net_metering_limit::T2=nothing,
    monthly_energy_rates::Array=[],
    monthly_demand_rates::Array=[],
    blended_annual_energy_rate::S=nothing,
    blended_annual_demand_rate::R=nothing,
    add_monthly_rates_to_urdb_rate::Bool=false,
    tou_energy_rates_per_kwh::Array=[],
    add_tou_energy_rates_to_urdb_rate::Bool=false,
    remove_tiers::Bool=false,
    demand_lookback_months::AbstractArray{Int64, 1}=Int64[],
    demand_lookback_percent::Float64=0.0,
    demand_lookback_range::Int=0,
    coincident_peak_load_active_time_steps::Vector{Vector{Int64}}=[Int64[]],
    coincident_peak_load_charge_per_kw::AbstractVector{<:Real}=Real[]
    ) where {
        T1 <: Union{Nothing, Int, Float64, Array},
        T2 <: Union{Nothing, Int, Float64, Array},
        S <: Union{Nothing, Int, Float64},
        R <: Union{Nothing, Int, Float64}
    }
```

!!! note
    The `NEM` boolean is determined by the ElectricUtility.net_metering_limit_kw. There is no need to pass in a `NEM`
    value.

"""
function ElectricTariff(;
    urdb_label::String="",
    urdb_response::Dict=Dict(),
    urdb_utility_name::String="",
    urdb_rate_name::String="",
    year::Int=2020,
    time_steps_per_hour::Int=1,
    NEM::Bool=false,
    wholesale_rate::T1=nothing,
    export_rate_beyond_net_metering_limit::T2=nothing,
    monthly_energy_rates::Array=[],
    monthly_demand_rates::Array=[],
    blended_annual_energy_rate::S=nothing,
    blended_annual_demand_rate::R=nothing,
    add_monthly_rates_to_urdb_rate::Bool=false,
    tou_energy_rates_per_kwh::Array=[],
    add_tou_energy_rates_to_urdb_rate::Bool=false,
    remove_tiers::Bool=false,
    demand_lookback_months::AbstractArray{Int64, 1}=Int64[],
    demand_lookback_percent::Float64=0.0,
    demand_lookback_range::Int=0,
    coincident_peak_load_active_time_steps::Vector{Vector{Int64}}=[Int64[]],
    coincident_peak_load_charge_per_kw::AbstractVector{<:Real}=Real[]
    ) where {
        T1 <: Union{Nothing, Int, Float64, Array},
        T2 <: Union{Nothing, Int, Float64, Array},
        S <: Union{Nothing, Int, Float64},
        R <: Union{Nothing, Int, Float64}
    }
    # TODO remove_tiers for multinode models
    nem_rate = Float64[]

    energy_tier_limits = Float64[]
    n_energy_tiers = 1
    monthly_demand_tier_limits = Float64[]
    n_monthly_demand_tiers = 1
    tou_demand_tier_limits = Float64[]
    n_tou_demand_tiers = 1
    time_steps_monthly = get_monthly_time_steps(year, time_steps_per_hour=time_steps_per_hour)

    u = nothing
    if !isempty(urdb_response)

        u = URDBrate(urdb_response, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(urdb_label)

        u = URDBrate(urdb_label, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(urdb_utility_name) && !isempty(urdb_rate_name)

        u = URDBrate(urdb_utility_name, urdb_rate_name, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(tou_energy_rates_per_kwh) && length(tou_energy_rates_per_kwh) == 8760*time_steps_per_hour

        tou_demand_rates = Float64[]
        tou_demand_ratchet_time_steps = []
        energy_rates = tou_energy_rates_per_kwh
        monthly_demand_rates = convert(Array{Float64}, monthly_demand_rates)

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end

    elseif !isempty(monthly_energy_rates)

        invalid_args = String[]
        if !(length(monthly_energy_rates) == 12)
            push!(invalid_args, "length(monthly_energy_rates) must equal 12, got length $(length(monthly_energy_rates))")
        end
        if !isempty(monthly_demand_rates) && !(length(monthly_demand_rates) == 12)
            push!(invalid_args, "length(monthly_demand_rates) must equal 12, got length $(length(monthly_demand_rates))")
        end
        if length(invalid_args) > 0
            error("Invalid argument values: $(invalid_args)")
        end

        if isempty(monthly_demand_rates)
            monthly_demand_rates = repeat([0.0], 12)
        end

        tou_demand_rates = Float64[]
        tou_demand_ratchet_time_steps = []
        energy_rates = Real[]
        for m in 1:12
            append!(energy_rates, [monthly_energy_rates[m] for ts in time_steps_monthly[m]])
        end

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end

    elseif !isnothing(blended_annual_energy_rate)

        tou_demand_rates = Float64[]
        tou_demand_ratchet_time_steps = []
        energy_rates = repeat(Real[blended_annual_energy_rate], 8760 * time_steps_per_hour)
        if !isnothing(blended_annual_demand_rate)
            monthly_demand_rates = repeat(Real[blended_annual_demand_rate], 12)
        end
        if isempty(monthly_demand_rates)
            monthly_demand_rates = repeat([0.0], 12)
        end

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end

    else
        error("Creating ElectricTariff requires at least urdb_label, urdb_response, monthly rates, annual rates, or tou_energy_rates_per_kwh.")
    end

    if !isnothing(u)  # use URDBrate

        if NEM
            t = get_tier_with_lowest_energy_rate(u)
            nem_rate = [-0.999 * x for x in u.energy_rates[t,:]]
        end

        energy_rates = u.energy_rates
        energy_tier_limits = u.energy_tier_limits
        n_energy_tiers = u.n_energy_tiers
        users_monthly_demand_rates = copy(monthly_demand_rates)
        monthly_demand_rates = u.monthly_demand_rates
        monthly_demand_tier_limits = u.monthly_demand_tier_limits
        n_monthly_demand_tiers = u.n_monthly_demand_tiers
        tou_demand_rates = u.tou_demand_rates
        tou_demand_tier_limits = u.tou_demand_tier_limits
        n_tou_demand_tiers = u.n_tou_demand_tiers

        if remove_tiers
            energy_rates, monthly_demand_rates, tou_demand_rates = remove_tiers_from_urdb_rate(u)
            energy_tier_limits, monthly_demand_tier_limits, tou_demand_tier_limits =
                Float64[], Float64[], Float64[]
            n_energy_tiers, n_monthly_demand_tiers, n_tou_demand_tiers = 1, 1, 1
        end

        tou_demand_ratchet_time_steps = u.tou_demand_ratchet_time_steps
        demand_lookback_months = u.demand_lookback_months
        demand_lookback_percent = u.demand_lookback_percent
        demand_lookback_range = u.demand_lookback_range
        fixed_monthly_charge = u.fixed_monthly_charge
        annual_min_charge = u.annual_min_charge
        min_monthly_charge = u.min_monthly_charge

        if add_monthly_rates_to_urdb_rate
            if length(monthly_energy_rates) == 12
                for tier in 1:size(energy_rates, 2), mth in 1:12, ts in time_steps_monthly[mth]
                    energy_rates[ts, tier] += monthly_energy_rates[mth]
                end
            end
            if length(users_monthly_demand_rates) == 12
                for tier in 1:size(monthly_demand_rates, 2), mth in 1:12
                    monthly_demand_rates[mth, tier] += users_monthly_demand_rates[mth]
                end
            end
        end

        if add_tou_energy_rates_to_urdb_rate && length(tou_energy_rates_per_kwh) == size(energy_rates, 1)
            for tier in 1:size(energy_rates, 2)
                energy_rates[1:end, tier] += tou_energy_rates_per_kwh
            end
        end
    else
        # need to reshape cost vectors to arrays (2nd dim is for tiers)
        energy_rates = reshape(energy_rates, :, 1)
        monthly_demand_rates = reshape(monthly_demand_rates, :, 1)
        tou_demand_rates = reshape(tou_demand_rates, :, 1)
    end

    #= export_rates
    There are three Export tiers and their associated export rates (negative values):
    1. NEM (Net Energy Metering)
    2. WHL (Wholesale)
    3. EXC (Excess, beyond NEM)

    Only one of NEM and Wholesale can be exported into due to the binary constraints.
    Excess can be exported into in the same time step as NEM.

    Excess is meant to be combined with NEM: NEM export is limited to the total grid purchased energy in a year and some
    utilities offer a compensation mechanism for export beyond the site load.
    The Excess tier is not available with the Wholesale tier.

    - if NEM then set ExportRate[:Nem, :] to energy_rate[tier_with_lowest_energy_rate, :]
    - user can provide either scalar wholesale rate or vector of time_steps,
    =#
    whl_rate = create_export_rate(wholesale_rate, length(energy_rates), time_steps_per_hour)
    if !isnothing(u) && sum(u.sell_rates) < 0
        whl_rate += u.sell_rates
    end
    exc_rate = create_export_rate(export_rate_beyond_net_metering_limit, length(energy_rates), time_steps_per_hour)

    if !NEM & (sum(whl_rate) >= 0)
        export_rates = Dict{Symbol, AbstractArray}()
        export_bins = Symbol[]
    elseif !NEM
        export_bins = [:WHL]
        export_rates = Dict(:WHL => whl_rate)
    elseif (sum(whl_rate) >= 0)
        export_bins = [:NEM]
        export_rates = Dict(:NEM => nem_rate)
        if sum(exc_rate) < 0
            push!(export_bins, :EXC)
            export_rates[:EXC] = exc_rate
        end
    else
        export_bins = [:NEM, :WHL]
        export_rates = Dict(:NEM => nem_rate, :WHL => whl_rate)
        if sum(exc_rate) < 0
            push!(export_bins, :EXC)
            export_rates[:EXC] = exc_rate
        end
    end

    coincpeak_periods = Int64[]
    if !isempty(coincident_peak_load_charge_per_kw)
        coincpeak_periods = collect(1:length(coincident_peak_load_charge_per_kw))
    end

    ElectricTariff(
        energy_rates,
        energy_tier_limits,
        n_energy_tiers,
        monthly_demand_rates,
        time_steps_monthly,
        monthly_demand_tier_limits,
        n_monthly_demand_tiers,
        tou_demand_rates,
        tou_demand_ratchet_time_steps,
        tou_demand_tier_limits,
        n_tou_demand_tiers,
        demand_lookback_months,
        demand_lookback_percent,
        demand_lookback_range,
        fixed_monthly_charge,
        annual_min_charge,
        min_monthly_charge,
        export_rates,
        export_bins,
        coincident_peak_load_active_time_steps,
        coincident_peak_load_charge_per_kw,
        coincpeak_periods
    )
end

function get_tier_with_lowest_energy_rate(u::URDBrate)
    """
    ExportRate should be lowest energy cost for tiered rates.
    Otherwise, ExportRate can be > FuelRate, which leads REopt to export all PV energy produced.
    """
    tier_with_lowest_energy_cost = 1
    if length(u.energy_tier_limits) > 1
        annual_energy_charge_sums = Float64[]
        for etier in u.energy_rates
            push!(annual_energy_charge_sums, sum(etier))
        end
        tier_with_lowest_energy_cost =
            findall(annual_energy_charge_sums .== minimum(annual_energy_charge_sums))[1]
    end
    return tier_with_lowest_energy_cost
end

"""
    function create_export_rate(e::Nothing, N::Int, ts_per_hour::Int=1)
No export rate provided by user: set to 0 dollars/kWh for all time
"""
function create_export_rate(e::Nothing, N::Int, ts_per_hour::Int=1)
    [0 for _ in range(1, stop=N) for ts in 1:ts_per_hour]
end

"""
    function create_export_rate(e::T, N::Int, ts_per_hour::Int=1) where T<:Real
Case for scaler export rate provided -> convert to array of time_steps
"""
function create_export_rate(e::T, N::Int, ts_per_hour::Int=1) where T<:Real
    repeat(float(-1*e), N * ts_per_hour)
end

"""
    function create_export_rate(e::AbstractArray{<:Real, 1}, N::Int, ts_per_hour::Int=1)

Check length of e and upsample if length(e) != N
"""
function create_export_rate(e::AbstractArray{<:Real, 1}, N::Int, ts_per_hour::Int=1)
    Ne = length(e)
    if Ne != Int(N/ts_per_hour) || Ne != N
        @error "Export rates do not have correct number of entries. Must be $(N) or $(Int(N/ts_per_hour))."
    end
    if Ne != N  # upsample
        export_rates = [-1*x for x in e for ts in 1:ts_per_hour]
    else
        export_rates = -1*e
    end
    return export_rates
end

"""
    get_monthly_time_steps(year::Int; time_steps_per_hour=1)

return Array{Array{Int64,1},1}, size = (12,)
"""
function get_monthly_time_steps(year::Int; time_steps_per_hour=1)
    a = Array[]
    i = 1
    for m in range(1, stop=12)
        n_days = daysinmonth(Date(string(year) * "-" * string(m)))
        stop = n_days * 24 * time_steps_per_hour + i - 1
        if m == 2 && isleapyear(year)
            stop -= 24 * time_steps_per_hour  # TODO support extra day in leap years?
        end
        steps = [step for step in range(i, stop=stop)]
        append!(a, [steps])
        i = stop + 1
    end
    return a
end

# TODO use this function only for URDBrate
function remove_tiers_from_urdb_rate(u::URDBrate)
    # tariff args: have to validate that there are no tiers
    if length(u.energy_tier_limits) > 1
        @warn "Energy rate contains tiers. Using the first tier!"
    end
    elec_rates = vec(u.energy_rates[1,:])

    if u.n_monthly_demand_tiers > 1
        @warn "Monthly demand rate contains tiers. Using the last tier!"
    end
    if u.n_monthly_demand_tiers > 0
        demand_rates_monthly = vec(u.monthly_demand_rates[:,u.n_monthly_demand_tiers])
    else
        demand_rates_monthly = vec(u.monthly_demand_rates)  # 0×0 Array{Float64,2}
    end

    if u.n_tou_demand_tiers > 1
        @warn "TOU demand rate contains tiers. Using the last tier!"
    end
    if u.n_tou_demand_tiers > 0
        demand_rates = vec(u.tou_demand_rates[:,u.n_tou_demand_tiers])
    else
        demand_rates = vec(u.tou_demand_rates)
    end

    return elec_rates, demand_rates_monthly, demand_rates
end


lattice = load_lattice(joinpath(@__DIR__, "Al40.json"))
atoms   = load_atoms(  joinpath(@__DIR__, "Al40.json"))
atoms   = attach_pseudos(atoms, Al="hgh/pbe/al-q3.hgh")

cases = standard_cases()
cases["adaptive_05"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.5, α_min=0.05), )

function load_Al40_potential(basis)
    guessfile = joinpath(@__DIR__, "Al40_initV.jld2")
    if !isfile(guessfile)
        ρ0 = guess_density(basis)
        _, ham = DFTK.energy_hamiltonian(basis, nothing, nothing; ρ=ρ0)
        V = DFTK.total_local_potential(ham)
        V .+= 3.0rand(Float64, size(V)...)

        jldopen(guessfile, "w") do file
            file["V"] = V
        end
    end

    jldopen(guessfile, "r") do file
        file["V"]
    end
end

config[@__DIR__] = (
    directory        =  @__DIR__,
    lattice          = lattice,
    atoms            = atoms,
    magnetic_moments = [],
    #
    kwargs_model     = (smearing=Smearing.Gaussian(), temperature=0.001),
    Ecut             = 20,
    kgrid            = [3, 3, 1],
    supersampling    = 1.5,
    potential        = load_Al40_potential,
    kwargs_scf       = (; ),
    mixing           = KerkerMixing(),
    cases            = cases,
)