module Errors

using JuLIP: Atoms, energy, forces, virial
using IPFitting: LsqDB, Dat, configtype, err_weighthook, eval_obs,
                      observation, hasobservation, observations,
                      vec_obs, tfor_observations, truncate_string
using IPFitting.DB: matrows
using FileIO, Printf, Base.Threads

using Statistics: quantile
using LinearAlgebra: norm
using ProgressMeter, DataFrames, PrettyTables
using Random: shuffle!

export add_fits!, add_fits_serial!, rmse, mae, rmse_table, rmse_table_wg, mae_table,
       lsqerrors, splittraintest

# ------------------------- USER INTERFACE FUNCTIONS -------------------------

rmse(data::Vector{Dat}; fitkey="fit") = _fiterrs(data; fitkey=fitkey, p=2)
mae(data::Vector{Dat}; fitkey="fit") = _fiterrs(data; fitkey=fitkey, p=1)

rmse(D::Dict) = D["rmse"], D["relrmse"]
mae(D::Dict) = D["mae"], D["relmae"]

rmse_table(data::AbstractVector{Dat}) = rmse_table(rmse(data)...)
mae_table( data::AbstractVector{Dat}) =  mae_table( mae(data)...)

rmse_table(D::Dict) = rmse_table(rmse(D)...)

rmse_table_wg(D::Dict, weights) = rmse_table_wg(rmse(D)..., weights)

mae_table(D::Dict) =  mae_table( mae(D)...)

rmse_table(errs::Dict, errs_rel::Dict; configtypes=:) =
   _err_table(errs, errs_rel, "RMSE", configtypes)

rmse_table_wg(errs::Dict, errs_rel::Dict, weights; configtypes=:) =
      _err_table_wg(errs, errs_rel, "RMSE WEIGHTED TABLE", weights, configtypes)

mae_table(errs::Dict, errs_rel::Dict) = _err_table(errs, errs_rel, "MAE")

"""
`add_fits(IP, data::Vector{Dat}) -> nothing`

for each `d in data`, compute the observations with the IP (fitted to this
data) and store in `d.info["fit"][...]`

The key `"fit"` can be replaced with any string, through the kwarg
`fitkey`, e.g.,
```julia
data = ...
add_fits!(PIP4Benv, data; fitkey = "PIP4Benv")
add_fits!(GAP, data; fitkey = "GAP2010")
```
"""
function add_fits!(IP, configs::Vector{Dat}; fitkey = "fit")
   # create the nec essary dictionaries
   for cfg in configs
      cfg.info[fitkey] = Dict{String, Vector{Float64}}()
   end
   tfor_observations( configs,
         (n, okey, cfg, lck) -> begin
            obs = vec_obs(okey, eval_obs(okey, IP, cfg))
            lock(lck)
            cfg.info[fitkey][okey] = obs
            unlock(lck)
         end;
         msg="Add Fit info to configs")

   return nothing
end

function add_fits_serial!(IP, configs::Vector{Dat}; fitkey = "fit")
   # create the nec essary dictionaries
   for cfg in configs
      cfg.info[fitkey] = Dict{String, Vector{Float64}}()
   end
   for (okey, cfg, _) in observations(configs)
      obs = vec_obs(okey, eval_obs(okey, IP, cfg))
      cfg.info[fitkey][okey] = obs
   end
   return nothing
end

function _relerr(relerrs, ct, ot)
   if haskey(relerrs[ct], ot)
      if relerrs[ct][ot] isa Number
         return 100*relerrs[ct][ot]
      end
   end
   return NaN
end

function _err(errs, ct, ot)
   if  haskey(errs[ct], ot)
      if errs[ct][ot] isa Number
         return errs[ct][ot]
      end
   end
   return NaN
end

# function _err_table_wg(errs, relerrs, title, weights, configtypes=:)
#    lentitle = length(title)
#    print("┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n")
#    print("┃ "); printstyled(title; bold=true);
#    print(repeat(' ', 70-3-lentitle)); print("┃\n")
#    print("┣━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━┫\n")
#    print("┃ config type  ┃     E [meV]     │     F [eV/A]    │     V [meV]     ┃\n")
#    print("┠──────────────╂────────┬────────┼────────┬────────┼────────┬────────┨\n")
#    s_set = ""
#    for ct in keys(errs)  # ct in configtypes
#       if (configtypes != :)
#          if !(ct in configtypes)
#             continue
#          end
#       end
#
#       if haskey(weights, ct)
#          w = weights[ct]
#       else
#          w = weights["default"]
#       end
#
#       s = @sprintf("┃ %12s ┃ %6.2f ┊ %5.3f%% │ %6.3f ┊ %5.2f%% │ %6.1f ┊ %5.2f%% ┃\n",
#          truncate_string(ct, 12),
#          _err(errs, ct, "E")*1000 * w["E"], _relerr(relerrs, ct, "E") * w["E"],
#          _err(errs, ct, "F") * w["F"], _relerr(relerrs, ct, "F") * w["F"],
#          _err(errs, ct, "V")*1000 * w["V"], _relerr(relerrs, ct, "V") * w["V"])
#       if ct == "set"
#          s_set = s
#       else
#          print(s)
#       end
#    end
#    if s_set != ""
#       print("┠──────────────╂────────┼────────┼────────┼────────┼────────┼────────┨\n")
#       print(s_set)
#    end
#       print("┗━━━━━━━━━━━━━━┻━━━━━━━━┷━━━━━━━━┷━━━━━━━━┷━━━━━━━━┷━━━━━━━━┷━━━━━━━━┛\n")
# end

function _err_table(errs, relerrs, title, configtypes=:)
   #@show relerrs
   lentitle = length(title)
   datf = DataFrame()
   datf[!, Symbol("cf")] = Any[]
   datf[!, Symbol("E")]= Any[]
   datf[!, Symbol("F")]= Any[]
   datf[!, Symbol("Fr")]= Any[]
   datf[!, Symbol("V")]= Any[]

   #need to add relative errors!

   for ct in sort(collect(keys(errs)))
      if ct != "set"
         push!(datf, [ct, _err(errs, ct, "E")*1000, _err(errs, ct, "F"), _err(relerrs, ct, "F")*100, _err(errs, ct, "V")*1000])
      end
   end

   push!(datf, ["set", _err(errs, "set", "E")*1000, _err(errs, "set", "F"), _err(relerrs, "set", "F")*100, _err(errs, "set", "V")*1000])

   header = ["config type" "E [meV]" "F [eV/A]" "F [%]" "V[meV]"]

   pretty_table(datf, header, formatters = ft_printf("%5.3f"), body_hlines = [length(datf[!, 1])-1])

   # print("┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n")
   # print("┃ "); printstyled(title; bold=true);
   # print(repeat(' ', 70-3-lentitle)); print("┃\n")
   # print("┣━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━━┫\n")
   # print("┃ config type  ┃     E [meV]     │     F [eV/A]    │     V [meV]     ┃\n")
   # print("┠──────────────╂────────┬────────┼────────┬────────┼────────┬────────┨\n")
   # s_set = ""
   # for ct in keys(errs)  # ct in configtypes
   #    if (configtypes != :)
   #       if !(ct in configtypes)
   #          continue
   #       end
   #    end
   #
   #    s = @sprintf("┃ %12s ┃ %6.2f ┊ %5.3f%% │ %6.3f ┊ %5.2f%% │ %6.1f ┊ %5.2f%% ┃\n",
   #       truncate_string(ct, 12),
   #       _err(errs, ct, "E")*1000, _relerr(relerrs, ct, "E"),
   #       _err(errs, ct, "F"), _relerr(relerrs, ct, "F"),
   #       _err(errs, ct, "V")*1000, _relerr(relerrs, ct, "V"))
   #    if ct == "set"
   #       s_set = s
   #    else
   #       print(s)
   #    end
   # end
   # if s_set != ""
   #    print("┠──────────────╂────────┼────────┼────────┼────────┼────────┼────────┨\n")
   #    print(s_set)
   # end
   #    print("┗━━━━━━━━━━━━━━┻━━━━━━━━┷━━━━━━━━┷━━━━━━━━┷━━━━━━━━┷━━━━━━━━┷━━━━━━━━┛\n")
end

# ---------------------------------------------------------------------
# ------------------------- PRIVATE FUNCTIONS -------------------------
# ---------------------------------------------------------------------

"""
initialize an error dictionary
"""
function errdict(configtypes)
   D = Dict{String, Dict{String, Float64}}()
   for ct in configtypes
      D[ct] = Dict{String, Float64}()
   end
   D["set"] = Dict{String, Float64}()
   return D
end

function _fiterrs(configs::Vector{Dat}; fitkey="fit", p=2)
   configtypes = unique(configtype.(configs))
   err = errdict(configtypes)
   relerr = errdict(configtypes)
   nrm = errdict(configtypes)
   len = errdict(configtypes)

   # assemble the errors
   for (okey, dat, _) in observations(configs)
      ct = configtype(dat)
      if !haskey(err[ct], okey)
         len[ct][okey] = 0.0
         nrm[ct][okey] = 0.0
         err[ct][okey] = 0.0
      end
      if !haskey(err["set"], okey)
         len["set"][okey] = 0.0
         nrm["set"][okey] = 0.0
         err["set"][okey] = 0.0
      end
      w = err_weighthook(okey, dat)  # the default weight for this observation
      exval = dat.D[okey] * w
      fitval = dat.info[fitkey][okey] * w
      len[ct][okey] += length(exval)
      len["set"][okey] += length(exval)
      nrm[ct][okey] += norm(exval, p)^p
      nrm["set"][okey] += norm(exval, p)^p
      err[ct][okey] += norm(exval-fitval, p)^p
      err["set"][okey] += norm(exval-fitval, p)^p
   end

   # normalise correctly
   for ct in keys(err)
      for okey in keys(err[ct])
         err[ct][okey] = (err[ct][okey]/len[ct][okey])^(1/p)
         nrm[ct][okey] = (nrm[ct][okey]/len[ct][okey])^(1/p)
         relerr[ct][okey] = err[ct][okey] / nrm[ct][okey]
      end
   end

   return err, relerr
end

# TODO: the next fucntion is to be retired!!

@noinline function lsqerrors(db, c, Ibasis; cfgtypes = Colon(), Vref = nothing,
                                            Icfg=Colon())

   cfgtypes isa Colon ?  cfgtypes = configtypes(db) :
                         cfgtypes = collect(cfgtypes)

   # create the dict for the fit errors
   errs = Dict("rmse" => errdict([cfgtypes; "set"]),
               "relrmse" => errdict([cfgtypes; "set"]),
               "mae"  => errdict([cfgtypes; "set"]),
               "maxe" => errdict([cfgtypes; "set"]),
               "nrm2" => errdict([cfgtypes; "set"]),
               "nrminf" => errdict([cfgtypes; "set"]),
               "nrm1" => errdict([cfgtypes; "set"]) )
   lengths = Dict{String, Dict{String, Int}}()
   for ct in [cfgtypes; "set"]
      lengths[ct] = Dict{String, Int}()
   end

   for (okey, dat, icfg) in observations(db, Icfg)

      ct = configtype(dat)
      if !(ct in cfgtypes)
         continue
      end
      if !haskey(errs["rmse"][ct], okey)
         errs["rmse"][ct][okey] = 0.0
         errs["nrm2"][ct][okey] = 0.0
         errs["mae"][ct][okey] = 0.0
         errs["nrm1"][ct][okey] = 0.0
         errs["maxe"][ct][okey] = 0.0
         errs["nrminf"][ct][okey] = 0.0
         lengths[ct][okey] = 0
      end
      if !haskey(errs["rmse"]["set"], okey)
         errs["rmse"]["set"][okey] = 0.0
         errs["nrm2"]["set"][okey] = 0.0
         errs["mae"]["set"][okey] = 0.0
         errs["nrm1"]["set"][okey] = 0.0
         errs["maxe"]["set"][okey] = 0.0
         errs["nrminf"]["set"][okey] = 0.0
         lengths["set"][okey] = 0
      end
      # assemble the observation vector
      y = Float64[]
      o = observation(dat, okey)
      # correct the observation if we are fitting from a reference potential
      if Vref != nothing
         o -= vec_obs(okey, eval_obs(okey, Vref, dat))
      end
      append!(y, o)
      # get the lsq block
      irows = matrows(dat, okey)
      block = db.Ψ[irows, Ibasis]
      # the error on this particular data-point
      e = block * c - y
      # compute the various errors for this ct/okey combination
      w = err_weighthook(okey, dat)
      errs["rmse"][ct][okey] += norm(w .* e)^2
      errs["nrm2"][ct][okey] += norm(w .* y)^2
      errs["mae"][ct][okey]  += norm(w .* e, 1)
      errs["nrm1"][ct][okey] += norm(w .* y, 1)
      errs["maxe"][ct][okey] = max(errs["maxe"][ct][okey], norm(e, Inf))
      errs["nrminf"][ct][okey] = max(errs["nrminf"][ct][okey], norm(y, Inf))
      lengths[ct][okey] += length(y)
      errs["rmse"]["set"][okey] += norm(w .* e)^2
      errs["nrm2"]["set"][okey] += norm(w .* y)^2
      errs["mae"]["set"][okey]  += norm(w .* e, 1)
      errs["nrm1"]["set"][okey] += norm(w .* y, 1)
      errs["maxe"]["set"][okey] = max(errs["maxe"]["set"][okey], norm(e, Inf))
      errs["nrminf"]["set"][okey] = max(errs["nrminf"]["set"][okey], norm(y, Inf))
      lengths["set"][okey] += length(y)
   end

   for ct in [cfgtypes; "set"]
      for okey in keys(errs["rmse"][ct])
         len = lengths[ct][okey]
         errs["rmse"][ct][okey] = sqrt( errs["rmse"][ct][okey] / len )
         errs["nrm2"][ct][okey] = sqrt( errs["nrm2"][ct][okey] / len )
         errs["mae"][ct][okey] = errs["mae"][ct][okey] / len
         errs["nrm1"][ct][okey] = errs["nrm1"][ct][okey] / len
         errs["relrmse"][ct][okey] = errs["rmse"][ct][okey] / errs["nrm2"][ct][okey]
      end
   end

   return errs
