text:
### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ f041ef20-f968-11eb-1013-ef1b063dc4cc
begin
	using Pkg
	Pkg.activate(".")
	using Plots
	using SSpline
end

# ╔═╡ 0d9037b8-b0ff-4a4a-8d7f-80b29338ac19
html"""<style> main { max-width: 900px;} """

# ╔═╡ 470b7ee9-b8f7-4c8d-abfd-7fa871abc36f
x = collect(1:10);y = rand(10);

# ╔═╡ d0253bf6-4564-4c3f-9c80-3f59bf3f5370
md"""

(1)

    s[i](x)    = a[i] * (x-x[i])^3 + b[i] * (x-x[i])^2 + c[i] * (x-x[i]) + d[i]

----

    s'[i](x)   = 3 * a[i] * (x-x[i])^2 + 2 * b[i] * (x-x[i]) + c[i]

    s''[i](x)  = 6 * a[i] * (x-x[i]) + 2 * b[i]

    s'''[i](x) = 6 * a[i]

(2)

    s_i(x_i)     = y_i, i=0..n-1

																s_n-1 = y_n

(3)

    s_i(x_i+1)   = s_i+1(x_i+1),      i = 0..n-2

(4)

    s'_i(x_i+1)  = s'_i+1(x_i+1),     i = 0..n-2

(5)

    s''_i(x_i+1) = s''_i+1(x_i+1),    i = 0..n-2

Natural

    s''_0(x_0) = 0
    s''_n-1(x_n) = 0

Clamped

    s'_0(x_0) = y'_0
    s'_n-1(x_n) = y'_n

Periodic

    s'_0(x_0) = s'_n-1(x_n)
    s''_0(x_0) = s''_n-1(x_n)

NAK

    s'''_0(x_1) = s'''_1(x_1)
    s'''_n-2(x_n-1) = s'''_n-1(x_n-1)

----

	h = x_i+1 - x_i

from (1) `s[i](x) = a[i] * (x-x[i])^3 + b[i] * (x-x[i])^2 + c[i] * (x-x[i]) + d[i]`

    d_i = y_i

    σ_i = s''(x_i) = 6 * a_i * (x_i - x_i) + 2 * b_i = 2 * b_i,                          i = 0..n

    σ_i = 2 * b_i

    b_i = σ_i / 2

----

    σ_i+1 = 6 * a_i * h + 2 * b_i

    σ_i+1 = 6 * a_i * h + 2 * σ_i / 2

from (5) `s''_i(x_i+1) = s''_i+1(x_i+1),    i = 0..n-2`

	a_i = (σ_i+1 - σ_i) / (6 * h)

from (3) `s_i(x_i+1) = s_i+1(x_i+1),      i = 0..n-2`

    y_i+1 = a_i * h^3 + b_i * h^2 + c_i * h + y_i

    y_i+1 = (σ_i+1 - σ_i) / (6 * h) * h^3 + σ_i / 2 * h^2 + c_i * h + y_i

    y_i+1 - y_1 = (σ_i+1 - σ_1) / (6 * h) * h^3 + σ_i / 2 * h^2) + c_i * h

    (y_i+1 - y_1) / h = (σ_i+1 - σ_1) / 6 * h + σ_i / 2 * h + c_i

    (y_i+1 - y_1) / h - h * ( (σ_i+1 - σ_1) / 6 + σ_i / 2) = c_i

	c_i = (y_i+1 - y_1) / h - h * ( (σ_i+1 - σ_1) / 6 + 3σ_i / 6 )

    c_i = (y_i+1 - y_1) / h - h * (σ_i+1 - σ_1 + 3 * σ_i) / 6

    c_i = (y_i+1 - y_1) / h - h * (2 * σ_i + σ_i+1) / 6

from (4) `s'_i(x_i+1)  = s'_i+1(x_i+1),     i = 0..n-2`

	c_i+1 = 3 * a_i * h^2 + 2 * b_i * h + c_i

Insert a_i, b_i, c_i, and use

from (2) `s_i(x_i) = y_i, i=0..n-1)`

    c_i+1 = 3 * (σ_i+1 - σ_i) / (6 * h) * h^2 + 2 * σ_i / 2 * h + (y_i+1 - y_1) / h - h * (2 * σ_i + σ_i+1) / 6

"""

# ╔═╡ c82ec4aa-827e-40f0-ab10-543380e5d5d7
begin
	range = collect(1:0.1:10)

	nak = nakspline(x,y)
	kx, ky = interp(nak,range)

	nat = spline3(x,y,:natural)
	nx, ny = interp(nat,range)
end

# ╔═╡ 88c79600-cc20-4ad9-9a98-6d05e2f7512e
nakspline(x,y)

# ╔═╡ 10458fe7-350f-481f-99ca-29e5c185d3b9

# ╔═╡ Cell order:
# ╠═f041ef20-f968-11eb-1013-ef1b063dc4cc
# ╠═0d9037b8-b0ff-4a4a-8d7f-80b29338ac19
# ╠═470b7ee9-b8f7-4c8d-abfd-7fa871abc36f
# ╟─d0253bf6-4564-4c3f-9c80-3f59bf3f5370
# ╟─c82ec4aa-827e-40f0-ab10-543380e5d5d7
# ╠═88c79600-cc20-4ad9-9a98-6d05e2f7512e
# ╠═10458fe7-350f-481f-99ca-29e5c185d3b9


text:
module AMDB

using IntervalSets
using AutoHashEquals
using DataStructures: Accumulator
import Libz
import GZip

export
    # utilities
    process_compressed_database,
    below_maxiter,
    random_push!,
    deserialize_gz,
    # codes
    AMP,
    VMZ,
    # parsing
    parsedtype,
    Gender, male, female,
    AMDB_Date,
    MaybeEmpty,
    AutoEnum,
    LineParser,
    # types
    DateInterval,
    PersonSpell,
    PersonSpellParser,
    AMPSpell,
    AMPSpellParser,
    VMZSpell,
    VMZSpellParser,
    IndividualData,
    # reading
    read_person_spells!,
    read_AMP_spells!,
    read_VMZ_spells!,
    traverse_AMP_spells

include("utilities.jl")
include("AMP_codes.jl")
include("VMZ_codes.jl")
include("parsing.jl")
include("types.jl")
include("reading.jl")
include("analysis.jl")

end # module

ext: jl


text:
#########################################################################
# Making vectors → ranges
#########################################################################
function vector2range(x::AbstractVector{<:Real})
    length(x) < 3 && return x
    dx = x[2]-x[1]
    for i in 3:length(x)
        x[i]-x[i-1] ≠ dx && return x # if no constant step, return array as is
    end
    # do not check value equality, only difference equality
    return x[1]:dx:x[end]
end

function vector2range(t::AbstractVector{<:Dates.AbstractTime})
    tsamp = temporal_sampling(t)
    period = tsamp2period(tsamp)
    isnothing(period) && return t
    t1 = tsamp == :hourly ? t[1] : Date(t[1])
    tf = tsamp == :hourly ? t[end] : Date(t[end])
    r = t1:period:tf
    return r == t ? r : t # final safety check to ensure equal values
end

vector2range(r) = r


text:
"""
**Get information about a taxon at any level**

    taxon(name::String)

This function will look for a taxon by its (scientific) name in the GBIF
reference taxonomy.

Optional arguments are

- `rank::Union{Symbol,Nothing}=:SPECIES` -- the rank of the taxon you want. This
  is part of a controlled vocabulary, and can only be one of `:DOMAIN`,
  `:CLASS`, `:CULTIVAR`, `:FAMILY`, `:FORM`, `:GENUS`, `:INFORMAL`, `:ORDER`,
  `:PHYLUM,`, `:SECTION`, `:SUBCLASS`, `:VARIETY`, `:TRIBE`, `:KINGDOM`,
  `:SUBFAMILY`, `:SUBFORM`, `:SUBGENUS`, `:SUBKINGDOM`, `:SUBORDER`,
  `:SUBPHYLUM`, `:SUBSECTION`, `:SUBSPECIES`, `:SUBTRIBE`, `:SUBVARIETY`,
  `:SUPERCLASS`, `:SUPERFAMILY`, `:SUPERORDER`, and `:SPECIES`

- `strict::Bool=true` -- whether the match should be strict, or fuzzy

Finally, one can also specify other levels of the taxonomy, using  `kingdom`,
`phylum`, `class`, `order`, `family`, and `genus`, all of which can either be
`String` or `Nothing`.

If a match is found, the result will be given as a `GBIFTaxon`. If not, this
function will return `nothing` and give a warning.
"""
function taxon(name::String;
   rank::Union{Symbol,Nothing}=:SPECIES, strict::Bool=true,
   kingdom::Union{String,Nothing}=nothing, phylum::Union{String,Nothing}=nothing, class::Union{String,Nothing}=nothing,
   order::Union{String,Nothing}=nothing, family::Union{String,Nothing}=nothing, genus::Union{String,Nothing}=nothing)
   @assert rank ∈ [
      :DOMAIN, :CLASS, :CULTIVAR, :FAMILY, :FORM, :GENUS, :INFORMAL, :ORDER, :PHYLUM,
      :SECTION, :SUBCLASS, :VARIETY, :TRIBE, :KINGDOM, :SUBFAMILY , :SUBFORM,
      :SUBGENUS, :SUBKINGDOM, :SUBORDER, :SUBPHYLUM, :SUBSECTION , :SUBSERIES,
      :SUBSPECIES, :SUBTRIBE, :SUBVARIETY, :SUPERCLASS , :SUPERFAMILY, :SUPERORDER,
      :SPECIES
   ]
   args = Dict{String, Any}("name" => name, "strict" => strict)

   isnothing(rank) || (args["rank"] = String(rank))
   isnothing(kingdom) || (args["kingdom"] = String(kingdom))
   isnothing(phylum) || (args["phylum"] = String(phylum))
   isnothing(class) || (args["class"] = String(class))
   isnothing(order) || (args["order"] = String(order))
   isnothing(family) || (args["family"] = String(family))
   isnothing(genus) || (args["genus"] = String(genus))

   sp_s_url = gbifurl * "species/match"
   sp_s_req = HTTP.get(sp_s_url, query=args)
   if sp_s_req.status == 200
      body = JSON.parse(String(sp_s_req.body))
      # This will throw warnings for various reasons related to matchtypes
      matchtype = get(body, "matchType", "WELP...")
      # The first one will catch issues
      if matchtype == "WELP..."
         @warn "Impossible to get information for $(name) at level $(rank)"
         return nothing
      end
      if matchtype == "NONE"
         @warn "No match for $(name) at level $(rank) -- try with strict=false"
         return nothing
      end
      return GBIFTaxon(body)
   else
      @warn "Impossible to retrieve information for $(name) -- HTML error code $(sp_s_req.status)"
      return nothing
   end
   return nothing
end

"""
**Get information about a taxon at any level using taxonID**

    taxon(id::Int)

This function will look for a taxon by its taxonID in the GBIF
reference taxonomy.
"""
function taxon(id::Int)
   args = Dict{String, Any}("id" => id)

   sp_s_url = gbifurl * "species/$id"
   sp_s_req = HTTP.get(sp_s_url, query=args)
   if sp_s_req.status == 200
      body = JSON.parse(String(sp_s_req.body))
      return GBIFTaxon(body)
   else
      throw(ErrorException("Impossible to retrieve information for taxonID $(id) -- HTML error code $(sp_s_req.status)"))
   end

end

taxon(t::Pair) = taxon(t.second)


text:
module FakePackage

using InlineTest

greet() = "Hello World!"

const RUN = []

@testset "FakePackage inline begin-end" begin
    @test greet() == "Hello World!"
    @testset "inner" begin
        push!(RUN, 1)
        @test true
    end
end

@testset "FakePackage inline for" for i=1:1
    @test greet() == "Hello World!"
    push!(RUN, 2)
end

module FakeMacro ###########################################

using ReTest # check that @testset_macro is also exported by ReTest
RUN = []

macro fake_macro(x, rec::Bool=true)
    ex = quote
        @testset "macro begin x = $($x)" begin
            @test $x isa Int
            push!(FakeMacro.RUN, $x)
        end
        @testset "macro for $i" for i=1:1
            @testset "plus" begin
                @test i+$x == $x+i
            end
        end
    end
    if rec
        push!(ex.args,
              quote
              if $x == 2
                  @fake_macro 3 false
              end
              end)
    end
    ex
end

@testset_macro @fake_macro

@testset "test macro" begin
    @test true
    @fake_macro 2
end
end # FakeMacro ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

using .FakeMacro: @fake_macro

macro using_fake_macro(x)
    quote
        @testset "inner" begin
            @fake_macro $x
        end
    end
end

@testset_macro @using_fake_macro

@testset "test macro" begin
    @using_fake_macro 4
end

module Sub1 ################################################

using InlineTest
using ..FakePackage: RUN

@testset "check that Sub1 is noticed by ReTest" begin
    push!(RUN, 3)
    @test true
end

end # Sub1 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

module Sub2 ################################################
# Sub2 itself doesn't have any @testset, but has a submodule which has some

module SubSub

using InlineTest
using ...FakePackage: RUN

@testset "check that Sub2.SubSub is noticed by ReTest" begin
    push!(RUN, 4)
    @test true
end

end # SubSub
end # Sub2 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

end # FakePackage

max_issues_repo_name: rfourquet/InlineTest

max_forks_repo_name: rfourquet/InlineTest


text:
# DX = [tx,ty]
function odoAdd(X::Array{Float64,1}, DX::Array{Float64,1})
    # TODO -- this really needs to be cleaned up and optimized! Rather do separate Point2Point2 constraint.
    len = length(X)
    A = eye(3)
    #@show A[2,3] = X[2]
    A[1:len,3] = X

    B = eye(3)
    B[1:len,3] = DX
    #B[2,3] = DX[2]
    retval = zeros(len,1)
    AB = A*B
    retval[1:len,1] = AB[1:len,3]
    #retval[2,1] = AB[2,3]
    return retval
end