end

splittraintest(db::LsqDB, ptrain=0.8) = splittraintest(db.configs, ptrain)

function splittraintest(cfgs::AbstractVector{Dat}, ptrain=0.8)
   cfgtypes = unique(configtype.(cfgs))
   Itrain = Int[]
   Itest = Int[]
   for ct in cfgtypes
      # find all indices of configurations with the current configtype
      Ict = findall(configtype.(cfgs) .== ct)
      # randomly permute the indices
      shuffle!(Ict)
      # first 80% go into train, remaining 20% into test
      ntrain = ceil(Int, length(Ict) * ptrain)
      append!(Itrain, Ict[1:ntrain])
      append!(Itest, Ict[ntrain+1:end])
   end
   return Itrain, Itest
end

end


"""
Gallery of RTG models
"""

abstract type RTGmodel end

include("utils.jl")
include("actM1.jl")
include("rtgM1.jl")
include("rtgM2.jl")
include("rtgM3.jl")
include("rtgM4.jl")


@auto_hash_equals mutable struct Grid
    seed::Integer # to allow reproducibility
    buses::Vector{Bus}
    trans_lines::Vector{TransLine}
    substations::Vector{Substation}
    bus_conn::AbstractArray{Bool,2} # bus connectivity matrix
    sub_conn::AbstractArray{Int,2} # this is slightly more than a substation connectivity
                # matrix, as it counts how many connections are there between substations.

    function Grid(
        seed::Integer,
        bus = [],
        line = [],
        sub = [],
        conm = falses(0, 0),
        subm = zeros(Int, 0, 0),
    )
         return new(seed, bus, line, sub, conm, subm)
    end
end

buses(g::Grid) = g.buses
trans_lines(g::Grid) = g.trans_lines
substations(g::Grid) = g.substations
adjacency(g::Grid) = g.bus_conn
sub_connectivity(g::Grid) = g.sub_conn

"""
    add_load!(grid::Grid, args...; reconnect = false)

Add LoadBus to `grid` by calling the `LoadBus(coords, load, volt, pop, connected_to =
Set(), connections = Set())` method with `args...` as arguments. If `reconnect = true`, all
connections will be remade.
"""
function add_load!(grid::Grid, args...; reconnect = false)
    bus = LoadBus(length(grid.buses)+1, args...)
    add_bus!(grid, bus, reconnect = reconnect)
end

"""
    add_gen!(grid::Grid, args...; reconnect = false)

Add Genbus to `grid` by calling the `GenBus(coords, gen, volt, tech, connected_to =
Set(), connections = Set(), pfactor = -1, summgen = -1, wintgen = -1, gens = [])` method
with `args...` as arguments. If `reconnect = true`, all connections will be remade.
"""
function add_gen!(grid::Grid, args...; reconnect = false)
    bus = GenBus(length(grid.buses)+1, args...)
    add_bus!(grid, bus, reconnect = reconnect)
end