function evalPotential(odom::Odo, Xi::Array{Graphs.ExVertex,1}, Xid::Int; N::Int=100)
    warn("evalPotential(::Odo..) deprecated, use functor instead")
    rz,cz = size(odom.Zij)
    # implicit equation portion -- bi-directional pairwise function
    if Xid == Xi[1].index #odom.
        #Z = (odom.Zij\eye(rz)) # this will be used for group operations
        Z = - odom.Zij
        Xval = getVal(Xi[2])
    elseif Xid == Xi[2].index
        Z = odom.Zij
        Xval = getVal(Xi[1])
    else
        error("Bad evalPairwise Odo")
    end

    r,c = size(Xval)
    RES = zeros(r,c*cz)
    # increases the number of particles based on the number of modes in the measurement Z
    for i in 1:(c*cz)
        ent = diag(odom.Cov).*randn(size(vec(Z[:,floor(Int,i/(c+1)+1)])))
        # RES[:,i] = Xval[:,i] + ent+Z[:,floor(Int,i/(c+1)+1)]
        RES[:,i] = odoAdd(Xval[:,i], ent+Z[:,floor(Int,i/(c+1)+1)])
    end
    return RES
end

function evalPotential(odom::OdoMM, Xi::Array{Graphs.ExVertex,1}, Xid::Int; N::Int=100)
    rz,cz = size(odom.Zij)
    # implicit equation portion -- bi-directional pairwise function
    if Xid == Xi[1].index #odom.
        #Z = (odom.Zij\eye(rz)) # this will be used for group operations
        Z = -odom.Zij

        XvalMM = Array{Array{Float64,2},1}(2)
        for i in 2:3
            XvalMM[i-1] = getVal(Xi[i])
        end

        len1 = size(XvalMM[1],2)
        len2 = size(XvalMM[2],2)
        if len1 != len2
            error("odoMM -- must be the same size $(len1), $(len2)")
        #    tmp = kde!(XvalMM[2][:,:], "lcv")
        #    XvalMM[2], = sample(p,len1)
        #elseif len2 > len1
        #    tmp = kde!(XvalMM[1][:,:], "lcv")
        #    XvalMM[1], = sample(p,len2)
        end

        len = len1
        s = rand(len).>0.5 # assume fixed weight for now
        p = (1:len)[s]
        np = (1:len)[!s]
        Xval = zeros(size(XvalMM[1]))
        Xval[:,p] = XvalMM[1][:,p]
        Xval[:,np] = XvalMM[2][:,np]
    elseif Xid == Xi[2].index || Xid == Xi[3].index
        Z = odom.Zij
        Xval = getVal(Xi[1])
    else
        error("Bad evalPairwise OdoMM")
    end

    r,c = size(Xval)
    RES = zeros(r,c*cz)
    # increases the number of particles based on the number of modes in the measurement Z
    for i in 1:(c*cz)
        ent = diag(odom.Cov).*randn(size(vec(Z[:,floor(Int,i/(c+1)+1)])))
        RES[:,i] = odoAdd(Xval[:,i], ent+Z[:,floor(Int,i/(c+1)+1)])
    end
    return RES
end

function rangeAdd(X::Array{Float64,1}, DX::Array{Float64,1})
    A = eye(2)
    A[1,2] = X[1]
    B = eye(2)
    B[1,2] = DX[1]
    return [(A*B)[1,2]]
end

function evalPotential(rang::Ranged, Xi::Array{Graphs.ExVertex,1}, Xid::Int; N::Int=100)
    if Xid == Xi[1].index #rang.
        Z = -rang.Zij
        Xval = getVal(Xi[2])
    elseif Xid == Xi[2].index
        Z = rang.Zij
        Xval = getVal(Xi[1])
    else
        error("Bad evalPairwise Ranged")
    end

    r,c = size(Xval)
    cz = size(rang.Zij,1)
    RES = zeros(r,c*cz)

    for i in 1:(c*cz) # for each mode in the measurement
        ent = rang.Cov[1]*randn(size(vec(Z[:,floor(Int,i/(c+1)+1)])))
        RES[:,i] = rangeAdd(Xval[:,i], ent+Z[floor(Int,i/(c+1)+1)])
    end
    return RES
end


text:
using DifferentialEquations
using Plots

function lotka_volterra(du,u,p,t)
  x, y = u
  a, b, c, d = p
  du[1] = dx = a*x - b*x*y
  du[2] = dy = -c*y + d*x*y
end

# Initial condition
u0 = [1.0, 1.0]

# Simulation interval and intermediary points
tspan = (0.0, 10.0)
tsteps = 0.0:0.1:10.0

# LV equation parameter. p = [α, β, δ, γ]
p = [1.5, 1.0, 3.0, 1.0]

# Setup the ODE problem, then solve
prob = ODEProblem(lotka_volterra, u0, tspan, p)
sol = solve(prob, Tsit5())

# Plot the solution
plot(sol)
savefig("LV_ode.png")

function loss(p)
  sol = solve(prob, Tsit5(), p=p, saveat = tsteps)
  loss = sum(abs2, sol.-1)
  return loss, sol
end

callback = function (p, l, pred)
  display(l)
  plt = plot(pred, ylim = (0, 6))
  display(plt)
  # Tell sciml_train to not halt the optimization. If return true, then
  # optimization stops.
  return false
end

result_ode = DiffEqFlux.sciml_train(loss, p,
                                    cb = callback,
                                    maxiters = 100)

prob_final = ODEProblem(lotka_volterra, u0, tspan, result_ode)
sol_final = solve(prob_final,Tsit5())
plot(sol_final)
savefig("LV_ode_optimized.png") # Both state components close to 1

# Start julia interpreter and type the following
# include("neural_ode.jl") # This will take a while
# After finishing
# Now you can change u0, tspan, p and re-run
# prob = ODEProblem(lotka_volterra,u0,tspan,p)
# sol=solve(prob);plot(sol); # This will be superfast (JIT)

max_stars_repo_name: Romit-Maulik/Tutorials-Demos-Practice

max_issues_repo_name: Romit-Maulik/Tutorials-Demos-Practice

max_forks_repo_name: Romit-Maulik/Tutorials-Demos-Practice


text:
"""
Return 6 surfaces of B-spline solid
"""
function _bsplinesurfaces(M::AbstractBSplineManifold)
    a = controlpoints(M)
    P1, P2, P3 = bsplinespaces(M)
    I1, I2, I3 = bsplineunity(P1), bsplineunity(P2), bsplineunity(P3)
    n1, n2, n3 = dim(P1), dim(P2), dim(P3)

    t_⚀ = minimum(I1)
    t_⚁ = minimum(I2)
    t_⚂ = minimum(I3)
    t_⚃ = maximum(I1)
    t_⚄ = maximum(I2)
    t_⚅ = maximum(I3)

    B_⚀ = [bsplinebasis(i,P1,t_⚀) for i in 1:n1]
    B_⚁ = [bsplinebasis(i,P2,t_⚁) for i in 1:n2]
    B_⚂ = [bsplinebasis(i,P3,t_⚂) for i in 1:n3]
    B_⚃ = [bsplinebasis(i,P1,t_⚃) for i in 1:n1]
    B_⚄ = [bsplinebasis(i,P2,t_⚄) for i in 1:n2]
    B_⚅ = [bsplinebasis(i,P3,t_⚅) for i in 1:n3]

    a_⚀ = sum(a[i1,:,:,:]*B_⚀[i1] for i1 in 1:n1)
    a_⚁ = sum(a[:,i2,:,:]*B_⚁[i2] for i2 in 1:n2)
    a_⚂ = sum(a[:,:,i3,:]*B_⚂[i3] for i3 in 1:n3)
    a_⚃ = sum(a[i1,:,:,:]*B_⚃[i1] for i1 in 1:n1)
    a_⚄ = sum(a[:,i2,:,:]*B_⚄[i2] for i2 in 1:n2)
    a_⚅ = sum(a[:,:,i3,:]*B_⚅[i3] for i3 in 1:n3)

    M_⚀ = BSplineSurface([P2,P3], a_⚀)
    M_⚁ = BSplineSurface([P1,P3], a_⚁)
    M_⚂ = BSplineSurface([P1,P2], a_⚂)
    M_⚃ = BSplineSurface([P2,P3], a_⚃)
    M_⚄ = BSplineSurface([P1,P3], a_⚄)
    M_⚅ = BSplineSurface([P1,P2], a_⚅)

    return (M_⚀, M_⚁, M_⚂, M_⚃, M_⚄, M_⚅)
end


text:
"""
`tbuild_coarse_transition_matrix`: Contruct a transition matrix amongst the bins
(multithreaded version).

### Arguments
* `mutation!` - an in place mutation function
* `bin_id` - bin identification function
* `x0_vals` - an array of starting values
* `bin0_vals` - an array of the bins corresponding to `x0_vals`
* `n_bins` - total number of bins
* `n_trials` - number of independent trials for each x0 starting value
"""
function tbuild_coarse_transition_matrix(mutation!, bin_id, x0_vals, n_bins, n_trials)

   n_x0 = length(x0_vals);
   row_vals = zeros(Int, n_x0*n_trials);
   col_vals = zeros(Int,n_x0*n_trials);

   Threads.@threads for k in 1:n_trials
      for l in 1:n_x0
         X = deepcopy(x0_vals[l]);
         row_vals[n_x0*(k-1) + l] = bin_id(X);
         mutation!(X);
         col_vals[n_x0*(k-1) + l] = bin_id(X);
      end
   end

   K = sparse(row_vals,col_vals,ones(size(row_vals)),n_bins, n_bins);
   for i in 1:n_bins
      rowZ = sum(K[i,:]);
      @. K[i,:] = K[i,:]/rowZ;
   end

   return K
end


text:
using Weave
path = dirname(@__FILE__)
examples=[
    "benders_quatile_regression.jl"
]
cd(path)
for ex in examples
    weave(ex)
    # weave(ex, doctype = "github")
    # weave(ex, doctype = "notebook")
    # Weave.convert_doc(ex, ex[1:end-3]*".ipynb")
    # weave(ex, doctype = "md2pdf")
end


text:
#Figure 5.9, 5.10
using Plots, SparseArrays, JLD2, LinearAlgebra, MultiscaleGraphSignalTransforms, PyCall
py"""
import numpy
xs = numpy.load('../data/texture_mask.npy')
"""
xs = PyArray(py"xs"o)

matrix_gt = reshape(xs,(512,512))[1:4:512,1:4:512]
heatmap(matrix_gt,ratio=1, yaxis =:flip, axis = false, color = :grays, colorbar = false)
max1 = maximum(matrix_gt[:])
min1 = minimum(matrix_gt[:])
matrix_gt = (matrix_gt[:] .- min1)./(max1 - min1)

JLD2.@load "../data/handcut_images.jld2"
heatmap(tmp["block5"],ratio=1, yaxis =:flip, axis = false, color = :grays, colorbar = false)
matrix = tmp["block5"]

########################

sigma = 0.001
m, n = size(matrix)

############### Get the neighbors for affinity matrix computing
r = 5; # specify radius of neighbors
l = 2*r + 1

temp_x, temp_y = fill(1,l^2), fill(1,l^2)
temp_xy = CartesianIndices((1:l,1:l))
for i = 1:l^2
    temp_x[i], temp_y[i] = temp_xy[i][1], temp_xy[i][2]
end
temp_ind = ((temp_x .- (r + 1)).^2 + (temp_y .- (r + 1)).^2).<= r^2
neighbor_x = temp_x[temp_ind] .- (r + 1)
neighbor_y = temp_y[temp_ind] .- (r + 1)
# for any index(x,y), (x + neighbor_x, y - neigbor_y) are the neighbors to calculate affinity

################ Create affinity matrix
W = fill(0., (m*n,m*n))
for i = 1:m*n
    cur = CartesianIndices((m,n))[i]
    for j = 1:length(neighbor_x)
        if 1 <= cur[1] + neighbor_x[j] <= m && 1 <= cur[2] + neighbor_y[j] <= n
            tempd = LinearIndices((m,n))[cur[1] + neighbor_x[j], cur[2] + neighbor_y[j]]
            W[i,tempd] = exp(-(matrix_gt[i] - matrix_gt[tempd])^2/sigma)
        end
    end
end
W = sparse(W)

############### Preprocess to get G and GP
G = GraphSig(W, f = reshape(matrix,(length(matrix_gt),1)))
GP = partition_tree_fiedler(G)
dmatrix = ghwt_analysis!(G, GP=GP)

#####################
############# Haar
BS_haar = bs_haar(GP)
dvec_haar = dmatrix2dvec(dmatrix, GP, BS_haar)

############# Walsh
BS_walsh = bs_walsh(GP)
dvec_walsh = dmatrix2dvec(dmatrix, GP, BS_walsh)

############# GHWT_c2f
dvec_c2f, BS_c2f = ghwt_c2f_bestbasis(dmatrix, GP)

############# GHWT_f2c
dvec_f2c, BS_f2c = ghwt_f2c_bestbasis(dmatrix, GP)

############# eGHWT
dvec_eghwt, BS_eghwt = ghwt_tf_bestbasis(dmatrix, GP)

function top_vectors_plot(dvec::Array{Float64, 2}, BS::BasisSpec, GP::GraphPart; clims::Tuple{Float64,Float64} = (-0.01,0.01))
    sorted_ind = sortperm(abs.(dvec[:]), rev = true);
    plot(9,layout=(3,3),framestyle=:none, legend=false)
    for i=1:9
        dvecT = fill(0., size(dvec))
        #dvecT[sorted_ind[i]] = dvec_ghwt[sorted_ind[i]]
        dvecT[sorted_ind[i]] = 1
        f = ghwt_synthesis(dvecT, GP, BS)
        #print((maximum(f), minimum(f)))
        heatmap!(reshape(f, size(matrix)), subplot=i, ratio=1, yaxis=:flip, axis=false, color = :grays, clims = clims)
    end
    current()
end

#Figure 5.10b
top_vectors_plot(dvec_eghwt, BS_eghwt, GP)

function approx_error2(DVEC::Array{Array{Float64,1},1})
    plot(xaxis = "Fraction of Coefficients Retained", yaxis = "Relative Approximation Error")
    frac = 0.3
    T = ["eGHWT Haar basis","eGHWT Walsh basis","GHWT_c2f", "GHWT_f2c", "eGHWT best basis"]
    L = [(:dashdot,:orange),(:dashdot,:blue),(:solid, :red),(:solid, :green),(:solid, :black)]
    for i = 1:5
        dvec = DVEC[i]
        N = length(dvec)
        dvec_norm = norm(dvec,2)
        dvec_sort = sort(dvec.^2) # the smallest first
        er = sqrt.(reverse(cumsum(dvec_sort)))/dvec_norm # this is the relative L^2 error of the whole thing, i.e., its length is N
        p = Int64(floor(frac*N)) + 1 # upper limit
        plot!(frac*(0:(p-1))/(p-1), er[1:p], yaxis=:log, xlims = (0.,frac), label = T[i], line = L[i])
    end
end

#Figure 5.10a
approx_error2([dvec_haar[:], dvec_walsh[:], dvec_c2f[:], dvec_f2c[:], dvec_eghwt[:]])


text:
export NBCopulaCSObs, NBCopulaCSModel
### first structures
mutable struct NBCopulaCSObs{T <: BlasReal, D, Link} # d changes, so must be mutable
    # data
    n::Int
    p::Int
    y::Vector{T}
    X::Matrix{T}
    V::SymmetricToeplitz{T}
    vec::Vector{T}
    # working arrays
    ∇CSV::SymmetricToeplitz{T}
    ∇β::Vector{T}   # gradient wrt β
    ∇resβ::Matrix{T}# residual gradient matrix d/dβ_p res_ij (each observation has a gradient of residual is px1)
    ∇ρ::Vector{T}
    ∇σ2::Vector{T}
    ∇r::Vector{T}   # gradient wrt r (NB)
    Hβ::Matrix{T}   # Hessian wrt β
    Hρ::Matrix{T}   # Hessian wrt ρ
    Hr::Matrix{T}   # Hessian wrt r (NB)
    Hσ2::Matrix{T}   # Hessian wrt ρ
    Hρσ2::Matrix{T}  # Hessian wrt ρ, σ2
    Hβσ2::Vector{T}  # Hessian wrt β and σ2
    # Hβρ::Vector{T}
    res::Vector{T}  # residual vector res_i
    t::Vector{T}    # t[k] = tr(V_i[k]) / 2
    q::Vector{T}    # q[k] = res_i' * V_i[k] * res_i / 2
    xtx::Matrix{T}  # Xi'Xi
    storage_n::Vector{T}
    storage_n2::Vector{T}
    storage_n3::Vector{T}
    storage_p1::Vector{T}
    storage_np::Matrix{T}
    storage_pp::Matrix{T}
    added_term_numerator::Matrix{T}
    added_term2::Matrix{T}
    η::Vector{T}    # η = Xβ systematic component
    μ::Vector{T}    # μ(β) = ginv(Xβ) # inverse link of the systematic component
    varμ::Vector{T} # v(μ_i) # variance as a function of the mean
    dμ::Vector{T}   # derivative of μ
    d::D            # distribution()
    link::Link      # link function ()
    wt::Vector{T}   # weights wt for GLM.jl
    w1::Vector{T}   # working weights in the gradient = dμ/v(μ)
    w2::Vector{T}   # working weights in the information matrix = dμ^2/v(μ)
end

function NBCopulaCSObs(
    y::Vector{T},
    X::Matrix{T},
    d::D,
    link::Link) where {T <: BlasReal, D, Link}
    n, p = size(X, 1), size(X, 2)
    @assert length(y) == n "length(y) should be equal to size(X, 1)"
    # working arrays
    V = SymmetricToeplitz(ones(T, n))
    vec = Vector{T}(undef, n)
    ∇CSV = SymmetricToeplitz(ones(T, n))
    ∇β  = Vector{T}(undef, p)
    ∇resβ  = Matrix{T}(undef, n, p)
    ∇ρ  = Vector{T}(undef, 1)
    ∇σ2  = Vector{T}(undef, 1)
    ∇r  = Vector{T}(undef, 1)
    Hβ  = Matrix{T}(undef, p, p)
    Hρ  = Matrix{T}(undef, 1, 1)
    Hr  = Matrix{T}(undef, 1, 1)
    Hσ2  = Matrix{T}(undef, 1, 1)
    Hρσ2 = Matrix{T}(undef, 1, 1)
    Hβσ2 = zeros(T, p)
    # Hβρ = Vector{T}(undef, p)
    res = Vector{T}(undef, n)
    t   = [tr(V)/2]
    q   = Vector{T}(undef, 1)
    xtx = transpose(X) * X
    storage_n = Vector{T}(undef, n)
    storage_n2 = Vector{T}(undef, n)
    storage_n3 = Vector{T}(undef, n)
    storage_p1 = Vector{T}(undef, p)
    storage_np = Matrix{T}(undef, n, p)
    storage_pp = Matrix{T}(undef, p, p)
    added_term_numerator = Matrix{T}(undef, n, p)
    added_term2 = Matrix{T}(undef, p, p)
    η = Vector{T}(undef, n)
    μ = Vector{T}(undef, n)
    varμ = Vector{T}(undef, n)
    dμ = Vector{T}(undef, n)
    wt = Vector{T}(undef, n)
    fill!(wt, one(T))
    w1 = Vector{T}(undef, n)
    w2 = Vector{T}(undef, n)
    # constructor
    NBCopulaCSObs{T, D, Link}(n, p, y, X, V, vec, ∇CSV, ∇β,
        ∇resβ, ∇ρ, ∇σ2, ∇r, Hβ, Hρ, Hr, Hσ2, Hρσ2, Hβσ2, # Hβρ,
        res, t, q, xtx, storage_n, storage_n2, storage_n3, storage_p1, storage_np,
        storage_pp, added_term_numerator, added_term2,
        η, μ, varμ, dμ, d, link, wt, w1, w2)
end

"""
    NBCopulaCSModel(gcs)

Negative Binomial copula variance component model, which contains a vector of
`NBCopulaCSObs` as data, model parameters, and working arrays.
"""
struct NBCopulaCSModel{T <: BlasReal, D, Link} <: MathProgBase.AbstractNLPEvaluator
    # data
    data::Vector{NBCopulaCSObs{T, D, Link}}
    Ytotal::T
    ntotal::Int     # total number of singleton observations
    p::Int          # number of mean parameters in linear regression
    # parameters
    β::Vector{T}    # p-vector of mean regression coefficients
    τ::Vector{T}    # inverse of linear regression variance parameter (IS THIS NEEDED??)
    ρ::Vector{T}            # autocorrelation parameter
    σ2::Vector{T}           # autoregressive noise parameter
    θ::Vector{T}
    r::Vector{T}   # r parameter for negative binomial
    # working arrays
    ∇β::Vector{T}   # gradient of beta from all observations
    ∇ρ::Vector{T}           # gradient of rho from all observations
    ∇σ2::Vector{T}          # gradient of sigmasquared from all observations
    ∇r::Vector{T}
    XtX::Matrix{T}  # X'X = sum_i Xi'Xi
    Hβ::Matrix{T}    # Hessian from all observations
    Hρ::Matrix{T}    # Hessian from all observations
    Hσ2::Matrix{T}    # Hessian from all observations
    Hr::Matrix{T}
    Hρσ2::Matrix{T}
    Hβσ2::Vector{T}
    Ainv::Matrix{T}
    Aevec::Matrix{T}
    M::Matrix{T}
    vcov::Matrix{T}
    ψ::Vector{T}
    # Hβρ::Vector{T}
    TR::Matrix{T}
    QF::Matrix{T}         # n-by-1 matrix with qik = res_i' Vi res_i
    storage_n::Vector{T}
    storage_m::Vector{T}
    storage_θ::Vector{T}
    d::Vector{D}
    link::Vector{Link}
    penalized::Bool
end

function NBCopulaCSModel(gcs::Vector{NBCopulaCSObs{T, D, Link}}; penalized::Bool = false) where {T <: BlasReal, D, Link}
    n, p = length(gcs), size(gcs[1].X, 2)
    β   = Vector{T}(undef, p)
    τ   = [one(T)]
    ρ = [one(T)]
    σ2 = [one(T)]
    θ = [one(T)]
    r = [one(T)]
    ∇β  = Vector{T}(undef, p)
    ∇ρ  = Vector{T}(undef, 1)
    ∇σ2  = Vector{T}(undef, 1)
    ∇r  = Vector{T}(undef, 1)
    XtX = zeros(T, p, p) # sum_i xi'xi
    Hβ  = Matrix{T}(undef, p, p)
    Hρ  = Matrix{T}(undef, 1, 1)
    Hσ2  = Matrix{T}(undef, 1, 1)
    Hr  = Matrix{T}(undef, 1, 1)
    Hρσ2 = Matrix{T}(undef, 1, 1)
    Hβσ2 = zeros(T, p)
    Ainv    = zeros(T, p + 3, p + 3)
    Aevec   = zeros(T, p + 3, p + 3)
    M       = zeros(T, p + 3, p + 3)
    vcov    = zeros(T, p + 3, p + 3)
    ψ       = Vector{T}(undef, p + 3)
    # Hβρ = Vector{T}(undef, p)
    TR  = Matrix{T}(undef, n, 1) # collect trace terms
    QF  = Matrix{T}(undef, n, 1)
    Ytotal = zero(T)
    ntotal = 0
    d = Vector{D}(undef, n)
    link = Vector{Link}(undef, n)
    for i in eachindex(gcs)
        ntotal  += length(gcs[i].y)
        Ytotal  += sum(gcs[i].y)
        BLAS.axpy!(one(T), gcs[i].xtx, XtX)
        TR[i, :] = gcs[i].t
        d[i] = gcs[i].d
        link[i] = gcs[i].link
    end
    storage_n = Vector{T}(undef, n)
    storage_m = Vector{T}(undef, 1)
    storage_θ = Vector{T}(undef, 1)
    NBCopulaCSModel{T, D, Link}(gcs, Ytotal, ntotal, p, β, τ, ρ, σ2, θ, r,
        ∇β, ∇ρ, ∇σ2, ∇r, XtX, Hβ, Hρ, Hσ2, Hr, Hρσ2, Hβσ2, Ainv, Aevec,  M, vcov, ψ,
        TR, QF, storage_n, storage_m, storage_θ, d, link, penalized)
end

function loglikelihood!(
    gc::NBCopulaCSObs{T, D, Link},
    β::Vector{T},
    ρ::T,
    σ2::T,
    r::T,
    needgrad::Bool = false,
    needhess::Bool = false;
    penalized::Bool = false
    ) where {T <: BlasReal, D, Link}
    # n, p = size(gc.X, 1), size(gc.X, 2)
    needgrad = needgrad || needhess
    if needgrad
        fill!(gc.∇β, 0)
        get_∇V!(gc)
    end
    needhess && fill!(gc.Hβ, 0)
    fill!(gc.∇β, 0.0)
    update_res!(gc, β)
    standardize_res!(gc)
    fill!(gc.∇resβ, 0.0) # fill gradient of residual vector with 0
    std_res_differential!(gc) # this will compute ∇resβ

    # form V
    get_V!(ρ, gc)

    #evaluate copula loglikelihood
    mul!(gc.storage_n, gc.V, gc.res, one(T), zero(T)) # storage_n = V[k] * res

    if needgrad
        BLAS.gemv!('T', σ2, gc.∇resβ, gc.storage_n, 1.0, gc.∇β) # stores ∇resβ*Γ*res (standardized residual)
    end
    q = dot(gc.res, gc.storage_n)

    # @show q
    c1 = 1 + 0.5 * gc.n * σ2
    c2 = 1 + 0.5 * σ2 * q
    # loglikelihood
    logl = QuasiCopula.component_loglikelihood(gc, r)
    logl += -log(c1)
    # @show logl
    logl += log(c2)
    # add L2 ridge penalty
    if penalized
        logl -= 0.5 * (σ2)^2
    end
    # @show logl
    if needgrad
        inv1pq = inv(c2)
        # gradient with respect to rho
        mul!(gc.storage_n, gc.∇CSV, gc.res, one(T), zero(T)) # storage_n = ∇ARV * res
        q2 = dot(gc.res, gc.storage_n) #

        gc.∇ρ .= inv(c2) * 0.5 * σ2 * q2

        # gradient with respect to sigma2
        gc.∇σ2 .= -0.5 * gc.n * inv(c1) .+ inv(c2) * 0.5 * q
        if penalized
            gc.∇σ2 .-= σ2
        end

        # gradient with respect to r
        gc.∇r .= nb_first_derivative(gc, ρ, σ2, r)

      if needhess
            # hessian for rho
            gc.Hρ .= -(inv(c2) * 0.5 * σ2 * q2)^2

            # hessian for sigma2
            gc.Hσ2 .= 0.25 * gc.n^2 * inv(c1)^2 - inv1pq^2 * (0.25 * q^2)
            if penalized
                gc.Hσ2 .-= 1.0
            end

            # hessian cross term for rho and sigma2
            gc.Hρσ2 .= 0.5 * q2 * inv1pq - 0.25 * σ2 * q * q2 * inv1pq^2

            # hessian cross term for beta and sigma2
            gc.Hβσ2 .= inv1pq * gc.∇β - 0.5 * q * inv1pq^2 * σ2 * gc.∇β

            BLAS.syrk!('L', 'N', -abs2(inv1pq), gc.∇β, 0.0, gc.Hβ) # only lower triangular
            fill!(gc.added_term_numerator, 0.0) # fill gradient with 0
            fill!(gc.added_term2, 0.0) # fill hessian with 0
            mul!(gc.added_term_numerator, gc.V, gc.∇resβ, one(T), zero(T))
            BLAS.gemm!('T', 'N', σ2, gc.∇resβ, gc.added_term_numerator, one(T), gc.added_term2)
            gc.added_term2 .*= inv1pq
            gc.Hβ .+= gc.added_term2
            gc.Hβ .+= QuasiCopula.glm_hessian(gc)
            # hessian for r
            gc.Hr .= nb_second_derivative(gc, ρ, σ2, r)
      end
      gc.∇β .= gc.∇β .* inv1pq
      gc.res .= gc.y .- gc.μ
      gc.∇β .+= QuasiCopula.glm_gradient(gc)
    end
    logl
end

function loglikelihood!(
    gcm::NBCopulaCSModel{T, D, Link},
    needgrad::Bool = false,
    needhess::Bool = false
    ) where {T <: BlasReal, D, Link}
    if needgrad
        fill!(gcm.∇β, 0.0)
        fill!(gcm.∇ρ, 0.0)
        fill!(gcm.∇σ2, 0.0)
        fill!(gcm.∇r, 0.0)
    end
    if needhess
        fill!(gcm.Hβ, 0.0)
        fill!(gcm.Hρ, 0.0)
        fill!(gcm.Hσ2, 0.0)
        fill!(gcm.Hρσ2, 0.0)
        fill!(gcm.Hβσ2, 0.0)
        fill!(gcm.Hr, 0.0)
        # fill!(gcm.Hβρ, 0.0)
    end
    logl = zeros(Threads.nthreads())
    Threads.@threads for i in eachindex(gcm.data)
        @inbounds logl[Threads.threadid()] += loglikelihood!(gcm.data[i], gcm.β,
        gcm.ρ[1], gcm.σ2[1], gcm.r[1], needgrad, needhess; penalized = gcm.penalized)
    end
    @inbounds for i in eachindex(gcm.data)
        if needgrad
            gcm.∇β .+= gcm.data[i].∇β
            gcm.∇ρ .+= gcm.data[i].∇ρ
            gcm.∇σ2 .+= gcm.data[i].∇σ2
            gcm.∇r .+= gcm.data[i].∇r
        end
        if needhess
            gcm.Hβ .+= gcm.data[i].Hβ
            gcm.Hρ .+= gcm.data[i].Hρ
            gcm.Hσ2 .+= gcm.data[i].Hσ2
            gcm.Hρσ2 .+= gcm.data[i].Hρσ2
            gcm.Hβσ2 .+= gcm.data[i].Hβσ2
            gcm.Hr .+= gcm.data[i].Hr
        end
    end
    return sum(logl)