function add_bus!(grid::Grid, bus::Union{LoadBus, GenBus}; reconnect = false)
    # The ids are tricky in that they have to correspond to the actual position of the bus
    # in the vector, but, at the same time, it is convenient to have all buses of the same
    # type grouped and ordered.
    loads = [b for b in grid.buses if isa(b, LoadBus)]
    gens = [b for b in grid.buses if isa(b, GenBus)]
    load_ids = [b.id for b in loads]
    gen_ids = [b.id for b in gens]
    ls = sortperm(load_ids)
    gs = sortperm(gen_ids)
    loads = loads[ls]
    gens = gens[gs]
    ml_id = isempty(loads) ? 0 : loads[end].id
    mg_id = isempty(gens) ? 0 : gens[end].id
    first = mg_id > ml_id ? LoadBus : GenBus
    second_empty = isempty(loads) || isempty(gens)
    if !second_empty && isa(bus, first)
        if first == LoadBus
            # Add new LoadBus
            push!(
                loads,
                LoadBus(
                    ml_id + 1,
                    bus.coords,
                    bus.load,
                    bus.voltage,
                    bus.population,
                    bus.connected_to,
                    bus.connections
                )
            )
            # Update IDs of all GenBuses
            new_gens = GenBus[]
            for g in gens
                push!(
                    new_gens,
                    GenBus(
                        g.id + 1,
                        g.coords,
                        g.generation,
                        g.voltage,
                        g.tech_type,
                        g.connected_to,
                        g.connections,
                        g.pfactor,
                        g.summgen,
                        g.wintgen,
                        g.gens
                    )
                )
            end
            # Update grid
            grid.buses = vcat(loads, new_gens)
        else
            # Add new GenBus
            push!(
                gens,
                GenBus(
                    mg_id + 1,
                    bus.coords,
                    bus.generation,
                    bus.voltage,
                    bus.tech_type,
                    bus.connected_to,
                    bus.connections,
                    bus.pfactor,
                    bus.summgen,
                    bus.wintgen,
                    bus.gens
                )
            )
            # Update the IDs of all LoadBuses
            new_loads = LoadBus[]
            for l in loads
                push!(
                    new_loads,
                    LoadBus(
                        l.id + 1,
                        l.coords,
                        l.load,
                        l.voltage,
                        l.population,
                        l.connected_to,
                        l.connections
                    )
                )
            end
            # Update grid
            grid.buses = vcat(gens, new_loads)
        end
    elseif bus.id != length(grid.buses) + 1
        if isa(bus, LoadBus)
            push!(
                grid.buses,
                LoadBus(
                    length(grid.buses) + 1,
                    bus.coords,
                    bus.load,
                    bus.voltage,
                    bus.population,
                    bus.connected_to,
                    bus.connections
                )
            )
        elseif isa(bus, GenBus)
            push!(
                buses(grid),
                GenBus(
                    length(grid.buses) + 1,
                    bus.coords,
                    bus.generation,
                    bus.voltage,
                    bus.tech_type,
                    bus.connected_to,
                    bus.connections,
                    bus.pfactor,
                    bus.summgen,
                    bus.wintgen,
                    bus.gens
                )
            )
        end
    else
        push!(grid.buses, bus)
    end
    if reconnect
        # Break previous connections
        grid.bus_conn = falses(length(grid.buses), length(grid.buses))
        for b in grid.buses
            empty!(b.connected_to)
            empty!(b.connections)
        end
        # Redo connections including the new bus
        connect!(grid)
    elseif size(grid.bus_conn)[1] > 0 # Grid has already been connected
        # Expand connectivity matrix to support the new bus
        bottom_row = zeros(length(grid.buses) - 1)
        right_col = zeros(length(grid.buses))
        grid.bus_conn = vcat(grid.bus_conn, bottom_row')
        grid.bus_conn = hcat(grid.bus_conn, right_col)
        # Check if there are connections to be made
        for b in bus.connected_to
            if b.id > 0
                grid.bus_conn[b.id, bus.id] = 1
                grid.bus_conn[bus.id, b.id] = 1
            end
        end
    end
end

"""
    add_substation!(grid::Grid, args...; reconnect = false)

Add Substation to `grid` by calling the `Substation(coords, volts, load, gen, pop,
con = Set(), group = [])` method with `args...` as arguments.
"""
function add_substation!(grid::Grid, args...)
    sub = Substation(length(grid.substations)+1, args...)
    push!(grid.substations, sub)
    # Break previous connections
    grid.sub_conn = zeros(length(grid.substations), length(grid.substations))
    for s in grid.substations
        s.connected_to = Set()
    end
    # Reconnect
    connect_subs!(grid)
end

Grid() = Grid(rand(0:round(Int, typemax(Int) / 2)))

"""
    twst!(grid, k)

Execute Tunable Weight Spannning Tree algorithm for the grid using weight 'k'.

REFERENCE: Soltan, Saleh, and Gil Zussman. "Generation of synthetic spatially
embedded power grid networks." arXiv:1508.04447 [cs.SY], Aug. 2015.
"""
function twst!(grid::Grid, k)
    node_locations = [b.coords for b in buses(grid)]
    n = length(buses(grid))
    ind = collect(1:n)
    permut = Vector{Int}(undef, n)
    unorm_prob = [haversine(bus.coords, mean(node_locations))^(-k) for bus in buses(grid)]
    seed!(grid.seed + 15) # Just so we don't keep returning to the same point.
    rand_range = 0:PREC:(1-PREC)
    for step in 1:n # sets node permutation
        probs = unorm_prob / sum(unorm_prob)
        rd = rand(rand_range)
        cumul = 0.0
        bus_ind = -1
        for i in 1:length(probs) # selects bus
            cumul += probs[i]
            if cumul >= rd
                bus_ind = i
                break
            end
        end
        permut[step] = ind[bus_ind]
        deleteat!(unorm_prob, bus_ind)
        deleteat!(ind, bus_ind)
    end

    bus_conn = zeros(Bool, n, n)
    for i in 2:length(permut) # builds connections
        distances = [
            haversine(node_locations[permut[i]],
            node_locations[permut[j]]) for j in 1:(i-1)
            ]
        ind = findmin(distances)[2] # nearest neighbor with j < i
        bus_conn[permut[i], permut[ind]] = true
        bus_conn[permut[ind], permut[i]] = true
    end

    grid.bus_conn = sparse(bus_conn)
end

"""
    reinforce!(grid, m, a, b, g, t, N)

Execute the Reinforcement procedure to increase robustness of a grid.

# Arguments:

* m: total number of connections
* a, b, g, t: paramaters of the model (see REFERENCE)
* N: number of nearest neighbors in the average distance computation

REFERENCE: Soltan, Saleh, and Gil Zussman. "Generation of synthetic spatially
embedded power grid networks." arXiv:1508.04447 [cs.SY], Aug. 2015.
"""
function reinforce!(grid::Grid, m, a, b, g, t, N)
    n = length(buses(grid))
    if N > (n - 1)
        throw(ArgumentError("ERROR in Reinforcement procedure (reinforce!()):",
              "Number of nearest neighbors has to be less than the number or nodes"))
    end
    if m < 0
        m = round(Int, 1.22 * n) # using Overbye's approximation here
    end
    av_dist = sizehint!(Real[], length(buses(grid)))
    for bus in buses(grid) # computes averge distances to neighbors
        dists = [haversine(bus, b2) for b2 in buses(grid)]
        sort!(dists)
        deleteat!(dists, 1) # minimum distance will always be the self-distance = 0
        push!(av_dist, mean(dists[1:N]))
    end
    unorm_probs = Real[]
    inds = Int[]
    degrees = sum(grid.bus_conn, dims=1) # this gives the node degrees
    degrees = float(reshape(degrees, length(degrees))) # Just so we can zip it later. Using
    # float here in order to avoid problems when doing Int^(-t).
    if n >= LARGE_GRID_N # only sample nodes with degree < 3
        unorm_probs = [av_dist[i]^(-a) for i in 1:n if degrees[i] < 3]
        inds = [i for i in 1:n if degrees[i] < 3]
    else # sample all nodes
        unorm_probs = [degrees[i]^(-t)*av_dist[i]^(-a) for i in 1:n]
        inds = collect(1:n)
    end
    seed!(grid.seed + 479) # Just so we don't keep returning to the same point.
    rand_range = 0:PREC:(1-PREC)
    count = 1
    while count <= (m - n + 1)
        probs = unorm_probs / sum(unorm_probs)
        rd = rand(rand_range)
        cumul = 0.0
        bus_ind = -1
        for i in 1:length(probs) #selects bus
            cumul += probs[i]
            if cumul >= rd
                bus_ind = i # selected bus is inds[bus_ind]
                break
            end
        end
        b_dists = [haversine(buses(grid)[inds[bus_ind]], bus) for bus in buses(grid)]
        b_probs = [dist^(-b) * deg^g for (dist, deg) in zip(b_dists, degrees)]
        b_probs = [bp == Inf ? 0 : bp for bp in b_probs] # If buses overlap we may
                      # have multiple Infs (besides the Inf from the self-distance).
                      # Ideally buses should never overlap!
        b_probs = b_probs / sum(b_probs)
        rd = rand(rand_range)
        cumul = 0.0
        b_ind = -1
        for i in 1:length(b_probs) # draws node to which connect
            cumul += b_probs[i]
            if cumul >= rd
                b_ind = i
                break
            end
        end
        if b_ind == inds[bus_ind] # can't connect a bus to itself
            if b_ind == n
                b_ind -= 1
            else
                b_ind += 1
            end
        end
        if grid.bus_conn[b_ind, inds[bus_ind]] == false # not yet connected
            grid.bus_conn[b_ind, inds[bus_ind]] = true
            grid.bus_conn[inds[bus_ind], b_ind] = true
            deleteat!(unorm_probs, bus_ind)
            deleteat!(inds, bus_ind)
            count += 1
        end
    end
end

"""
    connect_buses!(grid)

Create reference links between buses that are connected according to the
connectivity matrix of the grid.
"""
function connect_buses!(grid::Grid)
    for j in 1:(length(buses(grid)) - 1), i in (j + 1):length(buses(grid))
        if grid.bus_conn[i, j] == true
            push!(buses(grid)[i].connected_to, buses(grid)[j])
            push!(buses(grid)[j].connected_to, buses(grid)[i])
        end
    end
end

"""
    connect!(grid::Grid; k=2.5, m=-1, a=1, b=3.2, g=2.5, t=2, N=10)

Connect buses in an electric grid.

# Arguments:
* k: weight of the spanning tree (see REFERENCE)
* m: total number of connections (computed from the # of nodes if default)
* a, b, g, t: paramaters of the model (see REFERENCE)
* N: number of nearest neighbors in the average distance computation

REFERENCE: Soltan, Saleh, and Gil Zussman. "Generation of synthetic spatially
embedded power grid networks." arXiv:1508.04447 [cs.SY], Aug. 2015.
"""
function connect!(
    grid::Grid;
    k=2.5,
    m=-1,
    a=1,
    b=3.2,
    g=2.5,
    t=2.0,
    N=10,
)
    twst!(grid, k)
    reinforce!(grid, m, a, b, g, t, N)
    connect_buses!(grid)
end

"""
    create_lines!(grid::Grid; impedfunc = linear_imped, capfunc = volt_cap)

Create transmission lines for a synthetic grid. Line impedancies are determined
by impedfunc and line capacities are calculated via capfunc.
This function will only work after the grid has been connected through the
connect!() function.
"""
function create_lines!(grid::Grid; impedfunc = linear_imped, capfunc = volt_cap)
    seed!(grid.seed + 16) # Just so we don't keep returning to the same point.
    for j in 1:(length(buses(grid)) - 1), i in (j + 1):length(buses(grid))
        if grid.bus_conn[i, j] == true
            a, b = buses(grid)[i], buses(grid)[j]
            push!(grid.trans_lines, TransLine(a, b; impedfunc=impedfunc, capfunc=capfunc))
            push!(buses(grid)[i].connections, grid.trans_lines[end])
            push!(buses(grid)[j].connections, grid.trans_lines[end])
        end
    end
end

function merge!(grid::Grid, sub1::Substation, sub2::Substation)
    if sub1.load > 0 || sub2.load > 0
        tot_pop = (sub1.population + sub2.population)
        lat = (sub1.population * sub1.coords.lat + sub2.population
            * sub2.coords.lat) / tot_pop
        lon = (sub1.population * sub1.coords.lon + sub2.population
            * sub2.coords.lon) / tot_pop
        sub1.coords = LatLon(lat, lon)
    else
        tot_gen = sub1.generation + sub2.generation
        lat = (sub1.generation * sub1.coords.lat + sub2.generation
            * sub2.coords.lat) / tot_gen
        lon = (sub1.generation * sub1.coords.lon + sub2.generation
            * sub2.coords.lon) / tot_gen
        sub1.coords = LatLon(lat, lon)
    end
    for v in sub2.voltages
        if !(v in sub1.voltages)
            push!(sub1.voltages, v)
        end
    end
    sub1.load += sub2.load
    sub1.generation += sub2.generation
    sub1.population += sub2.population
    sub1.grouping = vcat(sub1.grouping, sub2.grouping)
    deleteat!(substations(grid), findfirst(substations(grid), sub2))
end

function merge!(grid::Grid, i1::Int, i2::Int)
    sub1 = substations(grid)[i1]
    sub2 = substations(grid)[i2]
    if sub1.load > 0 || sub2.load > 0
        tot_pop = (sub1.population + sub2.population)
        lat = (sub1.population * sub1.coords.lat + sub2.population
            * sub2.coords.lat) / tot_pop
        lon = (sub1.population * sub1.coords.lon + sub2.population
            * sub2.coords.lon) / tot_pop
        sub1.coords = LatLon(lat, lon)
    else
        tot_gen = sub1.generation + sub2.generation
        lat = (sub1.generation * sub1.coords.lat + sub2.generation
            * sub2.coords.lat) / tot_gen
        lon = (sub1.generation * sub1.coords.lon + sub2.generation
            * sub2.coords.lon) / tot_gen
        sub1.coords = LatLon(lat, lon)
    end
    for v in sub2.voltages
        if !(v in sub1.voltages)
            push!(sub1.voltages, v)
        end
    end
    sub1.load += sub2.load
    sub1.generation += sub2.generation
    sub1.population += sub2.population
    sub1.grouping = vcat(sub1.grouping, sub2.grouping)
    deleteat!(substations(grid), i2)
end

"""
    connect_subs!(grid)

Determine which substations are connected to each other through their buses.
"""
function connect_subs!(grid::Grid)
    n = length(substations(grid))
    grid.sub_conn = zeros(Int, n, n)
    aux = copy(grid.bus_conn)
    aux = convert(Matrix{Int}, aux)
    for s in 1:n
        for j in 1:length(buses(grid)), b in substations(grid)[s].grouping
            if grid.bus_conn[b.id, j] == true
                aux[b.id, j] = s
            end
        end
    end

    for i in 1:(length(buses(grid)) - 1)
        for j in (i + 1):length(buses(grid))
            if aux[i, j] != aux[j, i]
                grid.sub_conn[aux[i, j], aux[j, i]] += 1 # This approach
                grid.sub_conn[aux[j, i], aux[i, j]] += 1 # not only gives
                # which substation is connected to which, but also gives how many
                # connections are there between them, i.e., gives a measure of
                # strength for the connection.
                push!(
                    substations(grid)[aux[i, j]].connected_to,
                    substations(grid)[aux[j, i]]
                )
                push!(
                    substations(grid)[aux[j, i]].connected_to,
                    substations(grid)[aux[i, j]]
                )
            end
        end
    end
    grid.sub_conn = sparse(grid.sub_conn)
end

function Grid(
        sub::Bool;
        lat = (33,35),
        long = (-95,-93),
    )
    grid = Grid()
    place_loads_from_zips!(grid; latlim = lat, longlim = long)
    place_gens_from_data!(grid; latlim = lat, longlim = long)
    connect!(grid)
    count = count_bus_type(grid)
    if sub # This is just one very crude approach to having substations automatically
           # created. Good for testing.
        cluster!(
            grid,
            round(Int, 0.3 * count[1]),
            round(Int, 0.05 * count[2]),
            round(Int, 0.5 * count[2])
        )
    end
    create_lines!(grid)
    return grid
end


export

    ProgressDisplay,
    StoreSnapshots,
    PointMicrophone

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

mutable struct ProgressDisplay
    dt::Float64
    bar::Progress
    showvalues::Bool

    function ProgressDisplay(dt::Number = .1, showvalues = true)
        hook = new()
        hook.dt = Float64(dt)
        hook.showvalues = showvalues
        hook
    end
end

function hook_init!(hook::ProgressDisplay, state::State, sim::Simulator)
    hook.bar = Progress(sim.maxiter, dt = hook.dt)
end

function hook_update!(hook::ProgressDisplay, state::State, sim::Simulator)
    if hook.showvalues
        next!(hook.bar; showvalues = [(:time,state.t), (:iter,state.iter)])
    else
        next!(hook.bar)
    end
end

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

mutable struct StoreSnapshots
    dt::Float64
    last::Float64
    index::Int
    buffer
    snapshots

    function StoreSnapshots(dt::Number)
        hook = new()
        hook.dt = Float64(dt)
        hook
    end
end

function hook_init!(hook::StoreSnapshots, state::State{N,T}, sim::Simulator) where {N,T}
    count = floor(Int, sim.duration / hook.dt)
    hook.last = typemin(Float64)
    hook.index = 1
    hook.buffer = zeros(T, size(state.current))
    hook.snapshots = zeros(T, (size(state.current)..., count))
    hook
end

function hook_update!(hook::StoreSnapshots, state::State{N}, sim::Simulator) where N
    if state.t - (hook.last + hook.dt) >= 0
        if hook.index <= size(hook.snapshots, ndims(hook.snapshots))
            copy!(hook.buffer, state.current)
            copy!(view(hook.snapshots, ntuple(_->:,N)..., hook.index), hook.buffer)
            hook.index += 1
        end
        hook.last = state.t
    end
    hook
end

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

mutable struct PointMicrophone{N}
    position::NTuple{N,Float64}
    startiter::Int
    gridpos::NTuple{N,Int}
    signal

    function PointMicrophone(position::Vararg{Number,N}) where N
        hook = new{N}()
        hook.position = Float64.(position)
        hook.startiter = 1
        hook
    end
end

function hook_init!(hook::PointMicrophone, state::State{N,T}, sim::Simulator) where {N,T}
    hook.signal = zeros(T, sim.maxiter)
    hook.startiter = state.iter
    hook.gridpos = floor.(Int, hook.position ./ sim.wave.dx) .+ 1
    hook
end

function hook_update!(hook::PointMicrophone, state::State{N}, sim::Simulator) where N
    hook.signal[state.iter - hook.startiter] = state.current[hook.gridpos...]
    hook
end

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

mutable struct PointSpeaker{N,S<:AbstractVector}
    position::NTuple{N,Float64}
    gridpos::NTuple{N,Int}
    signal::S

    function PointSpeaker(position::NTuple{N,Number}, signal::S) where {N,S<:AbstractVector}
        hook = new{N,S}()
        hook.position = Float64.(position)
        hook.signal = signal
        hook
    end
end

function hook_init!(hook::PointSpeaker, state::State{N,T}, sim::Simulator) where {N,T}
    hook.gridpos = floor.(Int, hook.position ./ sim.wave.dx) .+ 1
    hook
end

function hook_update!(hook::PointSpeaker, state::State{N}, sim::Simulator) where N
    if state.iter <= length(hook.signal)
        state.current[hook.gridpos...] = hook.signal[state.iter]
    end
    hook
end


#! /bin/julia
# ~/.julia/dev/JuliaPlot/src/JuliaPlot.jl
using JuliaPlot
posstrings=readlines(`xrandr --listmonitors`)
pos = Tuple(
            parse.(Int64,
                   getindex.(
                             Ref(match(r".* (\d+)/\d+x(\d+)/\d+\+(\d+)\+(\d+) .*",posstring))
                             ,(3,4,1,2)
                            )
                  )
            for posstring in posstrings[2:end]
           )
multiMonitorJuliaPlot(;
                      pos = pos,
                     )


"""
    loaddata(file)

Loads methylation data.

`file` is a `RData` file containing a matrix named `betas.dasen` of size `(n x m)` for `n` samples and `m` probes.

Sample names and probe names taken from colnames and rownames respectively.

"""
function loaddata(file)

    R"""
        load($file)
        samples <- colnames(betas.dasen)
        probes  <- rownames(betas.dasen)
        beta    <- betas.dasen
    """

    @rget samples probes beta

    samples, probes, beta

end

"""
    loadsamplemeta(file)

Loads sample meta, a tab or commma separated file that must contain the following first threecolumns (in any order):

| Sample    | Sex   | Age   |  Filter  |
|-----------|-------|-------|----------|
|  Sample_1 | F     | 1.0   |   1      |
|  Sample_2 | M     | 2.0   |   0      |
|  Sample_3 | M     | 10.0  |   1      |

- `Sample` unique string identifier
- `Sex`  {M, F} or missing
- `Age` floating point age
- `Filter` optional column of `{1, 0}`s to exclude samples, when `Filter = 0` sample is excluded.

`Sample` should match `colnames` of RData file.

"""
loadsamplemeta(file) = loadmeta(file, [:Sample, :Sex, :Age], "sample")

"""
    loadprobemeta(file)

Loads probe meta, a tab or commma separated file that must contain the following columns (in any order):

| Probe      |
|------------|
| cg00000029 |
| cg00000103 |
| cg00000109 |

Can be used to override the `rownames` in `RData` file.

"""
loadprobemeta(file) = loadmeta(file, [:Probe], "probe")

"""
    loadmeta(file, mandatorycols, label)

    Helper function used to load metafiles
"""
function loadmeta(file, mandatorycols, label)
    meta = CSV.read(file, DataFrame)

    if length(intersect(mandatorycols, propertynames(meta))) != length(mandatorycols)
        missingcols = setdist(mandatorycols, propertynames(meta))
        colstring = join(names(meta), ", ")
        mcolstring = join(missingcols, ", ")
        error("Loading $label meta. Found cols: $colstring\nMissing: $mcolstring")
    end
    meta
end

"""
    loadall(samplemetafile, betafile)

Loads samplefmetafile and `RData` file `betafile` ensures samples details correspond and filters samples
"""
function loadall(samplemetafile, betafile, verbose=true)

    verbose && println("[GPM]\tLoaded sample meta file ", samplemetafile)
    verbose && println("[GPM]\tLoaded probe data file ", betafile)
    meta = loadsamplemeta(samplemetafile)
    samples, probes, beta = loaddata(betafile)

    if meta.Sample != samples
        error("Sample meta does not correspond to beta colnames")
    end

    verbose && println("[GPM]\tLoaded meta information for ", size(meta, 1), " samples.")
    verbose && println("[GPM]\tLoaded probe data on  ", size(beta, 1), " probes and ", size(beta, 2), " samples.")

    ### Filter samples
    if :Filter ∈ propertynames(meta)
        ind = meta.Filter .== 1

        if verbose
            println("[GPM]\tExcluding ", sum(.!ind), " samples:")
            display(meta[.!ind, :])
        end
        meta = meta[ind, :]
        beta = beta[:, ind]

    end

    ### Converts row major R to column major Julia
    meta, probes, Matrix(beta')
end


using CSV
using DataFrames
using DelimitedFiles

include("constants.jl")

#path_dat = "C://Users//rapiduser//Documents//GitHub//grace-arpae//UnitCommitment_BAU"
#cd(path_dat)

##
# Importing input data from the input spreadsheets
# Generators' specifications
DF_Generators = CSV.read(".//inputs//data_generators.csv", DataFrame);
# Generators location map: if a generator g is located in zone z Map_Gens[g,z]=1; and 0 otherwise
Map_Gens = readdlm(".//inputs//location_generators.csv", ','; header = true);

#Peaker Units' specifications and location
DF_Peakers = CSV.read(".//inputs//data_peakers.csv", DataFrame);
# Peakers location map: if a peaker p is located in zone z Map_Gens[p,z]=1; # and 0 otherwise
Map_Peakers = readdlm(".//inputs//location_peakers.csv", ','; header = true);

# Storage Units' specification and location
DF_Storage = CSV.read(".//inputs//data_storage.csv", DataFrame);
Map_Storage = readdlm(".//inputs//location_storage.csv", ','; header = true);

# Energy demand at each location
FUCR_Demands = readdlm(".//inputs//data_demand.csv", ','; header = true);
SUCR_Demands = readdlm(".//inputs//data_demand_updated.csv", ','; header = true);
BUCR_Demands = readdlm(".//inputs//data_demand_actual.csv", ','; header = true);

# solar generation data at each location
FUCR_SolarGs = readdlm(".//inputs//data_solar.csv", ','; header = true);
SUCR_SolarGs = readdlm(".//inputs//data_solar_updated.csv", ','; header = true);
BUCR_SolarGs = readdlm(".//inputs//data_solar_actual.csv", ','; header = true);

# wind energy data for each location
FUCR_WindGs = readdlm(".//inputs//data_wind.csv", ','; header = true);
SUCR_WindGs = readdlm(".//inputs//data_wind_updated.csv", ','; header = true);
BUCR_WindGs = readdlm(".//inputs//data_wind_actual.csv", ','; header = true);

#hydro generation data for each location
FUCR_HydroGs = readdlm(".//inputs//data_hydro.csv", ','; header = true);
SUCR_HydroGs = readdlm(".//inputs//data_hydro_updated.csv", ','; header = true);
BUCR_HydroGs = readdlm(".//inputs//data_hydro_actual.csv", ','; header = true);

#nuclear generation timeseries for each location
FUCR_NuclearGs = readdlm(".//inputs//data_nuclear.csv", ','; header = true);
SUCR_NuclearGs = readdlm(".//inputs//data_nuclear_updated.csv", ','; header = true);
BUCR_NuclearGs = readdlm(".//inputs//data_nuclear_actual.csv", ','; header = true);

#Cogenerators' generation timeseries for each location
FUCR_CogenGs = readdlm(".//inputs//data_cogen.csv", ','; header = true);
SUCR_CogenGs = readdlm(".//inputs//data_cogen_updated.csv", ','; header = true);
BUCR_CogenGs = readdlm(".//inputs//data_cogen_actual.csv", ','; header = true);

TranC = readdlm(".//inputs//LineCapacity.csv", ','; header = true);
TranS = readdlm(".//inputs//LineSusceptance.csv", ','; header = true);
Reserve_Reqs = readdlm(".//inputs//data_reserve_reqs.csv", ','; header = true);

FuelPrice = readdlm(".//inputs//data_fuel_price.csv", ','; header = true);
FuelPricePeakers = readdlm(".//inputs//data_fuel_price_peakers.csv", ','; header = true);

FuelPrice_head = FuelPrice[2];
FuelPrice = FuelPrice[1];
FuelPrice = Array{Float64}(FuelPrice[2:GENS+1, 4:368]);

FuelPricePeakers_head = FuelPricePeakers[2];
FuelPricePeakers = FuelPricePeakers[1];
FuelPricePeakers = Array{Float64}(FuelPricePeakers[2:PEAKERS+1, 4:368]);

# Reorganize data
FUCR_Demands_head = FUCR_Demands[2];
FUCR_Demands = FUCR_Demands[1];
FUCR_Demands = Array{Int64}(FUCR_Demands[1:60313, 4:(4+N_ZONES-1)]);

SUCR_Demands_head = SUCR_Demands[2];
SUCR_Demands = SUCR_Demands[1];
SUCR_Demands = Array{Int64}(SUCR_Demands[1:60313, 4:(4+N_ZONES-1)]);

BUCR_Demands_head = BUCR_Demands[2];
BUCR_Demands = BUCR_Demands[1];
BUCR_Demands = BUCR_Demands[1:8760, 3:(3+N_ZONES-1)];
BUCR_Demands = Array{Int64}(BUCR_Demands);

# Solar
FUCR_SolarGs_head = FUCR_SolarGs[2];
FUCR_SolarGs = FUCR_SolarGs[1];
FUCR_SolarGs = Array{Int64}(FUCR_SolarGs[1:8760, 3:(3+N_ZONES-1)]);

SUCR_SolarGs_head = SUCR_SolarGs[2];
SUCR_SolarGs = SUCR_SolarGs[1];
SUCR_SolarGs = SUCR_SolarGs[1:8760, 3:(3+N_ZONES-1)];
SUCR_SolarGs = Array{Int64}(SUCR_SolarGs);

BUCR_SolarGs_head = BUCR_SolarGs[2];
BUCR_SolarGs = BUCR_SolarGs[1];
BUCR_SolarGs = Array{Int64}(BUCR_SolarGs[1:8760, 3:(3+N_ZONES-1)]);

#Wind
FUCR_WindGs_head = FUCR_WindGs[2];
FUCR_WindGs = FUCR_WindGs[1];
FUCR_WindGs = Array{Int64}(FUCR_WindGs[1:8760, 3:(3+N_ZONES-1)]);

SUCR_WindGs_head = SUCR_WindGs[2];
SUCR_WindGs = SUCR_WindGs[1];
SUCR_WindGs = Array{Int64}(SUCR_WindGs[1:8760, 3:(3+N_ZONES-1)]);

BUCR_WindGs_head = BUCR_WindGs[2];
BUCR_WindGs = BUCR_WindGs[1];
BUCR_WindGs = Array{Int64}(BUCR_WindGs[1:8760, 3:(3+N_ZONES-1)]);

#Hydro
FUCR_HydroGs_head = FUCR_HydroGs[2];
FUCR_HydroGs = FUCR_HydroGs[1];
FUCR_HydroGs = Array{Int64}(FUCR_HydroGs[1:8760, 2:(2+N_ZONES-1)]);

SUCR_HydroGs_head = SUCR_HydroGs[2];
SUCR_HydroGs = SUCR_HydroGs[1];
SUCR_HydroGs = Array{Int64}(SUCR_HydroGs[1:8760, 2:(2+N_ZONES-1)]);

BUCR_HydroGs_head = BUCR_HydroGs[2];
BUCR_HydroGs = BUCR_HydroGs[1];
BUCR_HydroGs = Array{Int64}(BUCR_HydroGs[1:8760, 2:(2+N_ZONES-1)]);

#Nuclear
FUCR_NuclearGs_head = FUCR_NuclearGs[2];
FUCR_NuclearGs = FUCR_NuclearGs[1];
FUCR_NuclearGs = Array{Int64}(FUCR_NuclearGs[1:8760, 3:(3+N_ZONES-1)]);

SUCR_NuclearGs_head = SUCR_NuclearGs[2];
SUCR_NuclearGs = SUCR_NuclearGs[1];
SUCR_NuclearGs = Array{Int64}(SUCR_NuclearGs[1:8760, 3:(3+N_ZONES-1)]);

BUCR_NuclearGs_head = BUCR_NuclearGs[2];
BUCR_NuclearGs = BUCR_NuclearGs[1];
BUCR_NuclearGs = Array{Int64}(BUCR_NuclearGs[1:8760, 3:(3+N_ZONES-1)]);

#Cogen
FUCR_CogenGs_head = FUCR_CogenGs[2];
FUCR_CogenGs = FUCR_CogenGs[1];
FUCR_CogenGs = Array{Int64}(FUCR_CogenGs[1:8760, 3:(3+N_ZONES-1)]);

SUCR_CogenGs_head = SUCR_CogenGs[2];
SUCR_CogenGs = SUCR_CogenGs[1];
SUCR_CogenGs = Array{Int64}(SUCR_CogenGs[1:8760, 3:(3+N_ZONES-1)]);

BUCR_CogenGs_head = BUCR_CogenGs[2];
BUCR_CogenGs = BUCR_CogenGs[1];
BUCR_CogenGs = Array{Int64}(BUCR_CogenGs[1:8760, 3:(3+N_ZONES-1)]);

Map_Gens_head = Map_Gens[2];
Map_Gens = Map_Gens[1];
Map_Gens = Array{Int64}(Map_Gens[:,3:N_ZONES+2]);

Map_Peakers_head = Map_Peakers[2];
Map_Peakers = Map_Peakers[1];
Map_Peakers = Array{Int64}(Map_Peakers[:,3:N_ZONES+2]);

Map_Storage_head = Map_Storage[2];
Map_Storage = Map_Storage[1];
Map_Storage = Array{Int64}(Map_Storage[:,2:N_ZONES+1])

TranC_head = TranC[2];
TranC = TranC[1];
TranC = Array{Float64}(TranC[1:N_ZONES,2:(2+N_ZONES-1)]);

TranS_head = TranS[2];
TranS = TranS[1];
TranS = Array{Float64}(TranS[1:N_ZONES,2:(2+N_ZONES-1)]);

Reserve_Reqs_head = Reserve_Reqs[2];
Reserve_Reqs = Reserve_Reqs[1];
Reserve_Req_Up = Array{Float64}(Reserve_Reqs)


export MolOrbital, getMolOrbitals, Molecule, nnRepulsions

using LinearAlgebra: norm

"""

    MolOrbital{N} <: AbstractMolOrbital

`Struct` of molecular orbitals.

≡≡≡ Field(s) ≡≡≡

`symmetry::String`: The symmetry of the orbital. The default value is "A" for being
antisymmetric.

`energy::Float64`: Molecular energy.

`spin::String`: Spin function of the orbital. Available values: "Alpha", "Beta".

`occupancy::Real`: Occupation number.

`orbitalCoeffs::NTuple{N, Float64}`: coefficients of the basis functions to form the
molecular orbital.

≡≡≡ Initialization Method(s) ≡≡≡

    MolOrbital(energy::Float64, occupancy::Real, orbitalCoeffs::Array{Float64, 1},
               spin::String="Alpha", symmetry::String="A") -> MolOrbital{N}

"""
struct MolOrbital{N} <: AbstractMolOrbital
    symmetry::String
    energy::Float64
    spin::String
    occupancy::Real
    orbitalCoeffs::NTuple{N, Float64}

    function MolOrbital(energy::Float64, occupancy::Real, orbitalCoeffs::Vector{Float64},
                        spin::String="Alpha", symmetry::String="A")
        spin != "Alpha" && spin != "Beta" && error("Keyword argument \"spin\" can only"*
                                                   " be \"Alpha\" or \"Beta\"")
        new{length(orbitalCoeffs)}(symmetry, energy, spin, occupancy, orbitalCoeffs|>Tuple)
    end
end

"""

    getMolOrbitals(ens::Array{Float64, 1}, occus::Array{<:Real, 1}, C::Matrix{Float64},
                   spins::Array{String, 1},
                   symms::Array{String, 1}=repeat(["A"], length(occus))) ->
    Tuple{Vararg{getMolOrbitals}}

A function that returns the molecular orbitals.
"""
function getMolOrbitals(ens::Vector{Float64}, occus::Vector{<:Real}, C::Matrix{Float64},
                        spins::Vector{String},
                        symms::Vector{String}=repeat(["A"], length(occus)))
    MolOrbital.(ens, occus, [ C[:,i] for i = 1:size(C, 2) ], spins, symms) |> Tuple
end

"""

    Molecule{Nc, Ne, Nb} <:MolecularHartreeFockCoefficient{Nc, Ne}

Container for the information of a molecule.

≡≡≡ Field(s) ≡≡≡

`nuc::Tuple{Vararg{String}}`: Nuclei of the molecule.

`nucCoords::Tuple{Vararg{NTuple{3, Real}}}`: The coordinates of the nuclei.

`Ne::Int`: The number of electrons.

`orbital::Tuple{Vararg{MolOrbital}}`: Molecular orbitals.

`basis::Tuple{Vararg{FloatingGTBasisFuncs}}`: The basis set for the molecular orbitals.

`E0HF::Float64`: Hartree-Fock energy of the electronic Hamiltonian from the basis set.

`EnnR::Float64`: The nuclear-nuclear repulsion energy.

≡≡≡ Initialization Method(s) ≡≡≡

    Molecule(basis::Array{FloatingGTBasisFuncs, 1}, nuc::Array{String, 1},
             nucCoords::Array{<:AbstractArray, 1}, Ne::Int, E0HF::Float64,
             Emos::Array{Float64, 1}, occus::Array{<:Real, 1}, C::Array{Float64, 2},
             spins::Array{String, 1},
             symms::Array{String, 1}=repeat(["A"], length(occus))) ->
    Molecule{<:Any, Ne, <:Any}

`Emos` are the energies of corresponding molecular energies. `occus` are the occupation
numbers of the orbitals. `C` is the coefficient matrix, which does not need to be a square
matrix since the number of rows is the size of the (spatial) basis set whereas the number
of the columns represents the number of molecular orbitals. `spin` specifies the spin
functions of the orbitals, entries of which can be set to "Alpha" or "Beta". `symms` are
symmetries of the orbitals where the default entry value is "A" for being antisymmetric.

    Molecule(basis::Array{<:FloatingGTBasisFuncs, 1}, nuc::Array{String, 1},
             nucCoords::Array{<:AbstractArray, 1}, HFfVars::HFfinalVars) ->
    Molecule

Construct a `Molecule` from a basis set, nuclei information, and the result from the
corresponding Hartree-Fock SCF procedure, specifically a `HFfinalVars` `struct`.
"""
struct Molecule{Nc, Ne, Nb} <:MolecularHartreeFockCoefficient{Nc, Ne}
    nuc::Tuple{Vararg{String}}
    nucCoords::Tuple{Vararg{NTuple{3, Real}}}
    Ne::Int
    orbital::Tuple{Vararg{MolOrbital}}
    basis::Tuple{Vararg{FloatingGTBasisFuncs}}
    E0HF::Float64
    EnnR::Float64

    function Molecule(basis::Vector{<:FloatingGTBasisFuncs}, nuc::Vector{String},
                      nucCoords::Vector{<:AbstractArray}, Ne::Int, E0HF::Float64,
                      Emos::Vector{Float64}, occus::Vector{<:Real}, C::Matrix{Float64},
                      spins::Vector{String},
                      symms::Vector{String}=repeat(["A"], length(occus)))
        @assert length(nuc) == length(nucCoords)
        Nb = basisSize(basis) |> sum
        coeff = spins |> unique |> length
        @assert (coeff*Nb .==
                 length.([Emos, occus, C[1,:], spins, symms]) .==
                 coeff*length(C[:,1])) |> prod
        EnnR = nnRepulsions(nuc, nucCoords)
        iSorted = mapPermute(basis, sortBasisFuncs)
        basis = basis[iSorted]
        offset = pushfirst!(accumulate(+, collect(basisSize(basis)) .- 1), 0)
        iSortedExtended = Union{Int, Vector{Int}}[]
        append!(iSortedExtended, iSorted)
        for (i, idx) in zip(iSorted, 1:length(iSorted))
            if basis[i] isa BasisFuncs
                iSortedExtended[idx] = collect(i : i+length(basis[i])-1) .+ offset[i]
            else
                iSortedExtended[idx] += offset[i]
            end
        end
        C = C[vcat(iSortedExtended...), :]
        new{getCharge(nuc), Ne, Nb}(nuc |> Tuple,
                                    nucCoords .|> Tuple |> Tuple,
                                    Ne,
                                    getMolOrbitals(Emos, occus, C, spins, symms),
                                    deepcopy(basis) |> Tuple,
                                    E0HF,
                                    EnnR)
    end
end

function Molecule(basis::Vector{<:FloatingGTBasisFuncs}, nuc::Vector{String},
                  nucCoords::Vector{<:AbstractArray}, HFfVars::HFfinalVars)
    t = typeof(HFfVars)
    len = t.parameters[3]
    if t <: HFfinalVars{:UHF}
        Emos = HFfVars.Emo |> flatten
        occus = HFfVars.occu |> flatten
        C = hcat(HFfVars.C...)
        spins = vcat(fill("Alpha", len), fill("Beta", len))
    else
        Emos = HFfVars.Emo
        occus = HFfVars.occu
        C = HFfVars.C
        spins = fill("Alpha", len)
    end
    Molecule(basis, nuc, nucCoords, t.parameters[2], HFfVars.E0HF, Emos|>collect,
             occus|>collect, C, spins, fill("A", length(spins)))
end

"""

    nnRepulsions(nuc::Array{String, 1}, nucCoords::Array{<:AbstractArray, 1}) -> Float64

Calculate the nuclear-nuclear repulsion energy given the nuclei and their coordinates of a
molecule.
"""
function nnRepulsions(nuc::Vector{String}, nucCoords::Vector{<:AbstractArray})
    E = 0.0
    len = length(nuc)
    for i = 1:len, j=i+1:len
        E += getCharge(nuc[i]) * getCharge(nuc[j]) / norm(nucCoords[i] - nucCoords[j])
    end
    E
end


export move

function move(x::Array{Float64,2}; by=[0.0, 0.0, 0.0])
    x .+= vec(by)
    return x
end

function move(obj::T; by=[0.0, 0.0, 0.0]) where T <: Union{Shape, PostOpObj}
    # x, v, y, vol, type
    function func(out)
        x = out[:x]
        x = move(x, by=by)
        return repack!(out, [:x], [x])
    end
    if isa(obj, Shape)
        return PostOpObj(obj, func)
    elseif isa(obj, PostOpObj)
        push!(obj.operations, func)
        return obj
    else
        error("Not allowed.")
    end
end


# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "Elemental"
version = v"0.87.7"

# Collection of sources required to build Elemental
sources = [
    GitSource("https://github.com/elemental/Elemental.git",
              "477e503a7a840cc1a75173552711b980505a0b06"),
]

# Bash recipe for building across all platforms
script = raw"""
if [[ "$nbits" == "64" ]]; then
  INT64="ON"
else
  INT64="OFF"
fi

if [[ "$nbits" == "64" ]] && [[ "$target" != aarch64-* ]]; then
  BLAS_LAPACK_LIB="$libdir/libopenblas64_.$dlext"
  BLAS_LAPACK_SUFFIX="_64_"
  BLAS_INT64="ON"
else
  BLAS_LAPACK_LIB="$libdir/libopenblas.$dlext"
  BLAS_LAPACK_SUFFIX=""
  BLAS_INT64="OFF"
fi

mkdir "$WORKSPACE/srcdir/$SRC_NAME/build"
cd "$WORKSPACE/srcdir/$SRC_NAME/build"

cmake \
  -DCMAKE_INSTALL_PREFIX="$prefix" \
  -DCMAKE_TOOLCHAIN_FILE="$CMAKE_TARGET_TOOLCHAIN" \
  -DCMAKE_BUILD_TYPE="Release" \
  -DEL_USE_64BIT_INTS="$INT64" \
  -DEL_USE_64BIT_BLAS_INTS="$BLAS_INT64" \
  -DEL_DISABLE_PARMETIS="ON" \
  -DMETIS_TEST_RUNS_EXITCODE="0" \
  -DMETIS_TEST_RUNS_EXITCODE__TRYRUN_OUTPUT="" \
  -DMATH_LIBS="$BLAS_LAPACK_LIB" \
  -DBLAS_LIBRARIES="$BLAS_LAPACK_LIB" \
  -DLAPACK_LIBRARIES="$BLAS_LAPACK_LIB" \
  -DEL_BLAS_SUFFIX="$BLAS_LAPACK_SUFFIX" \
  -DEL_LAPACK_SUFFIX="$BLAS_LAPACK_SUFFIX" \
  "$WORKSPACE/srcdir/$SRC_NAME"

make "-j$nproc"
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()
platforms = expand_cxxstring_abis(platforms)
filter!(!Sys.iswindows, platforms)

# The products that we will ensure are always built
products = [
    LibraryProduct("libEl", :libEl),
    LibraryProduct("libElSuiteSparse", :libElSuiteSparse),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("METIS_jll"),
    Dependency("MPICH_jll"),
    Dependency("OpenBLAS_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


# This file is a part of ValueShapes.jl, licensed under the MIT License (MIT).

__precompile__(true)

"""
    ValueShapes

Provides a Julia API to describe the shape of values, like scalars, arrays
and structures.
"""
module ValueShapes

using Base: @propagate_inbounds

using ArgCheck
using ArraysOfArrays
using ChangesOfVariables
using Distributions
using ElasticArrays
using FillArrays
using InverseFunctions
using Random
using Statistics
using StatsBase
using Static

using DensityInterface

import MeasureBase
using MeasureBase: AbstractMeasure, StdMeasure, basemeasure, productmeasure, getdof
using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin
using MeasureBase: TransportFunction, PowerMeasure

import DistributionMeasures
using DistributionMeasures: DistributionMeasure, StandardDist

import ChainRulesCore
import IntervalSets
import Tables
import TypedTables

# Long-term, ChainRulesCore should be sufficient:
import ZygoteRules

using ChainRulesCore: AbstractTangent, Tangent, AbstractZero, NoTangent, ZeroTangent
using ChainRulesCore: AbstractThunk, ProjectTo, unthunk, backing

using Random123: Philox4x

include("tangent_utils.jl")
include("value_shape.jl")
include("value_accessor.jl")
include("scalar_shape.jl")
include("array_shape.jl")
include("const_value_shape.jl")
include("named_tuple_shape.jl")
include("functions.jl")
include("distributions.jl")
include("const_value_dist.jl")
include("named_tuple_dist.jl")
include("reshaped_dist.jl")

end # module


module GraphMLDatasets
    using InteractiveUtils: subtypes
    using SparseArrays: SparseMatrixCSC, sparse, findnz

    using CSV
    using DataDeps: DataDep, register, @datadep_str
    using HTTP
    using JLD2
    using JSON
    using LightGraphs: SimpleGraph, SimpleDiGraph, add_edge!, nv, ne
    using MAT
    using PyCall

    export
        Dataset,
        Planetoid,
        Cora,
        PPI,
        Reddit,
        QM7b,
        Entities,
        traindata,
        validdata,
        testdata,
        graphdata,
        rawdata,
        alldata,
        metadata

    include("./dataset.jl")
    include("./preprocess.jl")
    include("./interfaces.jl")
    include("./utils.jl")
    include("./register.jl")

    function __init__()
        init_dataset(Dataset)
    end
end


# This file was generated, do not modify it. # hide
LinMod = @pipeline(FeatureSelector(features=[:hp1]),
                   LinearRegressor());


module TestDay11

import AOC2021.Day11.part1, AOC2021.Day11.part2
using Test

test_input = """
5483143223
2745854711
5264556173
6141336146
6357385478
4167524645
2176841721
6882881134
4846848554
5283751526"""

@testset "Day 11" begin
    @testset "part 1" begin
        @test part1(test_input) == 1656
    end
    @testset "part 2" begin
        @test part2(test_input) == 195
    end
end

end # module


struct DefaultInit <: DiffEqBase.DAEInitializationAlgorithm end
struct NoInit <: DiffEqBase.DAEInitializationAlgorithm end

struct ShampineCollocationInit{T} <: DiffEqBase.DAEInitializationAlgorithm
	initdt::T
end
ShampineCollocationInit() = ShampineCollocationInit(nothing)

struct BrownFullBasicInit{T} <: DiffEqBase.DAEInitializationAlgorithm
	abstol::T
end
BrownFullBasicInit() = BrownFullBasicInit(1e-10)

## Notes

#=
differential_vars = [any(!iszero,x) for x in eachcol(M)]

A column should be zero for an algebraic variable, since that means that the
derivative term doesn't show up in any equations (i.e. is an algebraic variable).
The rows are not necessarily non-zero, for example a flux condition between two
differential variables. But if it's a condition that doesn't involve the algebraic
variable, then the system is not Index 1!

=#

## Expansion

function DiffEqBase.initialize_dae!(integrator::ODEIntegrator,
									initializealg = integrator.initializealg)
	_initialize_dae!(integrator, integrator.sol.prob,
					 initializealg,
					 Val(DiffEqBase.isinplace(integrator.sol.prob)))
end

## Default algorithms

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::DefaultInit, x::Val{true})
	_initialize_dae!(integrator, prob,
					BrownFullBasicInit(), x)
end

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::DefaultInit, x::Val{false})
	_initialize_dae!(integrator, prob,
					BrownFullBasicInit(), x)
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::DefaultInit, x::Val{false})
	if prob.differential_vars === nothing
		_initialize_dae!(integrator, prob,
						ShampineCollocationInit(), x)
	else
		_initialize_dae!(integrator, prob,
						BrownFullBasicInit(), x)
	end
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::DefaultInit, x::Val{true})
	if prob.differential_vars === nothing
		_initialize_dae!(integrator, prob,
						ShampineCollocationInit(), x)
	else
		_initialize_dae!(integrator, prob,
						BrownFullBasicInit(), x)
	end
end

## NoInit

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::NoInit, x::Val{true})
end

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::NoInit, x::Val{false})
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::NoInit, x::Val{false})
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::NoInit, x::Val{true})
end

## ShampineCollocationInit

#=
The method:

du = (u-u0)/h
Solve for `u`

=#

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::ShampineCollocationInit, ::Val{true})

	@unpack p, t, f = integrator
 	M = integrator.f.mass_matrix
	dtmax = integrator.opts.dtmax
	tmp = first(get_tmp_cache(integrator))
	u0 = integrator.u

	dt = t != 0 ? min(t/1000,dtmax) : dtmax # Haven't implemented norm reduction

 	nlequation! = function (out,u)
		update_coefficients!(M,u,p,t)
 		#M * (u-u0)/dt - f(u,p,t)
		@. tmp = (u - u0)/dt
		mul!(out,M,tmp)
		f(tmp,u,p,t)
		out .-= tmp
		nothing
 	end

	update_coefficients!(M,u0,p,t)
	algebraic_vars = [all(iszero,x) for x in eachcol(M)]
	algebraic_eqs  = [all(iszero,x) for x in eachrow(M)]
  (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return
	f(tmp,u0,p,t)
	tmp .= algebraic_eqs .* tmp

 	integrator.opts.internalnorm(tmp,t) <= integrator.opts.abstol && return

	integrator.u .= nlsolve(nlequation!, u0).zero
	recursivecopy!(integrator.uprev,integrator.u)
	if alg_extrapolates(integrator.alg)
		recursivecopy!(integrator.uprev2,integrator.uprev)
	end
  return nothing
end

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::ShampineCollocationInit, ::Val{false})

	@unpack p, t, f = integrator
	u0 = integrator.u
	M = integrator.f.mass_matrix
	dtmax = integrator.opts.dtmax

	dt = t != 0 ? min(t/1000,dtmax/10) : dtmax # Haven't implemented norm reduction

	update_coefficients!(M,u0,p,t)
	algebraic_vars = [all(iszero,x) for x in eachcol(M)]
	algebraic_eqs  = [all(iszero,x) for x in eachrow(M)]
  (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return
	du = f(u0,p,t)
	resid = du[algebraic_eqs]

  integrator.opts.internalnorm(resid,t) <= integrator.opts.abstol && return

	nlequation_oop = function (u)
		update_coefficients!(M,u,p,t)
		M * (u-u0)/dt - f(u,p,t)
	end

	nlequation! = (out,u) -> out .= nlequation_oop(u)

	integrator.u = nlsolve(nlequation!, u0).zero
	integrator.uprev = integrator.u
	if alg_extrapolates(integrator.alg)
		integrator.uprev2 = integrator.uprev
	end
  return
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::ShampineCollocationInit, ::Val{true})

	@unpack p, t, f = integrator
	u0 = integrator.u

	dtmax = integrator.opts.dtmax
	tmp = get_tmp_cache(integrator)[1]
	resid = get_tmp_cache(integrator)[2]

	dt = t != 0 ? min(t/1000,dtmax) : dtmax # Haven't implemented norm reduction

 	nlequation! = function (out,u)
 		#M * (u-u0)/dt - f(u,p,t)
		@. tmp = (u - u0)/dt
		f(out,tmp,u,p,t)
		nothing
 	end

	nlequation!(tmp,u0)
	f(resid,integrator.du,u0,p,t)
 	integrator.opts.internalnorm(resid,t) <= integrator.opts.abstol && return

	integrator.u .= nlsolve(nlequation!, u0).zero
	recursivecopy!(integrator.uprev,integrator.u)
	if alg_extrapolates(integrator.alg)
		recursivecopy!(integrator.uprev2,integrator.uprev)
	end
  return
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::ShampineCollocationInit, ::Val{false})

	@unpack p, t, f = integrator
	u0 = integrator.u
	dtmax = integrator.opts.dtmax

	dt = t != 0 ? min(t/1000,dtmax/10) : dtmax # Haven't implemented norm reduction

	nlequation_oop = function (u)
		f((u-u0)/dt,u,p,t)
	end

	nlequation! = (out,u) -> out .= nlequation_oop(u)

	resid = f(integrator.du,u0,p,t)
	integrator.opts.internalnorm(resid,t) <= integrator.opts.abstol && return

	integrator.u = nlsolve(nlequation!, u0).zero
	integrator.uprev = integrator.u
	if alg_extrapolates(integrator.alg)
		integrator.uprev2 = integrator.uprev
	end
end

## BrownFullBasic

#=
The method:

Keep differential variables constant
Solve for the algebraic variables

=#

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::BrownFullBasicInit, ::Val{true})
	@unpack p, t, f = integrator
	u = integrator.u
	M = integrator.f.mass_matrix
	update_coefficients!(M,u,p,t)
	algebraic_vars = [all(iszero,x) for x in eachcol(M)]
	algebraic_eqs  = [all(iszero,x) for x in eachrow(M)]
  (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return
	tmp = get_tmp_cache(integrator)[1]

	f(tmp,u,p,t)

	tmp .= algebraic_eqs .* tmp

	integrator.opts.internalnorm(tmp,t) <= alg.abstol && return
	alg_u = @view u[algebraic_vars]

	nlequation = (out, x) -> begin
		alg_u .= x
		f(tmp, u, p, t)
		out .= @view tmp[algebraic_eqs]
	end

	r = nlsolve(nlequation, u[algebraic_vars])
	alg_u .= r.zero

	recursivecopy!(integrator.uprev,integrator.u)
	if alg_extrapolates(integrator.alg)
		recursivecopy!(integrator.uprev2,integrator.uprev)
	end

	return
end

function _initialize_dae!(integrator, prob::ODEProblem,
						 alg::BrownFullBasicInit, ::Val{false})
	@unpack p, t, f = integrator

	u0 = integrator.u
	M = integrator.f.mass_matrix
	update_coefficients!(M,u0,p,t)
	algebraic_vars = [all(iszero,x) for x in eachcol(M)]
	algebraic_eqs  = [all(iszero,x) for x in eachrow(M)]
  (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return

	du = f(u0,p,t)
	resid = du[algebraic_eqs]

  integrator.opts.internalnorm(resid,t) <= alg.abstol && return

	if u0 isa Number
		# This doesn't fix static arrays!
		u = [u0]
	else
		u = u0
	end

	alg_u = @view u[algebraic_vars]

	nlequation = (out,x) -> begin
		alg_u .= x
		du = f(u,p,t)
		out .= @view du[algebraic_eqs]
	end
	r = nlsolve(nlequation, u0[algebraic_vars])
	alg_u .= r.zero

	if u0 isa Number
		# This doesn't fix static arrays!
		integrator.u = first(u)
	else
		integrator.u = u
	end

	integrator.uprev = integrator.u
	if alg_extrapolates(integrator.alg)
		integrator.uprev2 = integrator.uprev
	end

	return
end

function _initialize_dae!(integrator, prob::DAEProblem,
						 alg::BrownFullBasicInit, ::Val{true})
	@unpack p, t, f = integrator
	differential_vars = prob.differential_vars
	u = integrator.u
	du = integrator.du

	tmp = get_tmp_cache(integrator)[1]
	f(tmp, du, u, p, t)

	if integrator.opts.internalnorm(tmp,t) <= alg.abstol
		return
	elseif differential_vars === nothing
		error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.")
	end

	nlequation = (out, x) -> begin
		@. du = ifelse(differential_vars,x,du)
		@. u  = ifelse(differential_vars,u,x)
		f(out, du, u, p, t)
	end

	r = nlsolve(nlequation, ifelse.(differential_vars,du,u))

	@. du = ifelse(differential_vars,r.zero,du)
	@. u  = ifelse(differential_vars,u,r.zero)

	recursivecopy!(integrator.uprev,integrator.u)
	if alg_extrapolates(integrator.alg)
		recursivecopy!(integrator.uprev2,integrator.uprev)
	end

	return
end

function _initialize_dae!(integrator, prob::DAEProblem,
						  alg::BrownFullBasicInit, ::Val{false})
	@unpack p, t, f = integrator
	differential_vars = prob.differential_vars

	if integrator.opts.internalnorm(f(integrator.du, integrator.u, p, t),t) <= alg.abstol
		return
	elseif differential_vars === nothing
		error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.")
	end

	if integrator.u isa Number && integrator.du isa Number
		# This doesn't fix static arrays!
		u = [integrator.u]
		du = [integrator.du]
	else
		u = integrator.u
		du = integrator.du
	end

	nlequation = (out,x) -> begin
		@. du = ifelse(differential_vars,x,du)
		@. u  = ifelse(differential_vars,u,x)
		out .= f(du, u, p, t)
	end

	r = nlsolve(nlequation, ifelse.(differential_vars,du,u))

	@. du = ifelse(differential_vars,r.zero,du)
	@. u  = ifelse(differential_vars,u,r.zero)

	if integrator.u isa Number && integrator.du isa Number
		# This doesn't fix static arrays!
		integrator.u = first(u)
		integrator.du = first(du)
	else
		integrator.u = u
		integrator.du = du
	end

	integrator.uprev = integrator.u
	if alg_extrapolates(integrator.alg)
		integrator.uprev2 = integrator.uprev
	end

	return
end


macro plot(plottable_table, plot_call = :(plot()), err_style = :(:yerr))
    if @capture(plot_call, fun_(as__; kws__))
        return esc(Expr(:call, :(GroupedErrors.plot_helper), plottable_table, fun, err_style, as..., kws...))
    elseif @capture(plot_call, fun_(as__))
        return esc(Expr(:call, :(GroupedErrors.plot_helper), plottable_table, fun, err_style, as...))
    else
        return esc(plottable_table)
    end
end

function plot_helper(t::ProcessedTable, f::Function, err_style, args...; kwargs...)
    s = t.table
    axis_type = t.kw[:axis_type]
    group = Tuple(columns(t.table, s) for s in listsplits(t))
    group_list = unique(Columns(group))
    plot_kwargs = [(plot_kwarg[1], reshape(plot_kwarg[2].(group_list), 1, :))
        for plot_kwarg in t.kw[:plot_kwargs]]
    if axis_type == :pointbypoint
        err_kwargs = []
        if eltype(columns(s, :x))<:Tuple
            x = getindex.(columns(s, :x), 1)
            push!(err_kwargs, (:xerr, getindex.(columns(s, :x), 2)))
        else
            x = columns(s, :x)
        end
        if eltype(columns(s, :y))<:Tuple
            y = getindex.(columns(s, :y), 1)
            push!(err_kwargs, (:yerr, getindex.(columns(s, :y), 2)))
        else
            y = columns(s, :y)
        end
        f(args..., x, y; group = group, err_kwargs...,kwargs..., plot_kwargs...)
    else
        x, y = columns(s, :x), columns(s, :y)
        err = :err in colnames(s) ? columns(s, :err) : nothing
        (axis_type == :continuous) && (err_style = :ribbon)
        kw_itr = Iterators.flatten([[(err_style, err)], kwargs, plot_kwargs])
        f(args..., x, y; group = group, kw_itr...)
    end
end

plot_helper(t, f::Function, err_style, args...; kwargs...) =
    plot_helper(ProcessedTable(t), f, err_style, args...; kwargs...)


# this script generates the SnoopCompile precompile list

using SnoopCompile

### Log the compiles
# This only needs to be run once (to generate "/tmp/approxfun_compiles.csv")

SnoopCompile.@snoop "/tmp/approxfun_compiles.csv" begin
    using ApproxFun
    a=Fun()
    b=Fun(exp)
    a+b
    a*b
    a/b
    a-b

    # following causes SEGFAULT
    # S = Chebyshev()
    # D = Derivative(S)
    # B = dirichlet(S)
    # A = [B;D^2+I]
    # u = A\[1.0,0.0,b]
end

### Parse the compiles and generate precompilation scripts
# This can be run repeatedly to tweak the scripts

# IMPORTANT: we must have the module(s) defined for the parcelation
# step, otherwise we will get no precompiles for the ApproxFun module
using ApproxFun

data = SnoopCompile.read("/tmp/approxfun_compiles.csv")

# Blacklist helps fix problems:
# - MIME uses type-parameters with symbols like :image/png, which is
#   not parseable
blacklist = ["MIME"]

# Use these two lines if you want to create precompile functions for
# individual packages
pc, discards = SnoopCompile.parcel(data[end:-1:1,2], blacklist=blacklist)
SnoopCompile.write("/tmp/precompile", pc)


using FITSIO, DataFrames, FITSIO

function sxtevtlist2df(hdu::TableHDU)
#    f=FITS(evtfile)

    ti=read(hdu,"TIME")
    ccdframe = read(hdu, "CCDFrame")
        x= read(hdu, "X")
        y= read(hdu, "Y")
        rawx= read(hdu, "RAWX")
        rawy= read(hdu, "RAWY")
        detx= read(hdu, "DETX")
        dety= read(hdu, "DETY")
        phas= read(hdu, "PHAS")
        stts= read(hdu, "STATUS")
        pha= read(hdu, "PHA")
        grade= read(hdu, "GRADE")
        pixsabove= read(hdu, "PixsAbove")
        pi_channel= read(hdu, "PI")
        ra= read(hdu, "RA")
        decl= read(hdu, "Dec")
        phaso= read(hdu, "PHASO")
        evtdf=DataFrame(TIME=ti,CCDFrame=ccdframe, X=x,Y=y,RAWX=rawx,RAWY=rawy,DETX=detx,DETY=dety,PHAS1=phas[1,:],PHAS2=phas[2,:],PHAS3=phas[3,:],PHAS4=phas[4,:],PHAS5=phas[5,:],PHAS6=phas[6,:],PHAS7=phas[7,:],PHAS8=phas[8,:],PHAS9=phas[9,:],STATUS=stts,PHA=pha,GRADE=grade,PixsAbove=pixsabove,PI=pi_channel,RA=ra,Dec=decl,PHASO1=phaso[1,:],PHASO2=phaso[2,:],PHASO3=phaso[3,:],PHASO4=phaso[4,:],PHASO5=phaso[5,:],PHASO6=phaso[6,:],PHASO7=phaso[7,:],PHASO8=phaso[8,:],PHASO9=phaso[9,:])
    return evtdf
end

function sxtbadpixlist2df(hdu::TableHDU)
#    f=FITS(evtfile)
        rawx= read(hdu, "RAWX")
        rawy= read(hdu, "RAWY")
        typ= read(hdu, "TYPE")
        yextnt= read(hdu, "YEXTENT")
        badflag= read(hdu, "BADFLAG")
        badpixdf=DataFrame(RAWX=rawx,RAWY=rawy,TYPE=typ,YEXTENT=yextnt,BADFLAG=badflag)
    return badpixdf
end


# import Base: *
# import Knet: getindex, setindex!

# Matmuls 2d and 3d arrays
# function *(a::AbstractArray{T,2}, b::AbstractArray{T,3}) where T<:Real
#     b_sizes = size(b)
#     a = a * reshape(b, b_sizes[1], :)
#     return reshape(a, :, b_sizes[2:end]...)
# end

# Matmuls 2d and 3d arrays for KnetArrays
# function *(a::KnetArray{T,2}, b::KnetArray{T,3}) where T<:Real
#     b_sizes = size(b)
#     a = a * reshape(b, b_sizes[1], :)
#     return reshape(a, :, b_sizes[2:end]...)
# end

# TODO :
# Since backprop doesn't work with this new import, we define it as a complex function consisting primitives that Autograd can take derivatives of. A primitive derivative of this function would speed things up.
# function matmul23(a::KnetArray{T,2}, b::KnetArray{T,3}) where T<:Real
#     b_sizes = size(b)
#     a = a * reshape(b, b_sizes[1], :)
#     return reshape(a, :, b_sizes[2:end]...)
# end

gelu(x) = x .* 0.5 .* (1.0 .+ erf.(x ./ sqrt(2.0)))

function matmul23(a, b)
    b_sizes = size(b)
    a = a * reshape(b, b_sizes[1], :)
    return reshape(a, :, b_sizes[2:end]...)
end

# Wrote these first, then realized we don't need them. Might come in handy later.
# function matmul23(a::AbstractArray{T,2}, b::AbstractArray{T,3}) where T<:Real
#     b_sizes = size(b)
#     a = a * reshape(b, b_sizes[1], :)
#     return reshape(a, :, b_sizes[2:end]...)
# end

# function matmul23(a::Param{KnetArray{T,2}}, b::KnetArray{T,3}) where T<:Real
#     matmul23(value(a), b)
# end

# function matmul23(a::Param{AbstractArray{T,2}}, b::AbstractArray{T,3}) where T<:Real
#     matmul23(value(a), b)
# end

# function matmul23(a::Param{KnetArray{T,2}}, b::AutoGrad.Result{KnetArray{T,3}}) where T<:Real
#     matmul23(value(a), value(b))
# end

# function matmul23(a::Param{AbstractArray{T,2}}, b::AutoGrad.Result{AbstractArray{T,3}}) where T<:Real
#     matmul23(value(a), value(b))
# end

# @primitive *(x1::KnetArray{T,2},x2::KnetArray{T,3}),dy

# Not using this anymore
# function getindex(A::KnetArray{Float32,3}, ::Colon, I::Real, ::Colon)
#     sizes = size(A)
#     A = reshape(A, :, sizes[3])
#     return A[(I-1)*sizes[1]+1:I*sizes[1],:]
# #    reshape(A, :, size(A,3))[(I-1)*size(A,1)+1:I*size(A,1),:]
# end

# Does not work
# function setindex!(A::KnetArray{Float32,3}, v, ::Colon, I::Real, ::Colon)
#     A = reshape(A, :, size(A,3))
# #     setindex!(A, v, (I-1)*size(A,1)+1:I*size(A,1), ::Colon)
#     A[(I-1)*size(A,1)+1:I*size(A,1),:] = v
# end

# std doesn't work!
std2(a, μ, ϵ) = sqrt.(Knet.mean(abs2.(a .- μ), dims=1) .+ ϵ)

# Legend
# V -> Vocab size, E -> Embedding size, S -> Sequence length, B -> Batch size
# H -> head_size, N -> num_heads

abstract type Layer end

# MAYBE TODO : sin-cos positionwise embeddings. This will reduce model size by max_seq_len * E

mutable struct Embedding <: Layer
    w
end

Embedding(vocabsize::Int,embed::Int; atype=Array{Float32}) = Embedding(param(embed,vocabsize, atype=atype))

function (e::Embedding)(x)
    e.w[:,x]
end

# If we need 0's as pads
#=
struct SegmentEmbedding <: Layer
    w
    atype
end

SegmentEmbedding(vocabsize::Int,embed::Int; atype=Array{Float32}) = SegmentEmbedding(param(embed,vocabsize, atype=atype), atype)

function (e::SegmentEmbedding)(x)
    x != 0 ? e.w[:,x] : e.atype(zeros(size(e.w,1)))
end
=#

mutable struct Linear <: Layer
    w
    b
end

Linear(input_size::Int, output_size::Int; atype=Array{Float32}) = Linear(param(output_size, input_size, atype=atype), param0(output_size, atype=atype))

function (l::Linear)(x)
    return l.w * x .+ l.b
end

mutable struct Linear3D <: Layer
    w
    b
end

Linear3D(input_size::Int, output_size::Int; atype=Array{Float32}) = Linear3D(param(output_size, input_size, atype=atype), param0(output_size, atype=atype))

function (l::Linear3D)(x)
    return matmul23(l.w, x) .+ l.b
end

# Absolutely no difference between Dense and Linear! Except one has dropout and activation function.
mutable struct Dense <: Layer
    linear
    pdrop
    func
end

function Dense(input_size::Int, output_size::Int; pdrop=0.0, func=identity, atype=Array{Float32}, threeD=false)
    if threeD
        return Dense(Linear3D(input_size, output_size, atype=atype), pdrop, func)
    else
        return Dense(Linear(input_size, output_size, atype=atype), pdrop, func)
    end
end

function (a::Dense)(x)
    return a.func.(dropout(a.linear(x), a.pdrop))
end

mutable struct LayerNormalization <: Layer
    γ
    β
    ϵ
end

LayerNormalization(hidden_size::Int; epsilon=1e-12, atype=Array{Float32}) = LayerNormalization(Param(atype(ones(hidden_size))), param0(hidden_size, atype=atype), epsilon)

function (n::LayerNormalization)(x)
    μ = Knet.mean(x, dims=1)
    x = (x .- μ) ./ std2(x, μ, n.ϵ) # corrected=false for n
    return n.γ .* x .+ n.β
end

mutable struct EmbedLayer <: Layer
    wordpiece::Embedding
    positional::Embedding
#    segment::SegmentEmbedding
    segment::Embedding
    layer_norm::LayerNormalization
    seq_len::Int
    pdrop
end

function EmbedLayer(config)
    wordpiece = Embedding(config.vocab_size, config.embed_size, atype=config.atype)
    positional = Embedding(config.max_seq_len, config.embed_size, atype=config.atype)
    #segment = SegmentEmbedding(config.num_segment, config.embed_size, atype=config.atype)
    segment = Embedding(config.num_segment, config.embed_size, atype=config.atype)
    layer_norm = LayerNormalization(config.embed_size, atype=config.atype)
    return EmbedLayer(wordpiece, positional, segment, layer_norm, config.seq_len, config.pdrop)
end

function (e::EmbedLayer)(x, segment_ids) # segment_ids are SxB, containing 1 or 2, or 0 in case of pads.
    x = e.wordpiece(x)
    positions = zeros(Int64, e.seq_len, size(x,3)) .+ collect(1:e.seq_len) # size(x,3) is batchsize. Resulting matrix is SxB
    x = x .+ e.positional(positions)
    #x .+= reshape(hcat(e.segment.(segment_ids)...), (:, size(segment_ids,1),size(segment_ids,2)))
    x = x .+ e.segment(segment_ids)
    x = e.layer_norm(x)
    return dropout(x, e.pdrop)
end

function divide_to_heads(x, num_heads, head_size, seq_len)
    x = reshape(x, (head_size, num_heads, seq_len, :))
    x = permutedims(x, (1,3,2,4))
    return reshape(x, (head_size, seq_len, :)) # Reshape to 3D so bmm can handle it.
end

mutable struct SelfAttention <: Layer
    query::Linear3D # N*H x E
    key::Linear3D
    value::Linear3D
    linear::Linear3D
    num_heads::Int
    seq_len::Int
    embed_size::Int
    head_size::Int
    head_size_sqrt::Int
    attention_pdrop
    pdrop
end

function SelfAttention(config)
    config.embed_size % config.num_heads != 0 && throw("Embed size should be divisible by number of heads!")
    head_size = Int(config.embed_size / config.num_heads)
    head_size_sqrt = Int(sqrt(head_size))
    head_size_sqrt * head_size_sqrt != head_size && throw("Square root of head size should be an integer!")
    query = Linear3D(config.embed_size, head_size*config.num_heads, atype=config.atype) # H*N is always equal to E
    key = Linear3D(config.embed_size, head_size*config.num_heads, atype=config.atype)
    value = Linear3D(config.embed_size, head_size*config.num_heads, atype=config.atype)
    linear = Linear3D(config.embed_size, config.embed_size, atype=config.atype)
    return SelfAttention(query, key, value, linear, config.num_heads, config.seq_len, config.embed_size, head_size, head_size_sqrt, config.attention_pdrop, config.pdrop)
end

function (s::SelfAttention)(x, attention_mask)
    # We make all the batchsize ones colon, in case of batches smaller than batchsize.
    # x is ExSxB
    query = divide_to_heads(s.query(x), s.num_heads, s.head_size, s.seq_len) # H x S x N*B
    key = divide_to_heads(s.key(x), s.num_heads, s.head_size, s.seq_len)
    value = divide_to_heads(s.value(x), s.num_heads, s.head_size, s.seq_len)

    # Scaled Dot Product Attention
    query = bmm(permutedims(key, (2,1,3)), query)
    query = query ./ s.head_size_sqrt # Scale down. I init this value to avoid taking sqrt every forward operation.
    # Masking. First reshape to 4d, then add mask, then reshape back to 3d.
    query = reshape(reshape(query, (s.seq_len, s.seq_len, s.num_heads, :)) .+ attention_mask, (s.seq_len, s.seq_len, :))

    query = Knet.softmax(query, dims=1)
    query = dropout(query, s.attention_pdrop)
    query = bmm(value, query)
    query = permutedims(reshape(query, (s.head_size, s.seq_len, s.num_heads, :)), (1,3,2,4))

    query = reshape(query, (s.embed_size, s.seq_len, :)) # Concat
    return dropout(s.linear(query), s.pdrop) # Linear transformation at the end
    # In pytorch version dropout is after layer_norm!
end

mutable struct FeedForward <: Layer
    dense::Dense
    linear::Linear3D
    pdrop
end

function FeedForward(config)
    dense = Dense(config.embed_size, config.ff_hidden_size, func=eval(Meta.parse(config.func)), atype=config.atype, threeD=true)
    linear = Linear3D(config.ff_hidden_size, config.embed_size, atype=config.atype)
    return FeedForward(dense, linear, config.pdrop)
end

function (f::FeedForward)(x)
    x = f.dense(x)
    return dropout(f.linear(x), f.pdrop)
end

mutable struct Encoder <: Layer
    self_attention::SelfAttention
    layer_norm1::LayerNormalization
    feed_forward::FeedForward
    layer_norm2::LayerNormalization
end

function Encoder(config)
    return Encoder(SelfAttention(config), LayerNormalization(config.embed_size, atype=config.atype), FeedForward(config), LayerNormalization(config.embed_size, atype=config.atype))
end

function (e::Encoder)(x, attention_mask)
    x = e.layer_norm1(x .+ e.self_attention(x, attention_mask))
    return e.layer_norm2(x .+ e.feed_forward(x))
end

mutable struct Bert <: Layer
    embed_layer::EmbedLayer
    encoder_stack
    atype
end

function Bert(config)
    embed_layer = EmbedLayer(config)
    encoder_stack = Encoder[]
    for _ in 1:config.num_encoder
        push!(encoder_stack, Encoder(config))
    end
    return Bert(embed_layer, encoder_stack, config.atype)
end

# x and segment_ids are SxB integers
function (b::Bert)(x, segment_ids; attention_mask=nothing)
    # Init attention_mask if it's not given
    attention_mask = attention_mask == nothing ? ones(size(x)) : attention_mask
    attention_mask = reshape(attention_mask, (size(attention_mask,1), 1, 1, size(attention_mask,2))) # Make it 4d
    attention_mask = (1 .- attention_mask) .* -10000.0 # If integer was 0, now it is masking. ones(size(attention_mask))
    attention_mask = b.atype(attention_mask)

    x = b.embed_layer(x, segment_ids)
    for encoder in b.encoder_stack
        x = encoder(x, attention_mask)
    end
    return x
end

mutable struct Pooler <: Layer
    linear::Linear
end

Pooler(embed_size::Int; atype=Array{Float32}) = Pooler(Linear(embed_size, embed_size, atype=atype))

function (p::Pooler)(x)
    # TODO :
    # Gave up on getindex function for 3D matrices because I could not figure out how to write setindex! for backprop
#     x = reshape(x, :, size(x,3))
#    return tanh.(p.linear(x[:,1,:])) # Use only CLS token. Returns ExB
    return tanh.(p.linear(reshape(x, :, size(x,3))[1:size(x,1),:]))
end

mutable struct NSPHead <: Layer
    linear::Linear
end

NSPHead(embed_size::Int; atype=Array{Float32}) = NSPHead(Linear(embed_size, 2, atype=atype))

(n::NSPHead)(x) = n.linear(x)

mutable struct MLMHead <: Layer
    dense::Dense
    layer_norm::LayerNormalization
    linear::Linear3D
end

function MLMHead(config, embedding_matrix)
    dense = Dense(config.embed_size, config.embed_size, func=eval(Meta.parse(config.func)), pdrop=0.0, atype=config.atype, threeD=true)
    layer_norm = LayerNormalization(config.embed_size, atype=config.atype)
    linear = Linear3D(config.embed_size, config.vocab_size, atype=config.atype)
    # TODO : Do this a shared weight
    #linear.w = permutedims(embedding_matrix, (2,1))
    return MLMHead(dense, layer_norm, linear)
end

function (m::MLMHead)(x)
    x = m.dense(x)
    x = m.layer_norm(x)
    return m.linear(x)
end

mutable struct BertPreTraining <: Layer
    bert::Bert
    pooler::Pooler
    nsp::NSPHead
    mlm::MLMHead
end

function BertPreTraining(config)
    bert = Bert(config)
    pooler = Pooler(config.embed_size, atype=config.atype)
    nsp = NSPHead(config.embed_size, atype=config.atype)
    mlm = MLMHead(config, bert.embed_layer.wordpiece.w) # TODO : Dont forget about embedding matrix
    return BertPreTraining(bert, pooler, nsp, mlm)
end

# We do not need a predictor, since this is only for pretraining
function (b::BertPreTraining)(x, segment_ids, mlm_labels, nsp_labels; attention_mask=nothing) # mlm_labels are SxB, so we just flatten them.
    x = b.bert(x, segment_ids, attention_mask=attention_mask)
    nsp_preds = b.nsp(b.pooler(x)) # 2xB
    mlm_preds = b.mlm(x) # VxSxB
    mlm_preds = reshape(mlm_preds, size(mlm_preds, 1), :) # VxS*B
    nsp_loss = nll(nsp_preds, nsp_labels)
    mlm_labels = reshape(mlm_labels, :) # S*B
    mlm_loss = nll(mlm_preds[:,mlm_labels.!=-1], mlm_labels[mlm_labels.!=-1])
    return mlm_loss + nsp_loss
end

function (b::BertPreTraining)(dtrn)
    lvals = []
    for (x, attention_mask, segment_ids, mlm_labels, nsp_labels) in dtrn
        push!(lvals, b(x, segment_ids, mlm_labels, nsp_labels, attention_mask=attention_mask))
    end
    return Knet.mean(lvals)
end

mutable struct BertClassification <: Layer
    bert::Bert
    pooler::Pooler
    linear::Linear
    pdrop
end

function BertClassification(config, num_of_classes)
    bert = Bert(config)
    pooler = Pooler(config.embed_size, atype=config.atype)
    linear = Linear(config.embed_size, num_of_classes, atype=config.atype)
    return BertClassification(bert, pooler, linear, config.pdrop)
end

function (b::BertClassification)(x, segment_ids; attention_mask=nothing)
    x = b.bert(x, segment_ids, attention_mask=attention_mask)
    x = dropout(b.pooler(x), b.pdrop) # 2xB
    return b.linear(x)
end

function (b::BertClassification)(x, segment_ids, y; attention_mask=nothing)
    return nll(b(x, segment_ids, attention_mask=attention_mask), y)
end

function (b::BertClassification)(dtrn)
    lvals = []
    for (x, attention_mask, segment_ids, y) in dtrn
        push!(lvals, b(x, segment_ids, y, attention_mask=attention_mask))
    end
    return Knet.mean(lvals)
end

mutable struct BertConfig
    embed_size::Int
    vocab_size::Int
    ff_hidden_size::Int
    max_seq_len::Int
    seq_len::Int
    num_segment::Int
    num_heads::Int
    num_encoder::Int
    batchsize::Int
    atype
    pdrop
    attention_pdrop
    func
end

function load_from_torch_base(model, num_encoder, atype, torch_model)
    # Embed Layer
    model.bert.embed_layer.wordpiece.w = Param(atype(permutedims(torch_model["bert.embeddings.word_embeddings.weight"][:cpu]()[:numpy](), (2,1))))
    model.bert.embed_layer.positional.w = Param(atype(permutedims(torch_model["bert.embeddings.position_embeddings.weight"][:cpu]()[:numpy](), (2,1))))
    model.bert.embed_layer.segment.w = Param(atype(permutedims(torch_model["bert.embeddings.token_type_embeddings.weight"][:cpu]()[:numpy](), (2,1))))
    model.bert.embed_layer.layer_norm.γ = Param(atype(torch_model["bert.embeddings.LayerNorm.gamma"][:cpu]()[:numpy]()))
    model.bert.embed_layer.layer_norm.β = Param(atype(torch_model["bert.embeddings.LayerNorm.beta"][:cpu]()[:numpy]()))

    # Encoder Stack
    for i in 1:num_encoder
        model.bert.encoder_stack[i].self_attention.query.w = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.self.query.weight"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.query.b = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.self.query.bias"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.key.w = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.self.key.weight"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.key.b = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.self.key.bias"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.value.w = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.self.value.weight"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.value.b = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.self.value.bias"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.linear.w = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.output.dense.weight"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].self_attention.linear.b = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.output.dense.bias"][:cpu]()[:numpy]()))

        model.bert.encoder_stack[i].layer_norm1.γ = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.output.LayerNorm.gamma"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].layer_norm1.β = Param(atype(torch_model["bert.encoder.layer.$(i-1).attention.output.LayerNorm.beta"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].feed_forward.dense.linear.w = Param(atype(torch_model["bert.encoder.layer.$(i-1).intermediate.dense.weight"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].feed_forward.dense.linear.b = Param(atype(torch_model["bert.encoder.layer.$(i-1).intermediate.dense.bias"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].feed_forward.linear.w = Param(atype(torch_model["bert.encoder.layer.$(i-1).output.dense.weight"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].feed_forward.linear.b = Param(atype(torch_model["bert.encoder.layer.$(i-1).output.dense.bias"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].layer_norm2.γ = Param(atype(torch_model["bert.encoder.layer.$(i-1).output.LayerNorm.gamma"][:cpu]()[:numpy]()))
        model.bert.encoder_stack[i].layer_norm2.β = Param(atype(torch_model["bert.encoder.layer.$(i-1).output.LayerNorm.beta"][:cpu]()[:numpy]()))
    end

    # Pooler
    model.pooler.linear.w = Param(atype(torch_model["bert.pooler.dense.weight"][:cpu]()[:numpy]()))
    model.pooler.linear.b = Param(atype(torch_model["bert.pooler.dense.bias"][:cpu]()[:numpy]()))

    return model
end

function load_from_torch_pretraining(model, num_encoder, atype, torch_model)
    model = load_from_torch_base(model, num_encoder, atype, torch_model)

    # NSP Head
    model.nsp.linear.w = Param(atype(torch_model["cls.seq_relationship.weight"][:cpu]()[:numpy]()))
    model.nsp.linear.b = Param(atype(torch_model["cls.seq_relationship.bias"][:cpu]()[:numpy]()))

    # MLM Head.
    model.mlm.dense.linear.w = Param(atype(torch_model["cls.predictions.transform.dense.weight"][:cpu]()[:numpy]()))
    model.mlm.dense.linear.b = Param(atype(torch_model["cls.predictions.transform.dense.bias"][:cpu]()[:numpy]()))
    model.mlm.layer_norm.γ = Param(atype(torch_model["cls.predictions.transform.LayerNorm.gamma"][:cpu]()[:numpy]()))
    model.mlm.layer_norm.β = Param(atype(torch_model["cls.predictions.transform.LayerNorm.beta"][:cpu]()[:numpy]()))
    model.mlm.linear.w = Param(atype(torch_model["cls.predictions.decoder.weight"][:cpu]()[:numpy]()))
    model.mlm.linear.b = Param(atype(torch_model["cls.predictions.bias"][:cpu]()[:numpy]()))

    return model
end

function load_from_torch_classification(model, num_encoder, atype, torch_model)
    model = load_from_torch_base(model, num_encoder, atype, torch_model)

    model.linear.w = Param(atype(torch_model["classifier.weight"][:cpu]()[:numpy]()))
    model.linear.b = Param(atype(torch_model["classifier.bias"][:cpu]()[:numpy]()))

    return model
end


# Ex 2.1
open("ex1.acc") do file
    for ln in eachline(file)
        println(ln)
    end
end

# Ex 2.2
println("Enter a file name:")
fn = chomp(readline())

open(fn) do f
    for ln in eachline(f)
        println(ln)
    end
end

# Ex 2.3
println("Enter a file name:")
fn = chomp(readline())

nl = open(fn) do f
    i = 0
    for ln in eachline(f)
        i += 1
    end
    i
end

println(nl)

#Ex 2.4-9
for i = 1:3
    avg, sum, lc, neg, pos, zer, max_num, min_num = open("ex1.dat.$i") do f
        lc, sum, neg, pos, zer = 0, 0, 0, 0, 0
        max_num, min_num = -Inf, Inf
        for ln in eachline(f)
            ln_num = parse(Float64, chomp(ln))
            (ln_num > 0) ? pos += 1 : (ln_num < 0) ? neg += 1 : zer += 1
            (ln_num > max_num) && (max_num = ln_num)
            (ln_num < min_num) && (min_num = ln_num)
            lc += 1
            sum += ln_num
        end
        sum / lc, sum, lc, neg, pos, zer, max_num, min_num
    end
    println("Column ", i, " - Average: ", avg, ", Sum: ", sum,
    ", Line count: ", lc)
    println("Negative: ", neg, ", Positive: ", pos, ", Zeros: ", zer)
    println("Maximum number: ", max_num, ", Minimum number: ", min_num)
    println()
end

#=
avg = open("1to9") do f
    i, sum = 0, 0
    for ln in eachline(f)
        i += 1
        sum += parse(Float64, chomp(ln))
    end
    sum / i
end
println(avg)
=#

function get_number(num)
    println(string("Is your number ", string(num), "?"))
    chomp(readline())
end

function guess_number(mini, maxi)
    guess = Int(floor( (mini + maxi) / 2 ))
    response = get_number(guess)
    if response == "higher"
        guess_number(guess + 1, maxi)
    elseif response == "lower"
        guess_number(mini, guess)
    elseif response == "yes"
        println("Gotcha!")
        return
    end
end

guess_number(1,10)


module SIMDIndexingFillArrays

using Base: checkbounds

using ..SIMDIndexing: Index
using ..SIMDIndexing.FillArrays: AbstractFill, getindex_value

# It's probably semantically better to return a `Vec` here but I don't
# want to waste SIMD registers...  It's perhaps the best to add
# `VecFill` type but returning number works most of the time.

@inline function Base.getindex(F::AbstractFill, k::Index)
    @boundscheck checkbounds(F, k)
    getindex_value(F)
end

@inline function Base.getindex(F::AbstractFill{T, N}, kj::Vararg{<:Index, N}) where {T, N}
    @boundscheck checkbounds(F, kj...)
    getindex_value(F)
end

end  # module


@testset "loglikelihoods.jl" begin
    for m in DynamicPPL.TestUtils.DEMO_MODELS
        vi = VarInfo(m)

        vns = vi.metadata.m.vns
        if length(vns) == 1 && length(vi[vns[1]]) == 1
            # Only have one latent variable.
            DynamicPPL.setval!(vi, [1.0], ["m"])
        else
            DynamicPPL.setval!(vi, [1.0, 1.0], ["m[1]", "m[2]"])
        end

        lls = pointwise_loglikelihoods(m, vi)

        if isempty(lls)
            # One of the models with literal observations, so we just skip.
            continue
        end

        loglikelihood = if length(keys(lls)) == 1 && length(m.args.x) == 1
            # Only have one observation, so we need to double it
            # for comparison with other models.
            2 * sum(lls[first(keys(lls))])
        else
            sum(sum, values(lls))
        end

        @test loglikelihood ≈ -324.45158270528947
    end
end


using Distributions

const DISTRIBUTIONS = [Degenerate, Gaussian, Bernoulli]

# The role of these tests is to ensure that the
# the general functionality always works and
# doesn't result in NaN or Inf weights.
# Not that the RBM performance is still good.
if :integration in TEST_GROUPS
    @testset "RBM Integration" begin
        n_vis = 100
        n_hid = 10

        for T in [Float32, Float64]
            rbm = RBM(T, Gaussian, Bernoulli, n_vis, n_hid)
            Boltzmann.test(rbm; debug=true)
        end
    end
end

if :benchmark in TEST_GROUPS
    @testset "RBM Benchmark" begin
        n_vis = 100
        n_hid = 10

        suite = BenchmarkGroup()

        for T in [Float32, Float64]
            rbm = RBM(T, Gaussian, Bernoulli, n_vis, n_hid)
            suite = benchmark!(rbm, suite; debug=true)
        end

        tune!(suite)
        results = run(suite, verbose=true, seconds=10)
    end
end


# This system constructs all the meshes from geometries. These meshes will then be used to be uploaded to OpenGL
abstract type Mesher <: System end

struct PolygonMesher <: Mesher end

Overseer.requested_components(::PolygonMesher) = (PolygonGeometry, Mesh)

struct FileMesher <: Mesher end

Overseer.requested_components(::FileMesher) = (FileGeometry, Mesh)

struct VectorMesher <: Mesher end

Overseer.requested_components(::VectorMesher) = (VectorGeometry, Mesh)

geometry_type(::Type{PolygonMesher}) = PolygonGeometry
geometry_type(::Type{VectorMesher}) = VectorGeometry
geometry_type(::Type{FileMesher}) = FileGeometry

function Overseer.update(::Union{M}, m::AbstractLedger) where {M<:Mesher}
    mesh = m[Mesh]
    geom = m[geometry_type(M)]
    it = @entities_in(geom && !mesh)

    length(it) == 0 && return

    for e in it
        parid = Entity(geom, findfirst(x -> x.geometry == e.geometry, geom.data))
        if parid == e # Geometry was not added to meshes yet
            mesh[e] = Mesh(BasicMesh(e.geometry))
        else
            mesh[e] = parid
        end
    end
end

struct FunctionMesher <: Mesher end

Overseer.requested_components(::FunctionMesher) = (FunctionGeometry, Mesh, Grid)

struct DensityMesher <: Mesher end

Overseer.requested_components(::DensityMesher) = (DensityGeometry, DensityColor, Mesh, Grid)

function Overseer.update(::Union{FunctionMesher,DensityMesher}, m::AbstractLedger)
    mesh = m[Mesh]
    grid = m[Grid]
    dens_c = m[DensityColor]
    for geom in (m[FunctionGeometry], m[DensityGeometry])
        for e in @entities_in(geom && grid && !mesh)
            points = grid[e].points
            vertices, ids = marching_cubes(geom[e].geometry, points, geom[e].iso)
            if isempty(vertices)
                continue
            end
            faces = [GeometryBasics.TriangleFace{GLint}(i, i + 1, i + 2)
                     for i in 1:3:length(vertices)]
            norms = length(faces) == 0 ? Vec3f0[] : normals(vertices, faces)
            mesh[e] = Mesh(BasicMesh(vertices, faces, norms))
            if e in dens_c
                d = dens_c[e]
                m[BufferColor][e] = BufferColor(map(x -> d.color[x...], ids))
            end
        end
    end
end

struct FunctionColorizer <: System end

Overseer.requested_components(::FunctionColorizer) = (FunctionColor, Mesh, BufferColor)

function Overseer.update(::FunctionColorizer, m::AbstractLedger)
    for e in @entities_in(m, FunctionColor && Mesh && !BufferColor)
        e[BufferColor] = BufferColor(e.color.(e.mesh.vertices))
    end
end