end


text:
# ---
# title: 558. Logical OR of Two Binary Grids Represented as Quad-Trees
# id: problem558
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories:
# link: <https://leetcode.com/problems/logical-or-of-two-binary-grids-represented-as-quad-trees/description/>
# hidden: true
# ---
#
# A Binary Matrix is a matrix in which all the elements are either **0** or
# **1**.
#
# Given `quadTree1` and `quadTree2`. `quadTree1` represents a `n * n` binary
# matrix and `quadTree2` represents another `n * n` binary matrix.
#
# Return _a Quad-Tree_ representing the `n * n` binary matrix which is the
# result of **logical bitwise OR** of the two binary matrixes represented by
# `quadTree1` and `quadTree2`.
#
# Notice that you can assign the value of a node to **True** or **False** when
# `isLeaf` is **False** , and both are **accepted** in the answer.
#
# A Quad-Tree is a tree data structure in which each internal node has exactly
# four children. Besides, each node has two attributes:
#
#   * `val`: True if the node represents a grid of 1's or False if the node represents a grid of 0's.
#   * `isLeaf`: True if the node is leaf node on the tree or False if the node has the four children.
#
#
#
#     class Node {
#         public boolean val;
#         public boolean isLeaf;
#         public Node topLeft;
#         public Node topRight;
#         public Node bottomLeft;
#         public Node bottomRight;
#     }
#
# We can construct a Quad-Tree from a two-dimensional area using the following
# steps:
#
#   1. If the current grid has the same value (i.e all `1's` or all `0's`) set `isLeaf` True and set `val` to the value of the grid and set the four children to Null and stop.
#   2. If the current grid has different values, set `isLeaf` to False and set `val` to any value and divide the current grid into four sub-grids as shown in the photo.
#   3. Recurse for each of the children with the proper sub-grid.
#
# ![](https://assets.leetcode.com/uploads/2020/02/11/new_top.png)
#
# If you want to know more about the Quad-Tree, you can refer to the
# [wiki](https://en.wikipedia.org/wiki/Quadtree).
#
# **Quad-Tree  format:**
#
# The input/output represents the serialized format of a Quad-Tree using level
# order traversal, where `null` signifies a path terminator where no node exists
# below.
#
# It is very similar to the serialization of the binary tree. The only
# difference is that the node is represented as a list `[isLeaf, val]`.
#
# If the value of `isLeaf` or `val` is True we represent it as **1** in the list
# `[isLeaf, val]` and if the value of `isLeaf` or `val` is False we represent it
# as **0**.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/02/11/qt1.png)
# ![](https://assets.leetcode.com/uploads/2020/02/11/qt2.png)
#
#
#
#     Input: quadTree1 = [[0,1],[1,1],[1,1],[1,0],[1,0]]
#     , quadTree2 = [[0,1],[1,1],[0,1],[1,1],[1,0],null,null,null,null,[1,0],[1,0],[1,1],[1,1]]
#     Output: [[0,0],[1,1],[1,1],[1,1],[1,0]]
#     Explanation: quadTree1 and quadTree2 are shown above. You can see the binary matrix which is represented by each Quad-Tree.
#     If we apply logical bitwise OR on the two binary matrices we get the binary matrix below which is represented by the result Quad-Tree.
#     Notice that the binary matrices shown are only for illustration, you don't have to construct the binary matrix to get the result tree.
#     ![](https://assets.leetcode.com/uploads/2020/02/11/qtr.png)
#
#
# **Example 2:**
#
#
#
#     Input: quadTree1 = [[1,0]]
#     , quadTree2 = [[1,0]]
#     Output: [[1,0]]
#     Explanation: Each tree represents a binary matrix of size 1*1. Each matrix contains only zero.
#     The resulting matrix is of size 1*1 with also zero.
#
#
# **Example 3:**
#
#
#
#     Input: quadTree1 = [[0,0],[1,0],[1,0],[1,1],[1,1]]
#     , quadTree2 = [[0,0],[1,1],[1,1],[1,0],[1,1]]
#     Output: [[1,1]]
#
#
# **Example 4:**
#
#
#
#     Input: quadTree1 = [[0,0],[1,1],[1,0],[1,1],[1,1]]
#     , quadTree2 = [[0,0],[1,1],[0,1],[1,1],[1,1],null,null,null,null,[1,1],[1,0],[1,0],[1,1]]
#     Output: [[0,0],[1,1],[0,1],[1,1],[1,1],null,null,null,null,[1,1],[1,0],[1,0],[1,1]]
#
#
# **Example 5:**
#
#
#
#     Input: quadTree1 = [[0,1],[1,0],[0,1],[1,1],[1,0],null,null,null,null,[1,0],[1,0],[1,1],[1,1]]
#     , quadTree2 = [[0,1],[0,1],[1,0],[1,1],[1,0],[1,0],[1,0],[1,1],[1,1]]
#     Output: [[0,0],[0,1],[0,1],[1,1],[1,0],[1,0],[1,0],[1,1],[1,1],[1,0],[1,0],[1,1],[1,1]]
#
#
#
#
# **Constraints:**
#
#   * `quadTree1` and `quadTree2` are both **valid** Quad-Trees each representing a `n * n` grid.
#   * `n == 2^x` where `0 <= x <= 9`.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
module Equalities

include("./../../constraints/geometric/constants.jl")
using JuMP
using .Constants

export first_dim_equality

function first_dim_equality(m)

    # 1 flbxi == flbxj
    @variable(m, if11_eq_flbx[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flbx[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flbx[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    #Lz1 >= xi-xj+e
    #L(1-z1) >= xj-xi-e
    @constraint(m, if_eq_flbx_11[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if11_eq_flbx][i, j] >=
        m[:flbx_front_left_bot_x][i] - m[:flbx_front_left_bot_x][j] + 0.0001
    )
    @constraint(m, if_eq_flbx_12[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if11_eq_flbx][i, j]) >=
        m[:flbx_front_left_bot_x][j] - m[:flbx_front_left_bot_x][i] - 0.0001
    )

    #Lz2 >= xj-xi+e
    #L(1-z2) >= xi-xj-e
    @constraint(m, if_eq_flbx_13[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if12_eq_flbx][i, j] >=
        m[:flbx_front_left_bot_x][j] - m[:flbx_front_left_bot_x][i] + 0.0001
    )
    @constraint(m, if_eq_flbx_14[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if12_eq_flbx][i, j]) >=
        m[:flbx_front_left_bot_x][i] - m[:flbx_front_left_bot_x][j] - 0.0001
    )

    @constraint(m, if1_eq_flbx_cons[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flbx][i, j] + m[:if12_eq_flbx][i, j] >=
        2 * m[:if1_eq_flbx][i, j]
    )

    # 2 flbxi + box_x_axis_cover_i == flbxj + box_x_axis_cover_j
    @variable(m, if11_eq_flbx_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flbx_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flbx_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    @constraint(m, if_eq_flbx_11_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if11_eq_flbx_ij_plus][i, j] >=
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) + 0.0001
    )
    @constraint(m, if_eq_flbx_12_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if11_eq_flbx_ij_plus][i, j]) >=
        (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - 0.0001
    )
    @constraint(m, if_eq_flbx_13_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if12_eq_flbx_ij_plus][i, j] >=
        (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) + 0.0001
    )
    @constraint(m, if_eq_flbx_14_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if12_eq_flbx_ij_plus][i, j]) >=
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - 0.0001
    )

    @constraint(m, if1_eq_flbx_cons_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flbx_ij_plus][i, j] + m[:if12_eq_flbx_ij_plus][i, j] >=
        2 * m[:if1_eq_flbx_ij_plus][i, j]
    )

    # 3 flbxi + box_x_axis_cover_i == flbxj
    @variable(m, if11_eq_flbx_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flbx_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flbx_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    @constraint(m, if_eq_flbx_11_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if11_eq_flbx_i_plus][i, j] >=
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - m[:flbx_front_left_bot_x][j] + 0.0001
    )
    @constraint(m, if_eq_flbx_12_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if11_eq_flbx_i_plus][i, j]) >=
        m[:flbx_front_left_bot_x][j] - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - 0.0001
    )
    @constraint(m, if_eq_flbx_13_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if12_eq_flbx_i_plus][i, j] >=
        m[:flbx_front_left_bot_x][j] - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) + 0.0001
    )
    @constraint(m, if_eq_flbx_14_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if12_eq_flbx_i_plus][i, j]) >=
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]) - m[:flbx_front_left_bot_x][j] - 0.0001
    )

    @constraint(m, if1_eq_flbx_cons_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flbx_i_plus][i, j] + m[:if12_eq_flbx_i_plus][i, j] >=
        2 * m[:if1_eq_flbx_i_plus][i, j]
    )

    # 4 flbxj + box_x_axis_cover_j == flbxi
    @variable(m, if11_eq_flbx_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flbx_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flbx_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    @constraint(m, if_eq_flbx_11_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if11_eq_flbx_j_plus][i, j] >=
        m[:flbx_front_left_bot_x][i] - (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) + 0.0001
    )
    @constraint(m, if_eq_flbx_12_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if11_eq_flbx_j_plus][i, j]) >=
        (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - m[:flbx_front_left_bot_x][i] - 0.0001
    )
    @constraint(m, if_eq_flbx_13_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * m[:if12_eq_flbx_j_plus][i, j] >=
        (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - m[:flbx_front_left_bot_x][i] + 0.0001
    )
    @constraint(m, if_eq_flbx_14_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        L * (1 - m[:if12_eq_flbx_j_plus][i, j]) >=
        m[:flbx_front_left_bot_x][i] - (m[:flbx_front_left_bot_x][j] + m[:box_x_axis_cover][j]) - 0.0001
    )

    @constraint(m, if1_eq_flbx_cons_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flbx_j_plus][i, j] + m[:if12_eq_flbx_j_plus][i, j] >=
        2 * m[:if1_eq_flbx_j_plus][i, j]
    )

    # ==== Y AXIS =====

    # 1 flbyi == flbyj
    @variable(m, if11_eq_flby[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flby[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flby[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    #Wz1 >= yi-yj+e
    #W(1-z1) >= yj-yi-e
    @constraint(m, if_eq_flby_11[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if11_eq_flby][i, j] >=
        m[:flby_front_left_bot_y][i] - m[:flby_front_left_bot_y][j] + 0.0001
    )
    @constraint(m, if_eq_flby_12[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if11_eq_flby][i, j]) >=
        m[:flby_front_left_bot_y][j] - m[:flby_front_left_bot_y][i] - 0.0001
    )

    #Wz2 >= yj-yi+e
    #W(1-z2) >= yi-yj-e
    @constraint(m, if_eq_flby_13[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if12_eq_flby][i, j] >=
        m[:flby_front_left_bot_y][j] - m[:flby_front_left_bot_y][i] + 0.0001
    )
    @constraint(m, if_eq_flby_14[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if12_eq_flby][i, j]) >=
        m[:flby_front_left_bot_y][i] - m[:flby_front_left_bot_y][j] - 0.0001
    )

    @constraint(m, if1_eq_flby_cons[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flby][i, j] + m[:if12_eq_flby][i, j] >=
        2 * m[:if1_eq_flby][i, j]
    )

    # 2 flbyi + box_y_axis_cover_i == flbyj + box_y_axis_cover_j
    @variable(m, if11_eq_flby_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flby_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flby_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    @constraint(m, if_eq_flby_11_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if11_eq_flby_ij_plus][i, j] >=
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) + 0.0001
    )
    @constraint(m, if_eq_flby_12_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if11_eq_flby_ij_plus][i, j]) >=
        (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - 0.0001
    )
    @constraint(m, if_eq_flby_13_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if12_eq_flby_ij_plus][i, j] >=
        (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) + 0.0001
    )
    @constraint(m, if_eq_flby_14_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if12_eq_flby_ij_plus][i, j]) >=
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - 0.0001
    )

    @constraint(m, if1_eq_flby_cons_ij_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flby_ij_plus][i, j] + m[:if12_eq_flby_ij_plus][i, j] >=
        2 * m[:if1_eq_flby_ij_plus][i, j]
    )

    # 3 flbyi + box_y_axis_cover_i == flbyj
    @variable(m, if11_eq_flby_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flby_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flby_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    @constraint(m, if_eq_flby_11_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if11_eq_flby_i_plus][i, j] >=
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - m[:flby_front_left_bot_y][j] + 0.0001
    )
    @constraint(m, if_eq_flby_12_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if11_eq_flby_i_plus][i, j]) >=
        m[:flby_front_left_bot_y][j] - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - 0.0001
    )
    @constraint(m, if_eq_flby_13_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if12_eq_flby_i_plus][i, j] >=
        m[:flby_front_left_bot_y][j] - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) + 0.0001
    )
    @constraint(m, if_eq_flby_14_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if12_eq_flby_i_plus][i, j]) >=
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]) - m[:flby_front_left_bot_y][j] - 0.0001
    )

    @constraint(m, if1_eq_flby_cons_i_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flby_i_plus][i, j] + m[:if12_eq_flby_i_plus][i, j] >=
        2 * m[:if1_eq_flby_i_plus][i, j]
    )

    # 4 flbyj + box_y_axis_cover_j == flbyi
    @variable(m, if11_eq_flby_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if12_eq_flby_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j], Bin)
    @variable(m, if1_eq_flby_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j] >= 0, Bin)

    @constraint(m, if_eq_flby_11_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if11_eq_flby_j_plus][i, j] >=
        m[:flby_front_left_bot_y][i] - (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) + 0.0001
    )
    @constraint(m, if_eq_flby_12_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if11_eq_flby_j_plus][i, j]) >=
        (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - m[:flby_front_left_bot_y][i] - 0.0001
    )
    @constraint(m, if_eq_flby_13_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * m[:if12_eq_flby_j_plus][i, j] >=
        (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - m[:flby_front_left_bot_y][i] + 0.0001
    )
    @constraint(m, if_eq_flby_14_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        W * (1 - m[:if12_eq_flby_j_plus][i, j]) >=
        m[:flby_front_left_bot_y][i] - (m[:flby_front_left_bot_y][j] + m[:box_y_axis_cover][j]) - 0.0001
    )

    @constraint(m, if1_eq_flby_cons_j_plus[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        m[:if11_eq_flby_j_plus][i, j] + m[:if12_eq_flby_j_plus][i, j] >=
        2 * m[:if1_eq_flby_j_plus][i, j]
    )

    @expression(m, guillotine_equalities_objective[i=1:length(boxes)-1, j=1:length(boxes)-1; i!=j],
        #X
        (1 - m[:if1_eq_flbx][i, j]) +
        (1 - m[:if1_eq_flbx_i_plus][i, j]) +
        (1 - m[:if1_eq_flbx_j_plus][i, j]) +
        (1 - m[:if1_eq_flbx_ij_plus][i, j]) +

        #Y
        (1 - m[:if1_eq_flby][i, j]) +
        (1 - m[:if1_eq_flby_i_plus][i, j]) +
        (1 - m[:if1_eq_flby_j_plus][i, j]) +
        (1 - m[:if1_eq_flby_ij_plus][i, j])
    )

    return m
end

end


text:
import .LibCImPlot: GetPlotPos, GetPlotSize # ,PlotToPixels

function SetNextPlotTicksX(values::Vector{<:Real}, n_ticks::Integer;
                           labels::Vector{String} = [""], show_default::Bool = false)

    eltype(values) !== Float64 && (values = Float64.(values))

    LibCImPlot.SetNextPlotTicksXdoublePtr(values, Cint(n_ticks), labels, show_default)
end

function SetNextPlotTicksX(x_min, x_max, n_ticks::Integer;
                           labels::Vector{String} = [""], show_default::Bool = false)

    typeof(x_min) !== Float64 && (x_min = Float64(x_min))
    typeof(x_max) !== Float64 && (x_max = Float64(x_max))

    LibCImPlot.SetNextPlotTicksXdouble(x_min, x_max, Cint(n_ticks), labels, show_default)
end

function SetNextPlotTicksY(values::Vector{<:Real}, n_ticks::Integer;
                           labels::Vector{String} = [""], show_default::Bool = false,
                           y_axis::Integer = 0)

    eltype(values) !== Float64 && (values = Float64.(values))

    LibCImPlot.SetNextPlotTicksYdoublePtr(values, Cint(n_ticks), labels, show_default,
                                          Cint(y_axis))
end

function SetNextPlotTicksY(y_min, y_max, n_ticks::Integer;
                           labels::Vector{String} = [""], show_default::Bool = false,
                           y_axis::Integer = 0)

    typeof(y_min) !== Float64 && (y_min = Float64(y_min))
    typeof(y_max) !== Float64 && (y_max = Float64(y_max))

    LibCImPlot.SetNextPlotTicksYdouble(y_min, y_max, Cint(n_ticks), labels, show_default,
                                       Cint(y_axis))
end

function GetPlotPos()
    out = Ref{ImVec2}()
    LibCImPlot.GetPlotPos(out)
    return out[]
end

function GetPlotSize()
    out = Ref{ImVec2}()
    LibCImPlot.GetPlotSize(out)
    return out[]
end

#=
function PlotToPixels(plt::ImPlotPoint, y_axis::Integer = -1)
    out = Ref{ImVec2}()
    LibCImPlot.PlotToPixels(out, plt, Cint(y_axis))
    return out[]
end
=#
LibCImPlot.SetLegendLocation(location, orientation) = LibCImPlot.SetLegendLocation(location, orientation, false)
LibCImPlot.SetLegendLocation(location) = LibCImPlot.SetLegendLocation(location, LibCImPlot.ImPlotOrientation_Vertical, false)

LibCImPlot.SetNextPlotLimits(xmin, xmax, ymin, ymax) = LibCImPlot.SetNextPlotLimits(xmin, xmax, ymin, ymax, ImGuiCond_Once)


text:
using NCDatasets
using Test

fname = tempname()

Dataset(fname,"c") do ds
    # Dimensions

    ds.dim["xi_rho"] = 6
    ds.dim["eta_rho"] = 6
    ds.dim["N"] = 30

    # Declare variables

    nclon_rho = defVar(ds,"lon_rho", Float64, ("xi_rho", "eta_rho"))
    nclon_rho.attrib["long_name"] = "longitude of RHO-points"
    nclon_rho.attrib["units"] = "degree_east"
    nclon_rho.attrib["field"] = "lon_rho, scalar"

    nclat_rho = defVar(ds,"lat_rho", Float64, ("xi_rho", "eta_rho"))
    nclat_rho.attrib["long_name"] = "latitude of RHO-points"
    nclat_rho.attrib["units"] = "degree_north"
    nclat_rho.attrib["field"] = "lat_rho, scalar"

end

ds = Dataset(fname,"a")
@test ds.dim["N"] == 30
NCDatasets.renameDim(ds,"N","NNN")
@test ds.dim["NNN"] == 30

@test "lon_rho" in keys(ds)
renameVar(ds,"lon_rho","longitude_rho")
@test "longitude_rho" in keys(ds)
close(ds)


text:
import Pkg
Pkg.activate(".")

include("src/Exchequer.jl")

using ArgParse
using .Exchequer

function parse()

    s = ArgParseSettings()

    @add_arg_table! s begin
        "--currency"
            help = "choose which currency to report in"
            arg_type = String
            default = "GBP"
            range_tester = x -> x in ["EUR", "GBP"]
    end

    return parse_args(s)

end

function main()

    # parse arguments
    args = parse()
    println(args)

    # update the downloaded files
    Exchequer.move_downloads()

    # load the dataframes
    df_expenses = Exchequer.load_expenses()
    df_incomes = Exchequer.load_incomes()
    df_fxs = Exchequer.load_fxs()
    df_net = Exchequer.load_net()

    # exchange the currencies
    Exchequer.exchange!(df_expenses, args["currency"], df_fxs)
    Exchequer.exchange!(df_incomes, args["currency"], df_fxs)

    # plot away
    Exchequer.plot_categories(df_expenses)
    Exchequer.plot_monthly(df_expenses, df_incomes)
    Exchequer.plot_weekly(df_expenses)
    Exchequer.plot_yearly(df_expenses, df_incomes)
    Exchequer.plot_net(df_net)

    return nothing

end

main()

nothing

max_stars_repo_name: mzgubic/exchequer

max_issues_repo_name: mzgubic/exchequer

max_forks_repo_name: mzgubic/exchequer


text:
using OrdinaryDiffEq, Flux, Optim, Random, Plots
using Zygote
using ForwardDiff
using LinearAlgebra
using Statistics
using ProgressBars, Printf
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load
using DiffEqSensitivity

Random.seed!(1234);

###################################
# Argments
is_restart = true;
n_epoch = 1000000;
n_plot = 10;
opt = ADAMW(0.001, (0.9, 0.999), 1.f-6);
datasize = 40;
tstep = 0.1;
n_exp_train = 70;
n_exp_val = 30;
n_exp = n_exp_train + n_exp_val;
noise = 1.f-2;
ns = 9;
nr = 15;
alg = Tsit5();
atol = 1e-5;
rtol = 1e-2;

np = nr * (2 * ns + 1)
p = (rand(Float32, np) .- 0.5) * 2 * sqrt(6 / (ns + nr))

lb = 1.f-5;
ub = 1.f2;
####################################

function p2vec(p)
    w_b = @view p[1:nr]
    w_in = reshape(@view(p[nr * (ns + 1) + 1:nr * (2 * ns + 1)]), ns, nr)

    w_out = reshape(@view(p[nr + 1:nr * (ns + 1)]), ns, nr)
    w_out[[1, 4, 7], :] .= 0

    w_out = @. -w_in * abs(w_out);
    w_in = clamp.(w_in, 0.f0, 4.f0);

    return w_in, w_b, w_out
end

function display_p(p)
    w_in, w_b, w_out = p2vec(p)
    println("species (column) reaction (row)")
    println("w_in")
    show(stdout, "text/plain", round.(w_in', digits=3))

    println("\nw_b")
    show(stdout, "text/plain", round.(exp.(w_b'), digits=3))

    println("\nw_out")
    show(stdout, "text/plain", round.(w_out', digits=3))
    println("\n")
end
display_p(p)

function display_grad(grad)

    grad_w_out = reshape(grad[nr * (ns + 1) + 1:nr * (2 * ns + 1)], ns, nr);

    println("\ngrad w_out")
    show(stdout, "text/plain", round.(grad_w_out', digits=6))
    println("\n")

end

function trueODEfunc(dydt, y, k, t)
    #= Species:
        1: DNA_A
        2: mRNA_A
        3: A
        4: DNA_B
        5: mRNA_B
        6: B
        7: DNA_C
        8: mRNA_C
        9: C

    Reactions: i = [1,2,3]
        i+0: DNA_i -> DNA_i + mRNA_i - Transcription of mRNA
        i+1: mRNA_i -> mRNA_i + i - Translation of proteins
        i+2: mRNA_i -> 0 - mRNA decay
        i+3: i -> 0 - Dacay of proteins

        i+12: mRNA_i + (i+1) -> (i+1) Cyclic - Regulation of mRNA =#

    R = zeros(15)

    R[1] = k[1] * y[1]
    R[2] = k[2] * y[2]
    R[3] = k[3] * y[2]
    R[4] = k[4] * y[3]

    R[5] = k[5] * y[4]
    R[6] = k[6] * y[5]
    R[7] = k[7] * y[5]
    R[8] = k[8] * y[6]

    R[9] = k[9] * y[7]
    R[10] = k[10] * y[8]
    R[11] = k[11] * y[8]
    R[12] = k[12] * y[9]

    R[13] = k[13] * y[8] * y[3]
    R[14] = k[14] * y[5] * y[9]
    R[15] = k[15] * y[2] * y[6]

    dydt[1] = 0
    dydt[4] = 0
    dydt[7] = 0

    dydt[2] = R[1] - R[3] - R[15]
    dydt[5] = R[5] - R[7] - R[14]
    dydt[8] = R[9] - R[11] - R[13]

    dydt[3] = R[2] - R[4]
    dydt[6] = R[6] - R[8]
    dydt[9] = R[10] - R[12]
end

# Generate data sets
u0_list = rand(Float32, (n_exp, ns));
# @. u0_list[:, [1, 4, 7]] .= 1.f0;
# @. u0_list[:, [3, 5, 8]] .*= 0.f0;

tspan = Float32[0.0, datasize * tstep];
tsteps = range(tspan[1], tspan[2], length=datasize);
k = [1.8, 2.1, 1.3, 1.5, 2.2, 2, 2, 2.5, 3.2, 3, 2.3, 2.5, 6, 4, 3];
ode_data_list = zeros(Float32, (n_exp, ns, datasize));
std_list = [];

u0 = u0_list[1, :];
prob_trueode = ODEProblem(trueODEfunc, u0, tspan, k);
ode_data = Array(solve(prob_trueode, alg, saveat=tsteps));

function max_min(ode_data)
    return maximum(ode_data, dims=2) .- minimum(ode_data, dims=2) .+ lb
end

for i = 1:n_exp
    u0 = u0_list[i, :]
    prob_trueode = ODEProblem(trueODEfunc, u0, tspan, k)
    ode_data = Array(solve(prob_trueode, alg, saveat=tsteps))

    ode_data += randn(size(ode_data)) .* ode_data .* noise
    ode_data_list[i, :, :] = ode_data

    push!(std_list, max_min(ode_data))
end

y_std = maximum(hcat(std_list...), dims=2) .+ lb;
show(stdout, "text/plain", round.(y_std', digits=8))

function crnn(du, u, p, t)
    w_in_x = w_in' * @. log(clamp(u, lb, ub))
    du .= w_out * (@. exp(w_in_x + w_b))
end

u0 = u0_list[1, :]
prob = ODEProblem(crnn, u0, tspan, saveat=tsteps, atol=atol, rtol=rtol)

# sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
sense = ForwardDiffSensitivity(convert_tspan=false)
batch = datasize
function predict_neuralode(u0, p)
    global w_in, w_b, w_out = p2vec(p)
    prob_ = remake(prob, u0=u0, p=p)
    pred = clamp.(Array(solve(prob_, alg, saveat=tsteps[1:batch], sensalg=sense)), lb, ub)
    return pred
end
predict_neuralode(u0, p)

function loss_neuralode(p, i_exp)
    pred = predict_neuralode(u0_list[i_exp, :], p)
    ode_data = clamp.(ode_data_list[i_exp, :, 1:batch], lb, ub)
    loss = mae(ode_data, pred)
    return loss
end

# Callback function to observe training
cbi = function (p, i_exp)
    ode_data = ode_data_list[i_exp, :, :]
    pred = predict_neuralode(u0_list[i_exp, :], p)
    list_plt = []
    for i = 1:ns
        plt = scatter(
            tsteps,
            ode_data[i, :],
            markercolor=:transparent,
            title=string(i),
            label=string("data_", i),
        )
        plot!(plt, tsteps, pred[i, :], label=string("pred_", i))
        push!(list_plt, plt)
    end
    plt_all = plot(list_plt..., legend=false)
    png(plt_all, string("figs/i_exp_", i_exp))
    return false
end

list_loss_train = []
list_loss_val = []
list_grad = []
iter = 1
cb = function (p, loss_train, loss_val, g_norm)
    global list_loss_train, list_loss_val, list_grad, iter
    push!(list_loss_train, loss_train)
    push!(list_loss_val, loss_val)
    push!(list_grad, g_norm)

    if iter % n_plot == 0
        display_p(p)

        @printf("min loss train %.4e val %.4e\n", minimum(list_loss_train), minimum(list_loss_val))

        list_exp = randperm(5)[1:1]
        println("update plot for ", list_exp)
        for i_exp in list_exp
            cbi(p, i_exp)
        end

        plt_loss = plot(list_loss_train, xscale=:log10, yscale=:log10, framestyle=:box, label="Training");
        plot!(plt_loss, list_loss_val, label="Validation");
        plot!(xlabel="Epoch", ylabel="Loss")
        png(plt_loss, "figs/loss");

        plt_grad = plot(list_grad, xscale=:log10, yscale=:log10, label="grad_norm")
        png(plt_grad, "figs/grad")

        @save "./checkpoint/mymodel.bson" p opt list_loss_train list_loss_val list_grad iter
    end
    iter += 1
end

if is_restart
    @load "./checkpoint/mymodel.bson" p opt list_loss_train list_loss_val list_grad iter
    iter += 1
    # opt = ADAMW(0.0001, (0.9, 0.999), 1.f-8);
end

i_exp = 1;
epochs = ProgressBar(iter:n_epoch);
loss_epoch = zeros(Float32, n_exp);
grad_norm = zeros(Float32, n_exp);
for epoch in epochs
    global p
    for i_exp in randperm(n_exp)
        global batch = rand(2:datasize)
        grad = ForwardDiff.gradient(x -> loss_neuralode(x, i_exp), p);
        grad_norm[i_exp] = norm(grad, 2)
        update!(opt, p, grad)
    end
    batch = datasize
    for i_exp in 1:n_exp_train + n_exp_val
        loss_epoch[i_exp] = loss_neuralode(p, i_exp)
    end
    loss_train = mean(loss_epoch[1:n_exp_train]);
    loss_val = mean(loss_epoch[n_exp_train + 1:end]);
    set_description(epochs, string(@sprintf("Loss train %.4e val %.4e", loss_train, loss_val)))
    cb(p, loss_train, loss_val, mean(grad_norm));
end


text:
# remaining life 4.5.2 MAWP Approach
using Gadfly
using DataFrames
using Cairo

# Data variables
pipe_code = "B31.3"
annex2c_tmin_category = "Straight Pipes Subject To Internal Pressure" # ["Cylindrical Shell","Spherical Shell","Hemispherical Head","Elliptical Head","Torispherical Head","Conical Shell","Toriconical Head","Conical Transition","Nozzles Connections in Shells",
# "Junction Reinforcement Requirements at Conical Transitions","Tubesheets","Flat head to cylinder connections","Bolted Flanges","Straight Pipes Subject To Internal Pressure","Boiler Tubes","Pipe Bends Subject To Internal Pressure",
# "MAWP for External Pressure","Branch Connections","API 650 Storage Tanks"]
equipment_group = "piping" # "vessel", "tank"
flaw_location = "external" # "External","Internal"
FCA_string = "internal"
metal_loss_categorization = "LTA" # "LTA" or "Groove-Like Flaw"
units = "lbs-in-psi" # "lbs-in-psi" or "nmm-mm-mpa"
remaining_life = true
tmm_forcing = true
tmm = .222
tnom = .365 # nominal or furnished thickness of the component adjusted for mill undertolerance as applicable.
trd = .365 # uniform thickness away from the local metal loss location established by thickness measurements at the time of the assessment.
FCAml = 0.05 # Future Corrosion Allowance applied to the region of metal loss.
FCA = 0.05 # Future Corrosion Allowance applied to the region away from the metal loss (see Annex 2C, paragraph 2C.2.8).
LOSS = 0.0 #the amount of uniform metal loss away from the local metal loss location at the time of the assessment.
Do = 10.75 # Outside Diameter
D = Do - 2*(tnom) # Inside Dia.
P = 740.0 # internal design pressure.
F = 0.0 # B31.8 design factor
T = 0.0 # temperature factor B31.8
S = 20000.0 # allowable stress.
E = 1.0 # weld joint efficiency or quality factor from the original construction code, if unknown use 0.7.
MA = 0.0 # mechanical allowances (thread or groove depth); for threaded components, the nominal thread depth (dimension h of ASME B.1.20.1) shall apply.
Yb31 = 0.4 # coefficient from ASME B31 Piping codes used for determining the pipe wall thickness, the coefficient can be determined from the following table that is valid for tmin < Do / 6 Annex 2C .
t = trd - LOSS - FCA # thickness of the shell or pipe adjusted for mill tolerance, LOSS and FCA , or cylinder thickness at a conical transition for a junction reinforcement calculation adjusted for mill tolerance, LOSS and FCA , as applicable.
tsl = 0.0 # supplemental thickness for mechanical loads other than pressure that result in longitudinal stress; this thickness is usually obtained from the results of a weight case in a stress analysis of the piping system (see paragraph 2C.2.7).
spacings = 0.5 # spacings determine by visual inspection to adequately ccategorizse the corrosion -----------+ may add to CTP_Grid function for plotting purposes
# Flaw dimensions
s = 6.0 # longitudinal extent or length of the region of local metal loss based on future corroded thickness,
c = 2.0  # circumferential extent or length of the region of local metal loss (see Figure 5.2 and Figure 5.10), based on future corroded thickness, tc .
Ec = 1.0 # circumferential weld joint efficiency. note if damage on weld see # 2C.2.5 Treatment of Weld and Riveted Joint Efficiency, and Ligament Efficiency
El = 1.0 # longitudinal weld joint efficiency. note if damage on weld see # 2C.2.5 Treatment of Weld and Riveted Joint Efficiency, and Ligament Efficiency
RSFa = 0.9 # remaining strength factor - consult API 579 is go lower than 0.9

# For all assessments determine far enough from structural discontinuity
# Flaw-To-Major Structural Discontinuity Spacing
L1msd = [12.0] # distance to the nearest major structural discontinuity.
L2msd = [12.0] # distance to the nearest major structural discontinuity.
L3msd = [12.0] # distance to the nearest major structural discontinuity.
L4msd = [12.0] # distance to the nearest major structural discontinuity.
L5msd = [12.0] # distance to the nearest major structural discontinuity.
Lmsd = minimum([L1msd,L2msd,L3msd,L4msd,L5msd])[1]
if (Lmsd >= (1.8*(sqrt(D*(trd - LOSS - FCA)))))
    print("Satisfied - Flaw is located far enough from structural discontinuity\n")
    lmsd_satisfied = 1
else
    print("Not satisfied - Flaw is too close to a structural discontinuity - Conduct a level 3 assessment\n")
    lmsd_satisfied = 0
end

# For all assessments - determine the inspection data grid
M1 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M2 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.100, 0.220, 0.280, 0.250, 0.240, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M3 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.215, 0.255, 0.215, 0.145, 0.275, 0.170, 0.240, 0.250, 0.250, 0.280, 0.290, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M4 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.170, 0.270, 0.190, 0.190, 0.285, 0.250, 0.225, 0.275, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M5 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
M6 = [0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300, 0.300]
CTPGrid = hcat(M6,M5,M4,M3,M2,M1) # build in descending order
CTPGrid = rotl90(CTPGrid) # rotate to correct orientation

# Groove Like Flaw dimensions
gl = .05 # length of the Groove-Like Flaw based on future corroded condition.
gw = .4 # width of the Groove-Like Flaw based on future corroded condition.
gr = 0.1 # radius at the base of a Groove-Like Flaw based on future corroded condition.
β = 40.0 # see (Figure 5.4) :: orientation of the groove-like flaw with respect to the longitudinal axis or a parameter to compute an effective fracture toughness for a groove being evaluated as a crack-like flaw, as applicable.

# STEP 1 – Determine the metal loss of the component, tloss
tloss = tnom - tmm

# STEP 2 – Determine the MAWPr for a series of increasing time increments using an effective corrosion
# allowance and the nominal thickness in the computation. The effective corrosion allowance is determined
# as follows:
time_series = collect(1.0:1.0:50.0) # years
crate = 0.04661 # inches per year
Csrate = 0.04661 # long growth
Ccrate = 0.046615 # circ growth
CAe = tnom .- (crate .* time_series)

# STEP 3 – Using the results from STEP 2, determine the remaining life from a plot of the MAWP versus
# time. The time at which the MAWPr curve intersects the equipment design pressure or equipment
# MAWP is the remaining life of the component.
#crate = 0.01
MAWPr_out = zeros(size(time_series,1))
for i = 1:length(time_series)
        let c = c, s=s, crate=crate, tloss = tloss
    c = c + Ccrate * time_series[i]
    s = s + Csrate * time_series[i]
    tmm = CAe[i]
    out = Part5LTALevel1(CTPGrid; remaining_life=remaining_life,tmm_forcing=tmm_forcing, tmm=tmm, annex2c_tmin_category=annex2c_tmin_category, pipe_code=pipe_code,equipment_group=equipment_group, flaw_location=flaw_location, FCA_string=FCA_string, metal_loss_categorization=metal_loss_categorization, units=units, Lmsd=Lmsd,tnom=tnom,
        trd=trd, FCA=FCA, FCAml=FCAml, LOSS=LOSS, Do=Do, D=D, P=P, S=S, E=E, MA=MA, Yb31=Yb31, t=t,tsl=tsl, F=F, T=T, spacings=spacings, s=s, c=c, El=El, Ec=Ec, RSFa=RSFa, gl=gl, gw=gw, gr=gr,β=β)

MAWPr_out[i] = out[18,2]
end # let
end

# create dataframe for plotting
df = hcat(MAWPr_out,fill(crate,size(MAWPr_out,1)),time_series)
df = DataFrame(df)
names = ["MAWPr","Corrosion Rates (ipy)","Time"]
rename!(df, names)
# x axis label
P = 8139.00
x_label = join(["Time to Calculated P ",string(P)," psig - Years"],"")

# Plot data
p = plot(df, x=df.Time, y=df.MAWPr, color="Corrosion Rates (ipy)", Geom.line, Scale.y_continuous(format=:plain),
Guide.title("API 579 Remaining Life - tmm revised for specific corrosion rates - Years for MAWPr to reach equipment Design P at a specific corrosion rate (ipy)"), Guide.xlabel(x_label),Guide.ylabel("MAWPr"),
yintercept=[P],
Geom.hline(size=0.5mm, style=:solid, color="red"),
Guide.YTicks(ticks=collect(0:200:10000)),
Guide.XTicks(ticks=time_series))
draw(PNG("C:/Users/Andrew.Bannerman/OneDrive - Shell/Documents/remaining_life_static_cr_rate.png",12inch,9inch),p)

# Part 5 Remaining life method - note tmm defined vs CTP's use MAWP RL approach
# this is for varying corrosion rates
crate = collect(0:0.0025:0.03)
time = collect(1.0:1.0:50.0)
MAWPr = zeros(size(time,1))
remaining_wt = .4
Csrate = 0.005 # long growth
Ccrate = 0.005 # circ growth
rl_out = Array{Float64}(undef, 0, 3)
i=1
j = 10
for i = 1:length(crate)
    let remaining_wt = remaining_wt, c = c, s=s
        for j in 1:length(time)
    tmm = remaining_wt - (crate[i] * time[j])
    c = c + Ccrate * time[j]
    s = s + Csrate * time[j]
    out = Part5LTALevel1(CTPGrid; tmm_forcing=tmm_forcing, tmm=tmm, annex2c_tmin_category=annex2c_tmin_category, equipment_group=equipment_group, flaw_location=flaw_location, FCA_string=FCA_string, metal_loss_categorization=metal_loss_categorization, units=units, Lmsd=Lmsd,tnom=tnom,
        trd=trd, FCA=FCA, FCAml=FCAml, LOSS=LOSS, Do=Do, D=D, P=P, S=S, E=E, MA=MA, Yb31=Yb31, t=t,tsl=tsl, spacings=spacings, s=s, c=c, El=El, Ec=Ec, RSFa=RSFa, gl=gl, gw=gw, gr=gr,β=β)
MAWPr[j] = out[18,2]
print("this is iteration i",i,"this is iteration j",j)
if j == length(time)
    global final_out = hcat(fill(crate[i],size(MAWPr,1)),MAWPr,time)
    global rl_out = vcat(rl_out,final_out)
    end
end # end j loop

end # end let
end # end i loop

df = DataFrame(rl_out)
names = ["Corrosion Rates (ipy)","MAWPr","Time"]
rename!(df, names)

x_label = join(["Time to Design P ",string(P)," psig - Years"],"")

# Plot data
p = plot(df, x=df.Time, y=df.MAWPr, color="Corrosion Rates (ipy)", Geom.line, Scale.y_continuous(format=:plain),
Guide.title("API 579 Remaining Life - tmm revised for varying corrosion rates - Years for MAWPr to reach equipment Design P at varying corrosion rates (ipy)"), Guide.xlabel(x_label),Guide.ylabel("MAWPr"),
yintercept=[P],
Geom.hline(size=0.5mm, style=:solid, color="red"),
Guide.YTicks(ticks=collect(0:100:1500)),
Guide.XTicks(ticks=time))
draw(PNG("C:/Users/Andrew.Bannerman/OneDrive - Shell/Documents/remaining_life.png",12inch,9inch),p)


text:
abstract type MessagePassing end

adjlist(m::T) where {T<:MessagePassing} = m.adjlist
message(m::T) where {T<:MessagePassing} = error("not implement")
update(m::T) where {T<:MessagePassing} = error("not implement")

function propagate(mp::T; aggr::Symbol=:add, kwargs...) where {T<:MessagePassing}
    M, cluster = neighboring(mp; kwargs...)
    M = pool(aggr, cluster, M)
    upd_args = Dict{Symbol,AbstractArray}(:M=>M)
    haskey(kwargs, :X) && (upd_args[:X] = kwargs[:X])
    Y = update(mp; upd_args...)
    return Y
end

function neighboring(mp::T; kwargs...) where {T<:MessagePassing}
    adj = adjlist(mp)
    msg_args = getdata(kwargs, 1, adj[1])
    M = message(mp; msg_args...)
    apply_messages(mp, M, adj; kwargs...)
end

function getdata(d, i::Integer, ne)
    result = Dict{Symbol,AbstractArray}()
    if haskey(d, :X)
        result[:x_i] = view(d[:X], :, i)
        result[:x_j] = view(d[:X], :, ne)
    end
    if haskey(d, :E)
        result[:e_ij] = view(d[:E], :, i, ne)
    end
    result
end

function apply_messages(mp, M::AbstractMatrix{T}, adj::AbstractVector{<:AbstractVector},
                        F::Integer=size(M,1), N::Integer=length(adj), n::Integer=length(adj[1]);
                        kwargs...) where {T<:Real}
    if F > 50
        # TODO: Benchmark how many N or F need to switch from sequential to milti-thread version.
        return thread_apply_messages(mp, M, adj, F, N, n; kwargs...)
    else
        ne_N = mapreduce(length, +, adj)
        Y = Matrix{T}(undef, F, ne_N)
        Y[:, 1:n] = M
        cluster = Vector{Int}(undef, ne_N)
        cluster[1:n] .= 1
        j = n
        @inbounds for i = 2:N
            ne = adj[i]
            n = length(ne)
            msg_args = getdata(kwargs, i, ne)
            Y[:, j+1:j+n] = message(mp; msg_args...)
            cluster[j+1:j+n] .= i
            j += n
        end
        return Y, cluster
    end
end

function thread_apply_messages(mp, M::AbstractMatrix{T}, adj::AbstractVector{<:AbstractVector},
                               F::Integer=size(M,1), N::Integer=length(adj), n::Integer=length(adj[1]);
                               kwargs...) where {T<:Real}
    starts = cumsum(map(length, adj))
    ne_N = starts[end]
    Y = Matrix{T}(undef, F, ne_N)
    Y[:, 1:n] = M
    cluster = Vector{Int}(undef, ne_N)
    cluster[1:n] .= 1
    @inbounds Threads.@threads for i = 2:N
        j = starts[i-1]
        k = starts[i]
        msg_args = getdata(kwargs, i, adj[i])
        Y[:, j+1:k] = message(mp; msg_args...)
        cluster[j+1:k] .= i
    end
    Y, cluster
end


text:
# tensoradd.jl
#
# Method for adding one tensor to another according to the
# specified labels, thereby possibly having to permute the
# data.

# Simple method
# --------------
function tensoradd(A::StridedArray,labelsA,B::StridedArray,labelsB,outputlabels=labelsA)
    dims=size(A)
    T=promote_type(eltype(A),eltype(B))
    C=similar(A,T,dims[indexin(outputlabels,labelsA)])
    tensorcopy!(A,labelsA,C,outputlabels)
    tensoradd!(1,B,labelsB,1,C,outputlabels)
end

# In-place method
#-----------------
function tensoradd!(alpha::Number,A::StridedArray,labelsA,beta::Number,C::StridedArray,labelsC)
    NA=ndims(A)
    perm=indexin(labelsC,labelsA)
    length(perm) == NA || throw(LabelError("invalid label specification"))
    isperm(perm) || throw(LabelError("labels do not specify a valid permutation"))
    for i = 1:NA
        size(A,perm[i]) == size(C,i) || throw(DimensionMismatch("destination tensor of incorrect size"))
    end
    NA==0 && (C[1]=beta*C[1]+alpha*A[1]; return C)
    perm==collect(1:NA) && return (beta==0 ? scale!(copy!(C,A),alpha) : Base.LinAlg.axpy!(alpha,A,scale!(C,beta)))
    beta==0 && return scale!(tensorcopy_native!(A,C,perm),alpha)
    tensoradd_native!(alpha,A,beta,C,perm)
end

# Implementation
#-----------------
@eval @ngenerate N typeof(C) $PERMUTEGENERATE function tensoradd_native!{TA,TC,N}(alpha::Number,A::StridedArray{TA,N},beta::Number,C::StridedArray{TC,N},perm)
    @nexprs N d->(stridesA_{d} = stride(A,perm[d]))
    @nexprs N d->(stridesC_{d} = stride(C,d))
    @nexprs N d->(dims_{d} = size(C,d))

    if beta==0
        fill!(C,zero(TC))
    end

    startA = 1
    local Alinear::Array{TA,N}
    if isa(A, SubArray)
        startA = A.first_index
        Alinear = A.parent
    else
        Alinear = A
    end
    startC = 1
    local Clinear::Array{TC,N}
    if isa(C, SubArray)
        startC = C.first_index
        Clinear = C.parent
    else
        Clinear = C
    end

    if length(C)<=4*TBASELENGTH
        @stridedloops(N, i, dims, indA, startA, stridesA, indC, startC, stridesC, @inbounds Clinear[indC]=beta*Clinear[indC]+alpha*Alinear[indA])
    else
        @nexprs N d->(minstrides_{d} = min(stridesA_{d},stridesC_{d}))

        # build recursive stack
        depth=ceil(Integer, log2(length(C)/TBASELENGTH))+2 # 2 levels safety margin
        level=1 # level of recursion
        stackpos=zeros(Int,depth) # record pos of algorithm at the different recursion level
        stackpos[level]=0
        stackblength=zeros(Int,depth)
        stackblength[level]=length(C)
        @nexprs N d->begin
            stackbdims_{d} = zeros(Int,depth)
            stackbdims_{d}[level] = dims_{d}
        end
        stackbstartA=zeros(Int,depth)
        stackbstartA[level]=startA
        stackbstartC=zeros(Int,depth)
        stackbstartC[level]=startC

        stackdC=zeros(Int,depth)
        stackdA=zeros(Int,depth)
        stackdmax=zeros(Int,depth)
        stacknewdim=zeros(Int,depth)
        stackolddim=zeros(Int,depth)

        while level>0
            pos=stackpos[level]
            blength=stackblength[level]
            @nexprs N d->(bdims_{d} = stackbdims_{d}[level])
            bstartA=stackbstartA[level]
            bstartC=stackbstartC[level]

            if blength<=TBASELENGTH || level==depth # base case
                @stridedloops(N, i, bdims, indA, bstartA, stridesA, indC, bstartC, stridesC, @inbounds Clinear[indC]=beta*Clinear[indC]+alpha*Alinear[indA])
                level-=1
            elseif pos==0
                # find which dimension to divide
                dmax=0
                maxval=0
                newdim=0
                olddim=0
                dC=0
                dA=0
                @nexprs N d->begin
                    newmax=bdims_{d}*minstrides_{d}
                    if bdims_{d}>1 && newmax>maxval
                        dmax=d
                        olddim=bdims_{d}
                        newdim=olddim>>1
                        dC=stridesC_{d}
                        dA=stridesA_{d}
                        maxval=newmax
                    end
                end
                stackolddim[level]=olddim
                stacknewdim[level]=newdim
                stackdmax[level]=dmax
                stackdC[level]=dC
                stackdA[level]=dA

                stackpos[level+1]=0
                @nexprs N d->(stackbdims_{d}[level+1] = (d==dmax ? newdim : bdims_{d}))
                stackblength[level+1]=div(blength,olddim)*newdim
                stackbstartA[level+1]=bstartA
                stackbstartC[level+1]=bstartC

                stackpos[level]+=1
                level+=1
            elseif pos==1
                olddim=stackolddim[level]
                newdim=stacknewdim[level]
                dmax=stackdmax[level]
                dC=stackdC[level]
                dA=stackdA[level]

                stackpos[level+1]=0
                @nexprs N d->(stackbdims_{d}[level+1] = (d==dmax ? olddim-newdim : bdims_{d}))
                stackblength[level+1]=div(blength,olddim)*(olddim-newdim)
                stackbstartA[level+1]=bstartA+dA*newdim
                stackbstartC[level+1]=bstartC+dC*newdim

                stackpos[level]+=1
                level+=1
            else
                level-=1
            end
        end
    end
    return C
end


text:
# Tests for Profit Russell DEA Models
@testset "ProfitRussellDEAModel" begin

    # Test using Book data
    X = [2; 4; 8; 12; 6; 14; 14; 9.412];
    Y = [1; 5; 8; 9; 3; 7; 9; 2.353];
    W = [1; 1; 1; 1; 1; 1; 1; 1];
    P = [2; 2; 2; 2; 2; 2; 2; 2];

    # Profit Russell
    deaprofitrus = deaprofitrussell(X, Y, W, P)

    @test typeof(deaprofitrus) == ProfitRussellDEAModel

    @test nobs(deaprofitrus) == 8
    @test ninputs(deaprofitrus) == 1
    @test noutputs(deaprofitrus) == 1
    @test ismonetary(deaprofitrus) == false

    @test efficiency(deaprofitrus) == efficiency(deaprofitrus, :Economic)
    @test efficiency(deaprofitrus, :Economic)   ≈ [2; 0.25; 0.0; 1/12; 2/3; 0.285714; 0.142857; 1.349979] atol = 1e-5
    @test efficiency(deaprofitrus, :Technical)  ≈ [0.0; 0.0; 0.0; 0.0; 11/30; 0.276786; 0.071429; 0.552205] atol = 1e-5
    @test efficiency(deaprofitrus, :Allocative) ≈ [2.0; 0.25; 0.0; 1/12; 0.3; 0.008929; 0.071429; 0.797773] atol = 1e-5
    @test normfactor(deaprofitrus) ≈ [4.0; 8.0; 16.0; 24.0; 12.0; 28.0; 28.0; 9.412]

    # With monetary = true
    deaprofitrusmon = deaprofitrussell(X, Y, W, P, monetary = true)

    @test ismonetary(deaprofitrusmon) == true

    @test efficiency(deaprofitrusmon) ≈ efficiency(deaprofitrus) .* normfactor(deaprofitrus)
    @test efficiency(deaprofitrusmon, :Technical) ≈ efficiency(deaprofitrus, :Technical) .* normfactor(deaprofitrus)
    @test efficiency(deaprofitrusmon, :Allocative) ≈ efficiency(deaprofitrus, :Allocative) .* normfactor(deaprofitrus)

    # Print
    show(IOBuffer(), deaprofitrus)

    # Test errors
    @test_throws DimensionMismatch deaprofitrussell([1; 2 ; 3], [4; 5], [1; 1; 1], [1; 1]) #  Different number of observations
    @test_throws DimensionMismatch deaprofitrussell([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], [1; 1; 1]) # Different number of observation in input prices
    @test_throws DimensionMismatch deaprofitrussell([1; 2; 3], [4; 5; 6], [1; 1; 1], [1; 2; 3; 4]) # Different number of observation in output prices
    @test_throws DimensionMismatch deaprofitrussell([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [1; 1; 1]) # Different number of input prices and inputs
    @test_throws DimensionMismatch deaprofitrussell([1; 2; 3 ], [4; 5; 6], [1; 1; 1], [1 1 1; 2 2 2; 3 3 3]) # Different number of output prices and outputs

end


text:
@generated function Base.:(-)(v::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
    vtyp = vtype(W, T)
    instrs = "%res = fneg nsz arcp contract afn reassoc $vtyp %0\nret $vtyp %res"
    quote
        $(Expr(:meta, :inline))
        Vec(llvmcall($instrs, _Vec{$W,$T}, Tuple{_Vec{$W,$T}}, data(v)))
    end
end
# for T ∈ [Float32,Float64]
#     W = 2
#     typ = LLVM_TYPES[T]
#     while W ≤ pick_vector_width(T)
#         vtyp = "<$W x $typ>"
#         instrs = "%res = fneg fast $vtyp %0\nret $vtyp %res"
#         @eval @inline Base.:(-)(v::Vec{$W,$T}) = Vec(llvmcall($instrs, _Vec{$W,$T}, Tuple{_Vec{$W,$T}}, data(v)))
#         W += W
#     end
# end

# @inline Base.:(-)(v::VecUnroll) = fmap(-, v)
# @inline Base.:(-)(v::VecTile) = fmap(-, v)
# @inline Base.:(-)(v::Vec{<:Any,<:FloatingTypes}) = -v
@inline Base.:(-)(v::Vec{1,<:FloatingTypes}) = -first(v)
@inline Base.:(-)(v::Vec{<:Any,<:NativeTypes}) = zero(v) - v

@inline Base.inv(v::Vec) = vdiv(one(v), v)

@inline Base.:(~)(v::AbstractSIMD{W,T}) where {W,T<:IntegerTypesHW} = v ⊻ vbroadcast(Val(W), -1 % T)


text:
type OCCAFunction
    dev::OCCA.Device;
    kernel::OCCA.Kernel;
end


text:
#MAIN ENTRY POINT
#final columns for the decompressed file

#merges crsp and compustat
# note refresh merge will be automatically run if any of the other two are run
function makeuniv(;
  refreshcomp::Bool = true,
  refreshcrsp::Bool = true,
  refreshmerge::Bool = true,
  datapath::String=DATA_PATH,
  univname::String = UNIV_NAME,
  injlsstream::Function = IN_JLS_STREAM,
  outjlsstream::Function = OUT_JLS_STREAM,
  validatemerged::Bool = true)::DataFrame

  local crsp::DataFrame
  local comp::DataFrame
  local univ::DataFrame

  #as of now, refreshmerged must be true if any of the other flags are true
  refreshmerge = refreshmerge || refreshcrsp || refreshcomp || validatemerged

  #makes a serialization object so we don't have to keep parsing the CSV
  if refreshmerge || (!isfile("$datapath\\$univname.jls.lz4"))
    crsp = prepcrsp(refreshcrsp=refreshcrsp)
    comp = prepcomp(refreshcomp=refreshcomp)

    univ = mergecrspcomp!(crsp, comp, validatemerged=validatemerged)
    outjlsstream("$datapath\\$univname.jls.lz4") do s
      serialize(s, univ)
    end

    #NOTE: FINAL VERSION
    #preprocesscomp(compa, compq)
    #serialize("$datapath\\$compname.jls", comp)
  else
    injlsstream("$datapath\\$univname.jls.lz4") do s
      univ = deserialize(s) #NOTE: DEBUG ONLY
    end
  end

  @info "CRSP and Compustat merged in file $(univname).jls.lz4."

  return univ
end

function mergecrspcomp!(crsp::DataFrame, comp::DataFrame;
    testoutput::Bool = TEST_OUTPUT,
    retainedcolsuniv::Vector{Symbol} = RETAINED_COLS_UNIV,
    validatemerged::Bool=true)::DataFrame

  local univ::DataFrame
  local scomp::SubDataFrame

  local Ncrsp::Int = size(crsp,1)
  #create a record id to form the dictionary
  comp.rid = Vector{MInt}(1:(size(comp,1)))
  comp.validinterval = ((begindate,enddate)->
    (begindate+Day(1)):Day(1):(enddate)).(comp.begindate, comp.enddate)

  crsp.rid = Vector{MInt}(undef, Ncrsp)

  #acquire the appropriate record id
  @time for scrsp::SubDataFrame ∈ groupby(crsp, :permno)
    scomp = view(comp, comp.lpermno .== scrsp[1, :permno], :)

    for rscrsp ∈ eachrow(scrsp)
      #ctr::Int = 0
      local ridaddress::NInt
      local rscrspdate::Date = rscrsp.date

      ridaddress = findfirst(
        interval->rscrspdate ∈ interval, scomp.validinterval)

      (!isnothing(ridaddress)) && (rscrsp.rid = scomp.rid[ridaddress])
    end
  end

  println("sum crsp rid: $(sum(skipmissing(crsp.rid)))")
  println("Matching complete: $((Ncrsp - sum((ismissing).(crsp.rid)))/Ncrsp)% matched.")
  #println("crsp records: $(Ncrsp)")

  #mNOTE: main join
  univ = join(crsp, comp, on=:rid, kind=:inner)

  if testoutput
    d = Date(2015,12,31)
    testoutput && CSV.write("output\\univ_$d.csv", univ[(cd->cd==d).(univ.date), :])
  end

  validatemerged && testmerged(univ)
  println("Total records: $(size(univ))")

  select!(univ, retainedcolsuniv)

  return univ
end

#do a bunch of integrity checks on the merged dataset
function testmerged(univ::DataFrame)

  for r ∈ eachrow(univ)
    #check date is valid
    d::Date = r.date
    @assert d ≥ r.linkeffdate
    @assert d ≤ r.linkenddate
    @assert d ∈ r.validinterval
    @assert d > r.adateq
    @assert d ≤ r.Nadateq
    @assert d ≤ lastdayofmonth(firstdayofmonth(r.adateq) + Month(6))

    @assert r.lpermno == r.permno
  end

  sort!(univ, [:permno, :adateq, :date])

  for suniv ∈ groupby(univ, :permno)
    maxfiscal::YearQuarter = YearQuarter(1900,1)
    maxadate::Date = Date(1900,1,1)

    for r ∈ eachrow(suniv)
      @assert r.date > maxadate
      maxadate = r.date

      @assert maxfiscal ≤ r.fyearqtr
      (maxfiscal ≠ r.fyearqtr) && (maxfiscal=r.fyearqtr)
    end

  end
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
mutable struct ScoredSchedulerEntry
    capacity::Int
    fuel::Int
    bantime::Int
    banremaining::Int
    weight::Int          # bantime multiplier, low = good
end

isbanned(d::ScoredSchedulerEntry)::Bool = d.banremaining > 0

"""
A Rewrite Scheduler that implements exponential rule backoff.
For each rewrite, there exists a configurable initial match limit.
If a rewrite search yield more than this limit, then we ban this rule
for number of iterations, double its limit, and double the time it
will be banned next time.

This seems effective at preventing explosive rules like
associativity from taking an unfair amount of resources.
"""
struct ScoredScheduler <: AbstractScheduler
    data::Dict{Rule, ScoredSchedulerEntry}
    G::EGraph
    theory::Vector{<:Rule}
end

shouldskip(s::ScoredScheduler, r::Rule)::Bool = s.data[r].banremaining > 0

function exprsize(e)
    if !(e isa Expr)
        return 1
    end

    start = Meta.isexpr(e, :call) ? 2 : 1

    c = 1 + length(e.args[start:end])
    for a ∈ e.args[start:end]
        c += exprsize(a)
    end

    return c
end

function ScoredScheduler(g::EGraph, theory::Vector{<:Rule})
    ScoredScheduler(g, theory, 8, 2, exprsize)
end

function ScoredScheduler(G::EGraph, theory::Vector{<:Rule}, fuel::Int, bantime::Int, complexity::Function)
    gsize = length(G.uf)
    data = Dict{Rule, ScoredSchedulerEntry}()

    # These numbers seem to fit
    for rule ∈ theory
        (l, r) = rule.left, rule.right
        l = l |> cleanast_rec |> remove_assertions |> unquote_sym
        r = r |> cleanast_rec |> remove_assertions |> unquote_sym

        cl = complexity(l)
        cr = complexity(r)
        # println("$rule HAS SCORE $((cl, cr))")
        if cl > cr
            w = 1   # reduces complexity
            fuel = Inf
        elseif cr > cl
            w = 3   # augments complexity
        else
            w = 2   # complexity is equal
        end
        # println(w)
        data[rule] = ScoredSchedulerEntry(fuel, fuel, bantime, 0, w)
    end

    return ScoredScheduler(data, G, theory)
end

# can saturate if there's no banned rule
cansaturate(s::ScoredScheduler)::Bool = all(kv -> !isbanned(last(kv)), s.data)

function readstep!(s::ScoredScheduler)
    for rule ∈ s.theory
        rd = s.data[rule]
        if rd.banremaining > 0
            rd.banremaining -= 1

            if rd.banremaining == 0
                rd.bantime *= 1 + rd.weight
                rd.capacity *= (4 - rd.weight)
                rd.fuel = rd.capacity
                # @info "unbanning rule" rule rd.weight rd.bantime rd.capacity
            end
        end
    end
end

function writestep!(s::ScoredScheduler, rule::Rule)
    rd = s.data[rule]

    # decrement fuel, ban rule if fuel is empty
    rd.fuel -= 1
    if rd.fuel == 0
        # @info "banning rule!" rule rd.weight rd.bantime rd.capacity
        rd.banremaining = rd.bantime
    end
end


text:
using BinaryBuilder

# Collection of sources required to build Nettle
name = "Nettle"
version = v"3.4"
sources = [
    "https://ftp.gnu.org/gnu/nettle/nettle-3.4.tar.gz" =>
    "ae7a42df026550b85daca8389b6a60ba6313b0567f374392e54918588a411e94",
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/nettle-*/
./configure --prefix=$prefix --host=$target
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 = [
    Windows(:i686),
    Windows(:x86_64),
    Linux(:i686, :glibc),
    Linux(:x86_64, :glibc),
    Linux(:aarch64, :glibc),
    Linux(:armv7l, :glibc),
    Linux(:powerpc64le, :glibc),
    MacOS(),
    FreeBSD(:x86_64),
]

# The products that we will ensure are always built
products(prefix) = [
    LibraryProduct(prefix, "libnettle", :libnettle),
    ExecutableProduct(prefix, "nettle-hash", :nettle_hash)
]

# Dependencies that must be installed before this package can be built
dependencies = [
]

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


text:
using Dates, Test
using TrustPilotCalc
using TrustPilotCalc: TPReview, TPScore

printstyled("Running tests:\n", color=:blue)
@testset "Testing base use cases" begin
let
    five_star_review = TPReview("seed", 5.0, Dates.now() - Dates.Year(2))
    one_star_review = TPReview("seed", 1.0, Dates.now())

    # 20x 5star reviews two years ago = trustscore 4.5.
    @test getTrustScore(fill(five_star_review, 20)) == TPScore(4.6, 4.5, five_star_review.created_at)
    # as above + 1-star review today
    @test getTrustScore([fill(five_star_review, 20); one_star_review]) == TPScore(3.3, 3.5, one_star_review.created_at)
    # single 1-star review today
    @test getTrustScore([one_star_review]) == TPScore(3.2, 3.0, one_star_review.created_at)
    # single 5-star review today
    @test getTrustScore([five_star_review]) == TPScore(3.7, 3.5, five_star_review.created_at)

end
end


text:
"""
Compute the planner's allocation by solving Bellman
equation.
"""
struct RecursiveAllocation{TP <: Model, TI <: Integer,
                           TVg <: AbstractVector, TVv <: AbstractVector,
                           TVp <: AbstractArray}
    model::TP
    mc::MarkovChain
    S::TI
    T::BellmanEquation
    μgrid::TVg
    xgrid::TVg
    Vf::TVv
    policies::TVp
end

"""
Initializes the {{ class_word }} from the calibration `Model`
"""
function RecursiveAllocation(model::Model, μgrid::AbstractArray)
    mc = MarkovChain(model.Π)
    G = model.G
    S = size(model.Π, 1) # Number of states
    # Now find the first best allocation
    Vf, policies, T, xgrid = solve_time1_bellman(model, μgrid)
    T.time_0 = true      # Bellman equation now solves time 0 problem
    return RecursiveAllocation(model, mc, S, T, μgrid, xgrid, Vf, policies)
end

"""
Solve the time 1 Bellman equation for calibration `Model` and initial grid `μgrid0`
"""
function solve_time1_bellman{TF <: AbstractFloat}(model::Model{TF}, μgrid::AbstractArray)
    μgrid0 = μgrid
    S = size(model.Π, 1)
    # First get initial fit
    PP = SequentialAllocation(model)
    c = Matrix{TF}(length(μgrid), 2)
    n = Matrix{TF}(length(μgrid), 2)
    x = Matrix{TF}(length(μgrid), 2)
    V = Matrix{TF}(length(μgrid), 2)
    for (i, μ) in enumerate(μgrid0)
        c[i, :], n[i, :], x[i, :], V[i, :] = time1_value(PP, μ)
    end
    Vf = Vector{LinInterp}(2)
    cf = Vector{LinInterp}(2)
    nf = Vector{LinInterp}(2)
    xprimef = Array{LinInterp}(2, S)
    for s in 1:2
        cf[s] = LinInterp(x[:, s][end:-1:1], c[:, s][end:-1:1])
        nf[s] = LinInterp(x[:, s][end:-1:1], n[:, s][end:-1:1])
        Vf[s] = LinInterp(x[:, s][end:-1:1], V[:, s][end:-1:1])
        for sprime in 1:S
            xprimef[s, sprime] = LinInterp(x[:, s][end:-1:1], x[:, s][end:-1:1])
        end
    end
    policies = [cf, nf, xprimef]
    # Create xgrid
    xbar = [maximum(minimum(x, 1)), minimum(maximum(x, 1))]
    xgrid = linspace(xbar[1], xbar[2], length(μgrid0))
    # Now iterate on bellman equation
    T = BellmanEquation(model, xgrid, policies)
    diff = 1.0
    while diff > 1e-6
        if T.time_0 == false
            Vfnew, policies =
                fit_policy_function(PP,
                (i_x, x, s) -> get_policies_time1(T, i_x, x, s, Vf), xgrid)
        elseif T.time_0 == true
            Vfnew, policies =
                fit_policy_function(PP,
                (i_x, B_, s0) -> get_policies_time0(T, i_x, B_, s0, Vf), xgrid)
        else
            error("T.time_0 is $(T.time_0), which is invalid")
        end
        diff = 0.0
        for s in 1:S
            diff = max(diff, maximum(abs, (Vf[s].(xgrid)-Vfnew[s].(xgrid))./Vf[s].(xgrid)))
        end
        print("diff = $diff \n")
        Vf = Vfnew
    end
    # Store value function policies and Bellman Equations
    return Vf, policies, T, xgrid
end

"""
Fits the policy functions PF using the points `xgrid` using interpolation
"""
function fit_policy_function(PP::SequentialAllocation,
                            PF::Function, xgrid::AbstractArray)
    S = PP.S
    Vf = Vector{LinInterp}(S)
    cf = Vector{LinInterp}(S)
    nf = Vector{LinInterp}(S)
    xprimef = Array{LinInterp}(S, S)
    for s in 1:S
        PFvec = Array{typeof(PP.model).parameters[1]}(length(xgrid), 3+S)
        for (i_x, x) in enumerate(xgrid)
            PFvec[i_x, :] = PF(i_x, x, s)
        end
        Vf[s] = LinInterp(xgrid, PFvec[:, 1])
        cf[s] = LinInterp(xgrid, PFvec[:, 2])
        nf[s] = LinInterp(xgrid, PFvec[:, 3])
        for sprime in 1:S
            xprimef[s, sprime] = LinInterp(xgrid, PFvec[:, 3+sprime])
        end
    end
    return Vf, [cf, nf, xprimef]
end

"""
Finds the optimal allocation given initial government debt `B_` and state `s_0`
"""
function time0_allocation(pab::RecursiveAllocation,
                          B_::AbstractFloat, s0::Integer)
    xgrid = pab.xgrid
    if pab.T.time_0 == false
        z0 = get_policies_time1(pab.T, i_x, x, s, pab.Vf)
    elseif pab.T.time_0 == true
        z0 = get_policies_time0(pab.T, B_, s0, pab.Vf)
    else
        error("T.time_0 is $(T.time_0), which is invalid")
    end
    c0, n0, xprime0 = z0[2], z0[3], z0[4:end]
    return c0, n0, xprime0
end

"""
Simulates Ramsey plan for `T` periods
"""
function simulate(pab::RecursiveAllocation,
                B_::AbstractFloat, s_0::Integer, T::Integer,
                sHist::Vector=QuantEcon.simulate(mc, s_0, T))
    model, S, policies = pab.model, pab.S, pab.policies
    β, Π, Uc = model.β, model.Π, model.Uc
    cf, nf, xprimef = policies[1], policies[2], policies[3]
    TF = typeof(model).parameters[1]
    cHist = Vector{TF}(T)
    nHist = Vector{TF}(T)
    Bhist = Vector{TF}(T)
    ΤHist = Vector{TF}(T)
    μHist = Vector{TF}(T)
    RHist = Vector{TF}(T-1)
    # time 0
    cHist[1], nHist[1], xprime = time0_allocation(pab, B_, s_0)
    ΤHist[1] = Τ(pab.model, cHist[1], nHist[1])[s_0]
    Bhist[1] = B_
    μHist[1] = 0.0
    # time 1 onward
    for t in 2:T
        s, x = sHist[t], xprime[sHist[t]]
        n = nf[s](x)
        c = [cf[shat](x) for shat in 1:S]
        xprime = [xprimef[s, sprime](x) for sprime in 1:S]
        ΤHist[t] = Τ(pab.model, c, n)[s]
        u_c = Uc(c, n)
        Eu_c = dot(Π[sHist[t-1], :], u_c)
        μHist[t] = pab.Vf[s](x)
        RHist[t-1] = Uc(cHist[t-1], nHist[t-1]) / (β * Eu_c)
        cHist[t], nHist[t], Bhist[t] = c[s], n, x / u_c[s]
    end
    return cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist
end


text:
# Run tests on functions in:
# src/by_reference_doc/users.jl
# src/by_reference_doc/follow.jl
# src/by_reference_doc/personalization.jl

@testset verbose = true "GET-request endpoints for users" begin

    # Input arguments from composite types defined in src/types.jl
    playlist_id = SpPlaylistId()
    user_id = SpUserId()

    # Get current user's profile
    @test ~isempty(Spotify.user_get_current_profile()[1])

    # Get user's top items
    # Currently shows 403 error, implemented in personalization.jl
    # @test_broken - This will denote the test as Broken if the test continues to fail
    # and alerts the user via an Error if the test succeeds.
    @test_broken ~isempty(Spotify.top_tracks()[1])
    @test_broken ~isempty(Spotify.top_artists()[1])

    # Get user's profile
    @test ~isempty(Spotify.user_get_profile(user_id)[1])

    # Get followed artists for current user
    @test ~isempty(Spotify.follow_artists()[1])

    # Check if current user follows a given artist
    @test ~isempty(Spotify.follow_check("artist", "7fxBPUc2bTUgl7GLuqjajk")[1])

    # Check if given users follow a given playlist
    @test ~isempty(Spotify.follow_check_playlist(playlist_id, user_id)[1])

end