text:
include("src/SchoolOutbreak.jl")

using Distributions, Plots, InvertedIndices, LinearAlgebra, StatsBase, Random
gr(fontfamily="Helvetica",foreground_color_legend = nothing,background_color_legend = nothing, titlefontsize=9,grid=false, tick_direction=:out)

# setup
struct ImproperUniform<: Distribution{Univariate,Discrete} end
Distributions.cdf(d::ImproperUniform, x::Integer)=1.0
Distributions.pdf(d::ImproperUniform, x::Integer)=1.0

rcomdist=Logistic(113.63,10.15)#Uniform(0,365)
hhfoi=-Inf
Gammameansd(mean,sd)=Gamma(mean^2/sd^2,sd^2/mean)
#gtimedist=Gammameansd(3,1)
gtimedist=Gammameansd(1.7,1) # H3N2 serial interval from Vink et al 2014 (https://doi.org/10.1093/aje/kwu209)
N0=[1,1,1,1,1].*5 # normalizing constant to avoid correlation between β and γ
C0=[0,30,30,30,30]

sus_covlabels=[:Grade, :Vaccine_curr, :NPI_mask,:NPI_wash]
inf_covlabels=[:Grade, :Vaccine_curr, :NPI_mask,:NPI_wash,:duringbreak]

classes=repeat([20,30,40],inner=100)
sizes=fill(5,300)
nclassmat=repeat(classes,1,1)
nstudentsmat=repeat(sizes,1,1).*nclassmat
strata=vcat.(1:300,sortslices.(repeat.((x->reduce(hcat,vec(collect.(x)))).(Iterators.product.(Ref(1:1),range.(1,classes,step=1)).|>collect),1,sizes),dims=2))
strata=vcat.(strata,repeat.(Ref([1]),1,size.(strata,2)))
strata=reduce(hcat,strata)

parameters=SchoolOutbreak.makeparameters(ones(5),ones(5)/2,ones(5)/2,0.02,
    ones(size(sus_covlabels)),ones(size(inf_covlabels)),gtimedist,20,rcomdist,strata,nclassmat,nstudentsmat,N0,C0,Inf)
@time students=SchoolOutbreak.studentsfromdata(parameters,false,fill(typemax(Int16)|>Int,size(strata,2)),eachcol(strata),
    Ref(zeros(length(sus_covlabels))),Ref(zeros(length(inf_covlabels))),nclassmat,nstudentsmat,-Inf,1);
@time simstudents=SchoolOutbreak.Students.([filter(x->x.stratum[1]==schoolid,students) for schoolid in 1:300],Ref(parameters));
# Ovetwrite beta
for schoolid in 1:length(simstudents)
    setfield!.(simstudents[schoolid].pairs[1:5:end,1:5:end],:pairwiseβ,Ref(simstudents[1].parameters.β[4]))
    setfield!.(simstudents[schoolid].pairs[1:5:end,1:5:end],:δ,Ref(simstudents[1].parameters.δ[4]))
end

#parameter settings
## infection hazard
###### infection hazard
using RCall
@rimport distr

## COVID-19
#incperiod=distr.Lnorm(1.434065,0.6612) # He et al.
#infprof=distr.Norm(0.53,sqrt(7)) # Aschroft et al.
#gtime=incperiod+infprof

## Pandemic influenza
incperiod=distr.Lnorm(0.34,0.42) # Lessler et al. 2009 (https://dx.doi.org/10.1016%2FS1473-3099(09)70069-6)
gtime=((mean,sd)->distr.Gammad(mean^2/sd^2, sd^2/mean))(1.7,1)
#

gtimevec=distr.p(gtime)((0:20))|>RCall.unsafe_vec|>diff
gtimevec=gtimevec./sum(gtimevec)

# set hazard function
symratio=(0:0.25:1)'
hazards_symsc=gtimevec.*(1 .- symratio.*(distr.p(incperiod)((1:20).-1)|>RCall.unsafe_vec))
surv=(0:0.1:0.5)'
hazards_randsc=[hazard.*((1 .-surv).^(0:19)) for hazard in eachcol(hazards_symsc)]
# enbed generation time distribution in parameters
function setgtime!(parameters,newpdf)
    SchoolOutbreak.updateArray!(parameters.pdfgtime,[0;newpdf;0])
    SchoolOutbreak.updateArray!(parameters.cdfgtime,[0;0;cumsum(newpdf)])
end
gtimes=[hazards_symsc[:,1]]#,hazards_symsc[:,3], hazards_randsc[3][:,2]]|>permutedims # no intervention, symptom screening, symptom+test

# Outbreak simulation
## Functions
incidence(outbreak)=length(outbreak)
incidence(outbreak::AbstractArray{<:AbstractVector})=incidence.(outbreak)
incidence(outbreak::AbstractArray{<:AbstractMatrix})=sum.(outbreak)

function outbreakscenarios!(simstudents,updateparameters,gtime)
    setgtime!(simstudents[1].parameters,gtime)
    outbreaks=SchoolOutbreak.simulateoutbreaks!.(Ref(simstudents), updateparameters, Ref(1:360))
    outbreaks#readoutbreakbyclass_reorder.(outbreaks,Ref(simstudents),20,Ref([1]))
end

include("src/hetLK.jl")
function hh_lkhratio(i,n,diOnset,rcomdist,rcomdays,rcom=0.02)
    sumi=sum(i)
    if sumi==1||i[1]==0 return(-Inf) end
    εh=-log.(1 .- rcom)
    εh*=rcomdays*pdf(rcomdist,diOnset)
    R = ones(2,2).*0.1
    ll_notfromschool = hetLK.ll(i,n,εh.*[0.0;ones(1)],R,0.)
    ll_fromschool = hetLK.ll(i,n,[1e20;εh.*ones(1)],R,0.) #+ log(εh[1])
    return(ll_notfromschool-ll_fromschool)
end

function ll(lparms,schooldata=householddata)
    parms=exp.(copy(lparms))
    @views SchoolOutbreak.updateArray!.(schooldata[1].parameters.β,[0;0;parms[1];0;0])
    @views SchoolOutbreak.updateArray!.(schooldata[1].parameters.γ,zeros(5))
    @views SchoolOutbreak.updateArray!.(schooldata[1].parameters.δ,zeros(5))
    @views SchoolOutbreak.updateArray!(schooldata[1].parameters.rcom,parms[2])
    #@views SchoolOutbreak.updateArray!(schooldata[1].parameters.suscoef,length(parms)>12&&parms[12:15])
    #@views SchoolOutbreak.updateArray!(schooldata[1].parameters.infcoef,parms[end-(length(parms)-12)÷2:end])
    -SchoolOutbreak.llfunc!(schooldata)
end

# repeated simulation
using Optim, LinearAlgebra, Random, Calculus
function repeatestimate(students,parameters, updateparameters,gtimes,assumedrcom=0.02,rcomdays=3,len=50)
    repeatedests=Tuple{Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64}[]
    for x in 1:len
        Random.seed!(x)
        #SchoolOutbreak.updateArray!.(getfield.(students[1:5:end],:HHllkh),-Inf)
        outbreakscenarios!.(Ref(simstudents),Ref(updateparameters),gtimes);
        householddat=SchoolOutbreak.StudentsData.([filter(x->x.stratum[1]==schoolid,students[1:5:end]) for schoolid in 1:300],Ref(parameters),Ref(Dict()));
        for data in householddat data.βmat.=Ref(data.parameters.β[4]) end
        for data in householddat data.δmat.=Ref(data.parameters.δ[4]) end

        hhfoi=[hh_lkhratio([Int(students[n].isinfected) sum(getfield.((@view students[n.+(1:4)]),:isinfected))]|>permutedims,[1 4]|>permutedims,students[n].onset,rcomdist,rcomdays,assumedrcom) for n in 1:5:length(students)]; # 3 days window period
        SchoolOutbreak.updateArray!.(getfield.(students[1:5:end],:HHllkh),hhfoi)
        for n in 1:length(householddat), member in householddat[n].members
            member.cdfcommunity.=cdf(rcomdist,member.onset)
            member.pdfcommunity.=pdf(rcomdist,member.onset)
        end
print("$x,")
        @time opt=optimize(p->ll(p,householddat),fill(-5.,3),fill(1.,3),fill(-1.,3))
        hess=Calculus.hessian(p->ll(p,householddat),opt.minimizer)
print(eigen(hess))
        sd=1.96.*sqrt.(diag(inv(hess)))
        #=βnorm=Normal(0.04,0.005)
        rcomnorm=Normal(0.02,0.001)
        ns=10000
        candidatepars=zip(log.(abs.(rand(βnorm,ns))),log.(abs.(rand(rcomnorm,ns))))
        lls=[ll(pars|>collect,householddata)-logpdf(βnorm,pars[1])-logpdf(βnorm,pars[1]) for pars in candidatepars]
        pos=collect(candidatepars)[sample(1:ns,weights(exp.(lls.-logsumexp(lls))),ns)]
        qs=quantile.([first.(pos).|>exp,last.(pos).|>exp],Ref([0.5,0.025,0.975]))=#
        push!(repeatedests,[vec((opt.minimizer.|>exp).*hcat(ones(3),exp.(sd),exp.(.-sd)));minimum(eigen(hess).values)]|>Tuple)
    end
    repeatedests
end
function estplot(repeatedests,colors)
    plots=[plot((1:50).+[0 0.33 0.66][:,1:size(repeatedests,2)],getindex.(repeatedests,1),ylim=(0.0,0.1),yerror=zip(getindex.(repeatedests,4).-getindex.(repeatedests,1),.-getindex.(repeatedests,7).+getindex.(repeatedests,1))|>collect,linewidth=0,markerstrokecolor=colors,color=colors,marker=:circle,markersize=1.5),
    plot((1:50).+[0 0.33 0.66][:,1:size(repeatedests,2)],getindex.(repeatedests,2),ylim=(0,1),yerror=zip(getindex.(repeatedests,5).-getindex.(repeatedests,2),.-getindex.(repeatedests,8).+getindex.(repeatedests,2))|>collect,linewidth=0,marker=:circle,markerstrokecolor=colors,color=colors,markersize=1.5),
    plot((1:50).+[0 0.33 0.66][:,1:size(repeatedests,2)],getindex.(repeatedests,3),ylim=(0,0.05),yerror=zip(getindex.(repeatedests,6).-getindex.(repeatedests,3),.-getindex.(repeatedests,9).+getindex.(repeatedests,3))|>collect,linewidth=0,marker=:circle,markerstrokecolor=colors,color=colors,markersize=1.5)]
    hline!.(plots,[[0.05],[0.5],[0.02]],color=:black,linestyle=:dash,label="")
end

## Simulation set 1
updateparameters=fill((β=[0,0,0,0.05,0.1], γ=zeros(5), δ=[0,0,0,0.5,0],
    rcom=0.02, suscoef=ones(4), infcoef=ones(5)),1)
repeatests1=repeatestimate.(Ref(students),Ref(parameters),Ref(updateparameters),Ref(gtimes),[0.02,0.02,0.,0.02],[1,3,3,7],50);
plot([plot(estplot(repeatests1[1],:gray40)...,title=["within-school transmissibility" "exponent parameter" "risk of infection from community"],layout=(1,3));
plot(estplot(reduce(hcat,repeatests1[[3,4]]),[1 2 3])...,xlabel="simulation runs",guidefontsize=9,layout=(1,3))]...,layout=(2,1),size=(650,400),legend=:none)

## Simulation set 2
updateparameters=fill((β=[0,0,0,0.08,0.1], γ=zeros(5), δ=[0,0,0,0.8,0],
    rcom=0.05, suscoef=ones(4), infcoef=ones(5)),1)
repeatests2=repeatestimate.(Ref(students),Ref(parameters),Ref(updateparameters),Ref(gtimes),[0.05,0.05,0.,0.05],[1,3,3,7],50);

plot([plot(estplot(repeatests2[1],:gray40)...,title=["within-school transmissibility" "exponent parameter" "risk of infection from community"],layout=(1,3));
plot(estplot(reduce(hcat,repeatests2[3:4]),[1 2 3])...,xlabel="simulation runs",guidefontsize=9,layout=(1,3))]...,layout=(2,1),legend=:none,size=(650,400))

## Save results
using JLD2
@save "repeatests.jld2" repeatests1 repeatests2


text:
"""
    can_do_fast_attack(knight_awake)

Return if a fast attack is possible.
"""
can_do_fast_attack(knight_awake) = !knight_awake

"""
    can_spy(knight_awake, archer_awake, prisoner_awake)

Return if the spy action is possible.
"""
can_spy(knight_awake, archer_awake, prisoner_awake) = knight_awake || archer_awake || prisoner_awake

"""
    can_signal_prisoner(archer_awake, prisoner_awake)

Return if signaling the prisoner is possible.
"""
can_signal_prisoner(archer_awake, prisoner_awake) = prisoner_awake && !archer_awake

"""
    can_free_prisoner(knight_awake, archer_awake, prisoner_awake, dog_present)

Return if the prisoner can be freed.
"""
can_free_prisoner(knight_awake, archer_awake, prisoner_awake, dog_present) = !knight_awake && !archer_awake && prisoner_awake || dog_present && !archer_awake

ext: jl


text:
# Benchmark Catlab.Graphs against LightGraphs and MetaGraphs.
using BenchmarkTools
const SUITE = BenchmarkGroup()

using Random
import LightGraphs, MetaGraphs
const LG, MG = LightGraphs, MetaGraphs

using Catlab, Catlab.CategoricalAlgebra, Catlab.Graphs
using Catlab.Graphs.BasicGraphs: TheoryGraph
using Catlab.WiringDiagrams: query
using Catlab.Programs: @relation

# Helpers
#########

# `bench_iter_edges` and `bench_has_edge` adapted from LightGraphs:
# https://github.com/JuliaGraphs/LightGraphs.jl/blob/master/benchmark/core.jl

function bench_iter_edges(g)
  count = 0
  for e in edges(g)
    s, t = src(g,e), tgt(g,e)
    count += 1
  end
  count
end
function bench_iter_edges_vectorized(g)
  count = 0
  for (s,t) in zip(src(g), tgt(g))
    count += 1
  end
  count
end

function bench_has_edge(g)
  Random.seed!(1)
  n = nv(g)
  srcs, tgts = rand(1:n, n÷4), rand(1:n, n÷4)
  count = 0
  for (s, t) in zip(srcs, tgts)
    if has_edge(g, s, t)
      count += 1
    end
  end
  count
end

function bench_iter_neighbors(g)
  count = 0
  for v in vertices(g)
    count += length(neighbors(g, v))
  end
  count
end

@inline Graphs.nv(g::LG.AbstractGraph) = LG.nv(g)
@inline Graphs.vertices(g::LG.AbstractGraph) = LG.vertices(g)
@inline Graphs.edges(g::LG.AbstractGraph) = LG.edges(g)
@inline Graphs.src(g::LG.AbstractGraph, e::LG.AbstractEdge) = LG.src(e)
@inline Graphs.tgt(g::LG.AbstractGraph, e::LG.AbstractEdge) = LG.dst(e)
@inline Graphs.has_edge(g::LG.AbstractGraph, args...) = LG.has_edge(g, args...)
@inline Graphs.neighbors(g::LG.AbstractGraph, args...) = LG.neighbors(g, args...)

function Graphs.connected_component_projection(g::LG.AbstractGraph)
  label = Vector{Int}(undef, LG.nv(g))
  LG.connected_components!(label, g)
end

abstract type FindTrianglesAlgorithm end
struct TriangleHomomorphism <: FindTrianglesAlgorithm end
struct TriangleQuery <: FindTrianglesAlgorithm end

""" Number of triangles in a graph.
"""
function ntriangles(g::T, ::TriangleHomomorphism) where T
  triangle = T(3)
  add_edges!(triangle, [1,2,1], [2,3,3])
  count = 0
  homomorphisms(triangle, g) do α;
    count += 1; return false
  end
  count
end
function ntriangles(g, ::TriangleQuery)
  length(query(g, ntriangles_query))
end
const ntriangles_query = @relation (;) begin
  E(src=v1, tgt=v2)
  E(src=v2, tgt=v3)
  E(src=v1, tgt=v3)
end

# Graphs
########

bench = SUITE["Graph"] = BenchmarkGroup()
clbench = bench["Catlab"] = BenchmarkGroup()
clvecbench = bench["Catlab-vectorized"] = BenchmarkGroup()
lgbench = bench["LightGraphs"] = BenchmarkGroup()

n = 10000
clbench["make-path"] = @benchmarkable path_graph(Graph,n)
lgbench["make-path"] = @benchmarkable begin
  g = LG.DiGraph()
  LG.add_vertices!(g, n)
  for v in 1:(n-1)
    LG.add_edge!(g, v, v+1)
  end
end

g = path_graph(Graph, n)
lg = LG.DiGraph(g)

clbench["iter-edges"] = @benchmarkable bench_iter_edges($g)
clvecbench["iter-edges"] = @benchmarkable bench_iter_edges_vectorized($g)
lgbench["iter-edges"] = @benchmarkable bench_iter_edges($lg)
clbench["has-edge"] = @benchmarkable bench_has_edge($g)
lgbench["has-edge"] = @benchmarkable bench_has_edge($lg)
clbench["iter-neighbors"] = @benchmarkable bench_iter_neighbors($g)
lgbench["iter-neighbors"] = @benchmarkable bench_iter_neighbors($lg)

n₀ = 2000
g₀ = path_graph(Graph, n₀)
g = ob(coproduct(fill(g₀, 5)))
lg = LG.DiGraph(g)
clbench["path-graph-components"] =
  @benchmarkable connected_component_projection($g)
lgbench["path-graph-components"] =
  @benchmarkable connected_component_projection($lg)

g₀ = star_graph(Graph, n₀)
g = ob(coproduct(fill(g₀, 5)))
lg = LG.DiGraph(g)
clbench["star-graph-components"] =
  @benchmarkable connected_component_projection($g)
lgbench["star-graph-components"] =
  @benchmarkable connected_component_projection($lg)

n = 100
g = wheel_graph(Graph, n)
lg = LG.DiGraph(g)
clbench["wheel-graph-triangles-hom"] =
  @benchmarkable ntriangles($g, TriangleHomomorphism())
clbench["wheel-graph-triangles-query"] =
  @benchmarkable ntriangles($g, TriangleQuery())

# Symmetric graphs
##################

bench = SUITE["SymmetricGraph"] = BenchmarkGroup()
clbench = bench["Catlab"] = BenchmarkGroup()
clvecbench = bench["Catlab-vectorized"] = BenchmarkGroup()
lgbench = bench["LightGraphs"] = BenchmarkGroup()

n = 10000
clbench["make-path"] = @benchmarkable path_graph(SymmetricGraph, n)
lgbench["make-path"] = @benchmarkable begin
  g = LG.Graph()
  LG.add_vertices!(g, n)
  for v in 1:(n-1)
    LG.add_edge!(g, v, v+1)
  end
end

g = path_graph(SymmetricGraph, n)
lg = LG.Graph(g)

clbench["iter-edges"] = @benchmarkable bench_iter_edges($g)
clvecbench["iter-edges"] = @benchmarkable bench_iter_edges_vectorized($g)
lgbench["iter-edges"] = @benchmarkable bench_iter_edges($lg)
clbench["has-edge"] = @benchmarkable bench_has_edge($g)
lgbench["has-edge"] = @benchmarkable bench_has_edge($lg)
clbench["iter-neighbors"] = @benchmarkable bench_iter_neighbors($g)
lgbench["iter-neighbors"] = @benchmarkable bench_iter_neighbors($lg)

n₀ = 2000
g₀ = path_graph(SymmetricGraph, n₀)
g = ob(coproduct(fill(g₀, 5)))
lg = LG.Graph(g)
clbench["path-graph-components"] =
  @benchmarkable connected_component_projection($g)
lgbench["path-graph-components"] =
  @benchmarkable connected_component_projection($lg)

g₀ = star_graph(SymmetricGraph, n₀)
g = ob(coproduct(fill(g₀, 5)))
lg = LG.Graph(g)
clbench["star-graph-components"] =
  @benchmarkable connected_component_projection($g)
lgbench["star-graph-components"] =
  @benchmarkable connected_component_projection($lg)

n = 100
g = wheel_graph(SymmetricGraph, n)
lg = LG.Graph(g)
clbench["wheel-graph-triangles-hom"] =
  @benchmarkable ntriangles($g, TriangleHomomorphism())
clbench["wheel-graph-triangles-query"] =
  @benchmarkable ntriangles($g, TriangleQuery())
lgbench["wheel-graph-triangles"] = @benchmarkable sum(LG.triangles($lg))

# Weighted graphs
#################

bench = SUITE["WeightedGraph"] = BenchmarkGroup()
clbench = bench["Catlab"] = BenchmarkGroup()
clvecbench = bench["Catlab-vectorized"] = BenchmarkGroup()
lgbench = bench["LightGraphs"] = BenchmarkGroup()

n = 10000
g = path_graph(WeightedGraph{Float64}, n; E=(weight=range(0,1,length=n-1),))
mg = MG.MetaDiGraph(g)

clvecbench["sum-weights"] = @benchmarkable sum(weight($g))
clbench["sum-weights"] = @benchmarkable begin
  total = 0.0
  for e in edges($g)
    total += weight($g, e)
  end
  total
end
lgbench["sum-weights"] = @benchmarkable begin
  total = 0.0
  for e in MG.edges($mg)
    total += MG.get_prop($mg, e, :weight)
  end
  total
end

clvecbench["increment-weights"] = @benchmarkable begin
  $g[:weight] .= $g[:weight] .+ 1.0
end

clbench["increment-weights"] = @benchmarkable begin
  for e in edges($g)
    $g[e,:weight] += 1.0
  end
end
lgbench["increment-weights"] = @benchmarkable begin
  for e in MG.edges($mg)
    MG.set_prop!($mg, e, :weight, MG.get_prop($mg, e, :weight) + 1.0)
  end
end

# Labeled graphs
################

bench = SUITE["LabeledGraph"] = BenchmarkGroup()
clbench = bench["Catlab"] = BenchmarkGroup()
lgbench = bench["LightGraphs"] = BenchmarkGroup()

@present TheoryLabeledGraph <: TheoryGraph begin
  Label::AttrType
  label::Attr(V,Label)
end
@acset_type LabeledGraph(TheoryLabeledGraph, index=[:src,:tgt]) <: AbstractGraph
@acset_type IndexedLabeledGraph(TheoryLabeledGraph, index=[:src,:tgt],
                                unique_index=[:label]) <: AbstractGraph

function discrete_labeled_graph(n::Int; indexed::Bool=false)
  g = (indexed ? IndexedLabeledGraph{String} : LabeledGraph{String})()
  add_vertices!(g, n, label=("v$i" for i in 1:n))
  g
end

function discrete_labeled_metagraph(n::Int; indexed::Bool=false)
  mg = MG.MetaDiGraph()
  for i in 1:n
    MG.add_vertex!(mg, :label, "v$i")
  end
  if indexed; MG.set_indexing_prop!(mg, :label) end
  mg
end

n = 5000
clbench["make-discrete"] = @benchmarkable discrete_labeled_graph($n)
lgbench["make-discrete"] = @benchmarkable discrete_labeled_metagraph($n)
clbench["make-discrete-indexed"] =
  @benchmarkable discrete_labeled_graph($n, indexed=true)
lgbench["make-discrete-indexed"] =
  @benchmarkable discrete_labeled_metagraph($n, indexed=true)

n = 10000
g = discrete_labeled_graph(n)
mg = discrete_labeled_metagraph(n)
clbench["iter-labels"] = @benchmarkable begin
  for v in vertices($g)
    label = $g[v,:label]
  end
end
lgbench["iter-labels"] = @benchmarkable begin
  for v in MG.vertices($mg)
    label = MG.get_prop($mg, v, :label)
  end
end

g = discrete_labeled_graph(n, indexed=true)
mg = discrete_labeled_metagraph(n, indexed=true)
Random.seed!(1)
σ = randperm(n)
clbench["indexed-lookup"] = @benchmarkable begin
  for i in $σ
    @assert incident($g, "v$i", :label) == i
  end
end
lgbench["indexed-lookup"] = @benchmarkable begin
  for i in $σ
    @assert $mg["v$i", :label] == i
  end
end


text:
module DataSeries

using  DataFrames, Distributions, StatsBase, Dates, Serialization, CSV

export coarseToFine, preProcessYahoo, flattenYears, flattenQuarters, getYahooData,
  dailyToEOM, dailyToEOQ, dailyToEOY

const YAHOO_DATE_FORMAT = "m/d/y"

const DEBUG_DATASERIES = true

function preProcessYahoo(path::String, name::String; makeReturns::Bool = true)::Nothing
  dateFormat::String = YAHOO_DATE_FORMAT
  DF::DataFrame = CSV.read("$path\\$name.csv")

  #fix the dates
  rename!(DF, :Date => :DATE)
  rename!(DF, Symbol("Adj Close") => :Adj_Close)

  if makeReturns
    DF[:DATE] = ((s::String)->Date(s, dateFormat)).(
      DF[:DATE])::V where V <: Union{Vector{Union{Date, Missing}},Vector{Date}}
    sort!(DF, [:DATE])
    DF[:returns] = similar(DF[:Adj_Close])
    DF[2:end,:returns] = log.(DF[2:end,:Adj_Close] ./  DF[1:(end-1),:Adj_Close])
  end
  #showcols(DF)

  stream::IOStream = open("$path\\$name.jls", "w")
  serialize(stream, DF)
  close(stream)
  return nothing
end

#this function processes yahoo data, with an optional shortcut of the file read
function getYahooData(path::String, name::String; refreshData::Bool = true,
  makeReturns::Bool = true)::DataFrame

  if !isfile("$path\\$name.jls") || refreshData
    preProcessYahoo(path, name; makeReturns = makeReturns)
  end

  stream::IOStream = open("$path\\$name.jls")
    DF::DataFrame = deserialize(stream)
  close(stream)

  return DF
end

#this function gets the data point for the last available day in a month
function convertToMaxMonth(coarseDate::Date, fineDates::Vector,
  fineData::Vector{Union{T, Missing}}, convertedFineDates::Vector)::T where T<:Real

  maxDate::Date = maximum(fineDates[convertedFineDates.==coarseDate])
  return fineData[something(findfirst(isequal(maxDate),fineDates),0)]
end

dailyToEOM(coarseDates::Vector,
  fineDates::Vector, fineData::Vector{R where R<: Union{T, Missing}}
  ) where T<: Real = fineToCoarse(coarseDates,
    fineDates, fineData, firstdayofmonth.(fineDates), convertToMaxMonth)

dailyToEOQ(coarseDates::Vector,
  fineDates::Vector, fineData::Vector{R where R<: Union{T, Missing}}
  ) where T<: Real= fineToCoarse(coarseDates,
  fineDates, fineData, firstdayofquarter.(fineDates), convertToMaxMonth)

dailyToEOY(years::Vector,
  fineDates::Vector, fineData::Vector{R where R<: Union{T, Missing}}
  ) where T<: Real= fineToCoarse(
  Vector{Union{Date, Missing}}(firstdayofyear.(Date.(years))),
  fineDates, fineData, firstdayofyear.(fineDates), convertToMaxMonth)

#fineToCoarse(coarseDF::DataFrame, fineDF::DataFrame, coarseField = :DATE, fine)

#this function takes in fine data and aggregates it upward
function fineToCoarse(coarseDates::Vector,
  fineDates::Vector, fineData::Vector{Union{T, Missing}},
  convertedFineDates::Vector,
  conversionFunc::Function)::Vector{Union{T, Missing}} where T<:Real

  #pre-allcoate the output vector
  out::Vector{Union{T, Missing}} = Vector{Union{T, Missing}}(undef, length(coarseDates))

  for i::Int ∈ 1:length(coarseDates)
    if coarseDates[i] ∈ convertedFineDates
#      println("ctf test: ", coarseDates[i])
      out[i] = conversionFunc(coarseDates[i], fineDates, fineData, convertedFineDates)
    else
      out[i] = missing
    end
  end

  return out
end

#This function takes lower resolution time series and paste values over a higher resolution period
#IN: Data vectors for coarse dates, coarse data, higher resolution dates. Also requires a many to one
# mapping function of fine dates to coarse dates
#OUT: A data vector of the coarse data mapped to the higher reoslution fine dates
function coarseToFine(coarseDates::Vector{Union{Date, Missing}}, coarseData::Vector{Union{T, Missing}},
  fineDates::Vector{V}, conversionFunc::Function
  )::Vector where {T<:Real, V<:Union{Date,Missing}}

  #pre-allcoate the output vector
  out::Vector{Union{T, Missing}} = Vector{Union{T, Missing}}(undef, length(fineDates))

  #create a lookup table for extacting coarse data given a coarse date
  dateToDataTable::Dict = Dict(coarseDates[i]=>coarseData[i] for i::Int ∈ 1:length(coarseDates))

  #get the data from the look-up table using the mapping function and the fine date input
  for i::Int ∈ 1:length(fineDates)
    out[i] = dateToDataTable[conversionFunc(fineDates[i])]
  end

  return out
end

flattenQuarters(coarseDates::Vector{Union{Date, Missing}}, coarseData::Vector{Union{T, Missing}},
  fineDates::Vector{Union{Date, Missing}}) where T<:Real = coarseToFine(coarseDates,
  coarseData, fineDates, Dates.firstdayofquarter)::Vector{Union{T, Missing}}

#wrapper function for coarseToFine which flattens yearly data to finer resolution
flattenYears(coarseDates::Vector{Union{Date, Missing}}, coarseData::Vector{Union{T, Missing}},
  fineDates::Vector{Date}) where T<:Real = coarseToFine(coarseDates,
  coarseData, fineDates, Dates.firstdayofyear)::Vector{Union{T, Missing}}

#helper method which takes years as an integer
flattenYears(years::Vector{Union{Int, Missing}}, coarseData::Vector{Union{T, Missing}},
  fineDates::Vector{Date}) where T<:Real = flattenYears(Vector{Union{Date, Missing}}(Date.(years)),
  coarseData, fineDates)::Vector{Union{T, Missing}}

end

max_stars_repo_path: Fall 17 Ivo Project 2.2/histtaxes/DataSeries.jl

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
# 0.18 deprecations. Remove before releasing 0.19
@deprecate path(t::Trie, str::AbstractString) partial_path(t::Trie, str::AbstractString)
@deprecate find_root find_root!
@deprecate top first
@deprecate reverse_iter Iterators.reverse
# Deprecations from #700
Base.@deprecate_binding DisjointSets DisjointSet
Base.@deprecate_binding IntDisjointSets IntDisjointSet
@deprecate DisjointSets(xs...) DisjointSet(xs)
# Enqueue and dequeue deprecations
@deprecate enqueue!(q::Queue, x) Base.push!(q, x)
@deprecate enqueue!(q::PriorityQueue, x) Base.push!(q, x)
@deprecate enqueue!(q::PriorityQueue, k, v) Base.push!(q, k=>v)
@deprecate dequeue!(q::Queue) Base.popfirst!(q)
@deprecate dequeue!(q::PriorityQueue) Base.popfirst!(q).first # maybe better: `val, _ = popfirst!(pq)`
@deprecate dequeue!(q::PriorityQueue, x) popat!(q, x).first
@deprecate dequeue_pair!(q::PriorityQueue) Base.popfirst!(q)
@deprecate dequeue_pair!(q::PriorityQueue, key) popat!(q, key)

function Base.peek(q::PriorityQueue)
    Expr(:meta, :noinline)
    Base.depwarn("`peek(q::PriorityQueue)` is deprecated, use `first(q)` instead.", :peek)
    first(q)
end


text:
module PPSeq

# === IMPORTS === #

using LinearAlgebra
using Statistics

import Base: size, rand, length, getindex, iterate
import Base.Iterators
import StatsBase: sample, pweights, denserank, mean
import Random
import SpecialFunctions: logabsgamma, logfactorial
import StatsFuns: softmax!, softmax, logaddexp, logsumexp, normlogpdf, normpdf
import Distributions
import Distributions: cdf
import PyPlot: plt

# === EXPORTS === #

export Spike, EventSummaryInfo, SeqHypers, SeqGlobals, SeqModel, DistributedSeqModel
export log_joint, log_prior, log_p_latents
export SymmetricDirichlet, RateGamma, ScaledInvChisq, NormalInvChisq
export specify_gamma, sample
export split_merge_sample!, gibbs_sample!, annealed_gibbs!
export masked_gibbs!, annealed_masked_gibbs!, Mask
export create_random_mask, split_spikes_by_mask, create_blocked_mask, sample_masked_spikes!

# === UTILS === #

include("./utils/misc.jl")
include("./utils/distributions.jl")

# === MODEL === #

 # Model structs
 # - SeqModel : holds full model
 # - SeqEvent : holds suff stats for a latent event
 # - SeqEventList : dynamically re-sized array of latent events.
 # - SeqPriors : holds prior distributions.
 # - SeqGlobals : holds global variables.
include("./model/structs.jl")

 # Convience methods for constructing / accessing SeqModel model struct.
include("./model/model.jl")

 # Distributed version of the model
include("./model/distributed.jl")

 # Methods for creating latent events, managing their storage.
include("./model/events.jl")

 # Adds and removes spikes from events, updates to sufficient statistics.
include("./model/add_remove.jl")

 # Evaluate various probability distributions for the model:
 # - predictive posterior
 # - prior distribution on global parameters
 # - distribution on latent events, given global parameters
 # - joint distribution on global parameters, latent events, observed spikes
include("./model/probabilities.jl")

 # Evaluate log likelihood on heldout data.
include("./model/masked_probabilities.jl")

# === PARAMETER INFERENCE === #

 # Collapsed Gibbs sampling.
include("./algorithms/gibbs.jl")

 # Masked Gibbs sampling (for cross-validation).
include("./algorithms/masked_gibbs.jl")

 # Distributed collapsed Gibbs sampling
include("./algorithms/distributed_gibbs.jl")

 # Collapsed Gibbs sampling with annealing.
include("./algorithms/annealed_gibbs.jl")

 # Split merge sampler.
include("./algorithms/split_merge.jl")

 # Easy one-stop-shop sampling function.
include("./algorithms/easy_sample.jl")

 # Equivalent easy one-stop-shop sampling function for masked sampling.
 include("./algorithms/easy_sample_masked.jl")

# === USER-FACING UTILS AND HELPER FUNCTIONS === #

 # Model configuration.
include("./utils/config.jl")

 # PyPlot visualizations.
include("./utils/visualization.jl")

 # Data analysis functions.
include("./utils/analysis.jl")

# === END OF MODULE === #

end


text:
# dbn_preprocess.jl
# 2019-11-12
# David Merrell
#
#

using CSV
using DataFrames
using Statistics

function ingest_timeseries!(timeseries_df)

    result = Vector{Matrix{Float64}}()

    #N = convert(Int64, maximum(timeseries_df.timeseries))

    for name in unique(timeseries_df.timeseries)
        ts = timeseries_df[timeseries_df.timeseries .== name, :]
        sort!(ts, [:timestep])
        push!(result, convert(Matrix{Float64}, ts[:,3:end]))
    end

    return result
end

function ingest_conf_adjacency!(adjacency_df)

    parent_confs = Vector{Dict{Int64,Float64}}()
    adj_mat = convert(Matrix{Float64}, adjacency_df)

    for j=1:size(adj_mat,2)
        ps = Dict{Int64,Float64}()
        for i=1:size(adj_mat,1)
            if adj_mat[i,j] > 0.0
                ps[i] = adj_mat[i,j]
            end
        end
        push!(parent_confs, ps)
    end

    return parent_confs
end

"""
    load_formatted_data(timeseries_data_path::String,
                        reference_graph_path::String)

Load data from files and convert to the types requisite
for our probabilistic program.

This function assumes the data (time series and adjacency matrices)
are stored in text-delimited files formatted in a very particular way.
"""
function load_formatted_data(timeseries_data_path::String,
	                     reference_graph_path::String)

    timeseries_df = DataFrame!(CSV.File(timeseries_data_path))
    ts_vec = ingest_timeseries!(timeseries_df)

    ref_adj_df = DataFrame!(CSV.File(reference_graph_path, header=false))
    ref_ps = ingest_conf_adjacency!(ref_adj_df)

    return ts_vec, ref_ps
end

function combine_X(X::Vector{Array{Float64,2}})

    len = sum([size(x)[1] - 1 for x in X])
    wid = size(X[1])[2]

    Xminus = zeros(len, wid)
    Xplus = zeros(len, wid)

    l_ind = 1
    for x in X
        r_ind = l_ind + size(x)[1] - 2
        Xminus[l_ind:r_ind, :] = x[1:size(x)[1]-1, :]
        Xplus[l_ind:r_ind, :] = x[2:size(x)[1], :]
        l_ind = r_ind + 1
    end

    return Xminus, Xplus
end

function standardize_X(X_arr)

    mus = Statistics.mean(X_arr, dims=1)
    sigmas = Statistics.std(X_arr, dims=1)

    X_arr =  (X_arr .- mus) ./ sigmas

    return X_arr, mus, sigmas
end

function vectorize_X(Xminus::Array{Float64,2},
		     Xplus::Array{Float64,2})

    Xplus_v = [Xplus[:,i] for i=1:size(Xplus)[2]]

    return Xminus, Xplus_v

end

max_stars_repo_name: gitter-lab/ssps

max_issues_repo_name: gitter-lab/ssps

max_forks_repo_name: gitter-lab/ssps


text:
"""
*Abstract* type for Holiday Calendars.
"""
abstract type HolidayCalendar end

Base.string(hc::HolidayCalendar) = string(typeof(hc))

Base.broadcastable(hc::HolidayCalendar) = Ref(hc)

function symtocalendar(sym::Symbol) :: HolidayCalendar
    if isdefined(BusinessDays, sym) && Core.eval(BusinessDays, sym) <: HolidayCalendar
        return Core.eval(BusinessDays, sym)()
    elseif isdefined(@__MODULE__, sym) && Core.eval(@__MODULE__, sym) <: HolidayCalendar
        return Core.eval(@__MODULE__, sym)()
    elseif isdefined(Main, sym) && Core.eval(Main, sym) <: HolidayCalendar
        return Core.eval(Main, sym)()
    else
        error("$sym is not a valid HolidayCalendar.")
    end
end

@inline strtocalendar(str::AbstractString) = symtocalendar(Symbol(str))

Base.convert(::Type{HolidayCalendar}, sym::Symbol) = symtocalendar(sym)
Base.convert(::Type{HolidayCalendar}, str::AbstractString) = strtocalendar(str)

"""
    isholiday(calendar, dt)

Checks if `dt` is a holiday based on a given `calendar` of holidays.

`calendar` can be an instance of `HolidayCalendar`,  a `Symbol` or an `AbstractString`.

Returns boolean values.
"""
isholiday(hc::HolidayCalendar, dt::Dates.Date) = error("isholiday for $(hc) not implemented.")

isholiday(calendar, dt::Dates.Date) = isholiday(convert(HolidayCalendar, calendar), dt)


text:
import Pkg

function main(repo_url::AbstractString;
              build_directory::AbstractString = pwd(),
              master_branch::AbstractString = "master",
              allow_unchanged_prerelease::Bool = true)::Nothing
    head_version = get_head_version(repo_url; build_directory = build_directory)
    @info("Head version: $(head_version)")
    master_version = get_upstream_version(repo_url; master_branch = master_branch)
    @info("Master version: $(master_version)")
    compare_versions(master_version,
                     head_version;
                     allow_unchanged_prerelease = allow_unchanged_prerelease)
    return nothing
end


text:
module StructuralTransformations

const UNVISITED = typemin(Int)
const UNASSIGNED = typemin(Int)

using Setfield: @set!, @set
using UnPack: @unpack

using SymbolicUtils
using SymbolicUtils.Code
using SymbolicUtils.Rewriters
using SymbolicUtils: similarterm, istree

using ModelingToolkit
using ModelingToolkit: ODESystem, var_from_nested_derivative, Differential,
                       states, equations, vars, Symbolic, diff2term, value,
                       operation, arguments, Sym, Term, simplify, solve_for,
                       isdiffeq, isdifferential,
                       get_structure, get_reduced_states, defaults

using ModelingToolkit.BipartiteGraphs
using ModelingToolkit.SystemStructures

using ModelingToolkit.DiffEqBase
using ModelingToolkit.StaticArrays
using ModelingToolkit: @RuntimeGeneratedFunction, RuntimeGeneratedFunctions

RuntimeGeneratedFunctions.init(@__MODULE__)

using SparseArrays

using NonlinearSolve

export tearing, dae_index_lowering
export build_torn_function, build_observed_function, ODAEProblem
export sorted_incidence_matrix

include("utils.jl")
include("pantelides.jl")
include("bipartite_tearing/modia_tearing.jl")
include("tearing.jl")
include("codegen.jl")

end # module


text:
/* Test arithmetic and comparisons. */

int main() {
    int x = 56;
    int y = -23;
    printInt(x+y);
    printInt(x-y);
    printInt(x*y);
    printInt(45/2);
    printInt(78%3);
    double z = -9.3;
    double w = 5.1;
    printBool(z+w > z-w);
    printBool(z/w <= z*w);
    return 0 ;
}

void printBool(boolean b) {
  if (b) {
    printString("true");
    return;
  } else {
    printString("false");
    return;
 }
}

max_stars_repo_name: jwilk/jtc

max_issues_repo_name: jwilk/jtc

max_forks_repo_name: jwilk/jtc


text:
using Delaunator
using Test

@show Delaunator.delaunator!([[0, 1], [1, 0], [1, 1]])
@show Delaunator.delaunator!([[0, 1], [1, 0], [0, 0], [1, 1]])
@show Delaunator.delaunator!([[516, 661], [369, 793], [426, 539], [273, 525], [204, 694], [747, 750], [454, 390]])
# import {test} from 'tape';
# import Delaunator from '../index.js';

# import points from './fixtures/ukraine.json';
# import issue13 from './fixtures/issue13.json';
# import issue43 from './fixtures/issue43.json';
# import issue44 from './fixtures/issue44.json';
# import robustness1 from './fixtures/robustness1.json';
# import robustness2 from './fixtures/robustness2.json';

# test('triangulates plain array', (t) => {
#     const d = new Delaunator([].concat(...points));
#     t.same(d.triangles, Delaunator.from(points).triangles);
#     t.end();
# });

# test('triangulates typed array', (t) => {
#     const d = new Delaunator(Float64Array.from([].concat(...points)));
#     t.same(d.triangles, Delaunator.from(points).triangles);
#     t.end();
# });

# test('constructor errors on invalid array', (t) => {
#     /* eslint no-new: 0 */
#     t.throws(() => {
#         new Delaunator({length: -1});
#     }, /Invalid typed array length/);
#     t.throws(() => {
#         new Delaunator(points);
#     }, /Expected coords to contain numbers/);
#     t.end();
# });

# test('produces correct triangulation', (t) => {
#     validate(t, points);
#     t.end();
# });

# test('produces correct triangulation after modifying coords in place', (t) => {
#     const d = Delaunator.from(points);

#     validate(t, points, d);
#     t.equal(d.trianglesLen, 5133);

#     const p = [80, 220];
#     d.coords[0] = p[0];
#     d.coords[1] = p[1];
#     const newPoints = [p].concat(points.slice(1));

#     d.update();
#     validate(t, newPoints, d);
#     t.equal(d.trianglesLen, 5139);

#     t.end();
# });

# test('issue #11', (t) => {
#     validate(t, [[516, 661], [369, 793], [426, 539], [273, 525], [204, 694], [747, 750], [454, 390]]);
#     t.end();
# });

# test('issue #13', (t) => {
#     validate(t, issue13);
#     t.end();
# });

# test('issue #24', (t) => {
#     validate(t, [[382, 302], [382, 328], [382, 205], [623, 175], [382, 188], [382, 284], [623, 87], [623, 341], [141, 227]]);
#     t.end();
# });

# test('issue #43', (t) => {
#     validate(t, issue43);
#     t.end();
# });

# test('issue #44', (t) => {
#     validate(t, issue44);
#     t.end();
# });

# test('robustness', (t) => {
#     validate(t, robustness1);
#     validate(t, robustness1.map(p => [p[0] / 1e9, p[1] / 1e9]));
#     validate(t, robustness1.map(p => [p[0] / 100, p[1] / 100]));
#     validate(t, robustness1.map(p => [p[0] * 100, p[1] * 100]));
#     validate(t, robustness1.map(p => [p[0] * 1e9, p[1] * 1e9]));
#     validate(t, robustness2.slice(0, 100));
#     validate(t, robustness2);
#     t.end();
# });

# test('returns empty triangulation for small number of points', (t) => {
#     let d = Delaunator.from([]);
#     t.same(d.triangles, []);
#     t.same(d.hull, []);
#     d = Delaunator.from(points.slice(0, 1));
#     t.same(d.triangles, []);
#     t.same(d.hull, [0]);
#     d = Delaunator.from(points.slice(0, 2));
#     t.same(d.triangles, []);
#     t.same(d.hull, [1, 0]); // [0, 1] is also correct
#     t.end();
# });

# test('returns empty triangulation for all-collinear input', (t) => {
#     const d = Delaunator.from([[0, 0], [1, 0], [3, 0], [2, 0]]);
#     t.same(d.triangles, []);
#     t.same(d.hull, [0, 1, 3, 2]); // [2, 3, 0, 1] is also correct
#     t.end();
# });

# test('supports custom point format', (t) => {
#     const d = Delaunator.from(
#         [{x: 5, y: 5}, {x: 7, y: 5}, {x: 7, y: 6}],
#         p => p.x,
#         p => p.y);
#     t.same(d.triangles, [0, 2, 1]);
#     t.end();
# });

# function orient([px, py], [rx, ry], [qx, qy]) {
#     const l = (ry - py) * (qx - px);
#     const r = (rx - px) * (qy - py);
#     return Math.abs(l - r) >= 3.3306690738754716e-16 * Math.abs(l + r) ? l - r : 0;
# }
# function convex(r, q, p) {
#     return (orient(p, r, q) || orient(r, q, p) || orient(q, p, r)) >= 0;
# }

function validate(t, points, d = Delaunator.from(points)) {
    # validate halfedges
    for (let i = 0; i < d.halfedges.length; i++) {
        const i2 = d.halfedges[i];
        if i2 !== -1 && d.halfedges[i2] != i
            t.fail('invalid halfedge connection');
        end
    end
    t.pass('halfedges are valid');

    # validate triangulation
    const hullAreas = [];
    for (let i = 0, len = d.hull.length, j = len - 1; i < len; j = i++) {
        const [x0, y0] = points[d.hull[j]];
        const [x, y] = points[d.hull[i]];
        hullAreas.push((x - x0) * (y + y0));
        const c = convex(points[d.hull[j]], points[d.hull[(j + 1) % d.hull.length]],  points[d.hull[(j + 3) % d.hull.length]]);
        if (!c) t.fail(`hull is not convex at ${j}`);
    }
    hullArea = kahansum(hullAreas)

    const triangleAreas = [];
    for (let i = 0; i < d.triangles.length; i += 3) {
        const [ax, ay] = points[d.triangles[i]];
        const [bx, by] = points[d.triangles[i + 1]];
        const [cx, cy] = points[d.triangles[i + 2]];
        triangleAreas.push(Math.abs((by - ay) * (cx - bx) - (bx - ax) * (cy - by)));
    }
    trianglesArea = sum(triangleAreas)

    err = abs((hullArea - trianglesArea) / hullArea)
    if err <= 2^-51
        # t.pass(`triangulation is valid: ${err} error`);
        return true
    else
        #t.fail(`triangulation is broken: ${err} error`);
        return false
    end
end

# Kahan and Babuska summation, Neumaier variant; accumulates less FP error
function kahansum(x)
    sum = x[0]
    err = 0
    for i = 1:length(x)-1
        k = x[i]
        m = sum + k
        err += abs(sum) >= abs(k) ? sum - m + k : k - m + sum
        sum = m
    end
    return sum + err
end


text:
module Tools

export uutourhoc,uutourhop,heaviside,set_h_boundary!

"""
    this function's inputs are uu and gamma

        uu has 3 rows
        the 1st row is u (velocity)
        the 2nd row is m (mass flow rate)
        the 3rd row is ̂E (total energy per volume)

        gamma is the heat capacity ratio

    this function's outputs is tourhoc

        tourhoc has 3 rows
        the 1st row is u (velocity)
        the 2nd row is ρ (density)
        the 3rd row is c (speed of sound)

    this function uses perfect gas state function to get tourhoc
"""
function uutourhoc(uu::Any, gamma::Float64)

    # tourhoc = Array{Float64,2}(UndefInitializer(),length(uu[:,1]),length(uu[1,:]))

    tourhoc = deepcopy( zeros(size(uu)))

    tourhoc[1,:] = uu[2,:]./uu[1,:]
    tourhoc[2,:] = uu[1,:]

    pp = ((uu[3,:]-0.5tourhoc[2,:].*tourhoc[1,:].^2)*(gamma-1))

    tourhoc[3,:] = sqrt.(gamma.*abs.(pp)./(tourhoc[2,:]))

    return tourhoc

end

"""
    this function's inputs are uu and gamma

        uu has 3 rows
        the 1st row is u (velocity)
        the 2nd row is m (mass flow rate)
        the 3rd row is ̂E (total energy per volume)

        gamma is the heat capacity ratio

    this function's outputs is tourhop

        tourhoc has 3 rows
        the 1st row is u (velocity)
        the 2nd row is ρ (density)
        the 3rd row is p (pressure)

    this function uses perfect gas state function to get tourhop
"""

function uutourhop(uu::Any, gamma::Float64)

    tourhop = deepcopy( zeros(size(uu)))

    tourhop[1,:] = uu[2,:]./uu[1,:]
    tourhop[2,:] = uu[1,:]

    tourhop[3,:] = ((uu[3,:]-0.5tourhop[2,:].*tourhop[1,:].^2)*(gamma-1))

    return tourhop

end

"""
this function's inputs is x

    x is the coordinate

this function's outputs is a Float value

    this function is the heaviside step function (half-maximum convention)
"""
heaviside(x::AbstractFloat) = ifelse(x < 0, zero(x), ifelse(x > 0, one(x), oftype(x,0.5)))

end


text:
using CLArrays, GLVisualize, GeometryTypes, GLAbstraction, StaticArrays

TY = Float32
N = 1024
const h    = TY(2*π/N)
const epsn = TY(h * .5)
const C    = TY(2/epsn)
const tau  = TY(epsn * h)
Tfinal = 50.

S(x,y) = exp(-x^2/0.1f0)*exp(-y^2/0.1f0)

ArrayType = CLArray
# real-space and reciprocal-space grids
# the real-space grid is just used for plotting!
X_cpu = convert.(TY, collect(linspace(-pi+h, pi, N)) .* ones(1,N))
X = ArrayType(X_cpu);
k = collect([0:N/2; -N/2+1:-1]);
Â = ArrayType(convert.(TY,kron(k.^2, ones(1,N)) + kron(ones(N), k'.^2)));

# initial condition
uc = ArrayType(TY(2.0)*(rand(TY, N, N)-TY(0.5)))

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

pow3(u) = complex((u * u * u) - u)
function take_step!(u, Â, t_plot, fftplan!, ifftplan!, u3fft, uc, tmp)
    u3fft .= pow3.(u)
    fftplan! * u3fft
    uc .= complex.(u)
    fftplan! * uc
    @. tmp .= ((1f0+C*tau*Â) .* uc .- tau/epsn * (Â .* u3fft)) ./ (1f0+(epsn*tau)*Â.^2f0+C*tau*Â)
    ifftplan! * tmp
    u .= real.(tmp)
    nothing
end
function normalise_af!(u, out)
    out .= u .- minimum(u)
    out .= out ./ maximum(out)
    nothing
end
#################################################################
#################################################################

n = 1
T_plot = 0.01; t_plot = 0.0
ceil(Tfinal / tau)
up = copy(uc)
ucc = complex.(uc)
fftplan! = plan_fft!(ucc)
ifftplan! = plan_ifft!(ucc)
u3fft = similar(ucc)
tmp = similar(ucc)

w = glscreen(resolution = (N, N))
normalise_af!(uc,up)
up .= up .* 0.1f0
robj = visualize(
    reinterpret(Intensity{Float32}, Array(up)),
    stroke_width = 0f0,
    levels = 20f0,
    color_map = Colors.colormap("RdBu", 100),
    color_norm = Vec2f0(0, 0.1)
).children[]
_view(robj, position = Vec3f0(0, 0.5, 2.5), lookat = Vec3f0(0))
io, buffer = GLVisualize.create_video_stream(homedir()*"/Desktop/pd1.mkv", w)
center!(w, :orthographic_pixel)
idx = 0
while isopen(w)
    take_step!(uc, Â, t_plot, fftplan!, ifftplan!, u3fft, ucc, tmp)
    t_plot += tau
    normalise_af!(uc, up)
    up .= up .* 0.1f0
    GLAbstraction.update!(robj[:intensity], reinterpret(SVector{1, Float32}, Array(up)))
    GLWindow.poll_glfw()
    GLWindow.reactive_run_till_now()
    GLWindow.render_frame(w)
    GLWindow.swapbuffers(w)
    if idx == 4
        GLVisualize.add_frame!(io, w, buffer)
        idx = 0
    end
    idx += 1
end
close(io)
GLWindow.destroy!(w)


text:
t1 = time_ns() #tic()

include("test_baseline_correct.jl")
include("test_workflow.jl")
include("test_delete_slice_2D.jl")
include("test_delete_slice_3D.jl")
include("test_filterContinuous_DArray.jl")
include("test_filterContinuous_SharedArray.jl")

include("run_examples.jl")
include("run_spectrum.jl")
t2 = time_ns() #toc()
println(float(t2-t1)*1e9)


text:
function issue48test(lib::PolyhedraLibrary)
    v = [-5 -15; 2 3; -4 2; -2 5]
    V = vrep(v)
    p = polyhedron(V, lib)
    b = [70, 16, 8, 15]
    A = [-17  1
          -3  2
           1  2
          18 -7]
    inequality_fulltest(p, A, b, BitSet([]))
    generator_fulltest(p, v)

    ind = [3, 4, 6, 1, 7, 8, 9, 2, 5]
    v = [3 13; -1 0; 16 9; 21 1; -5 -15; 9 0; 10 4; 5 -4; 0 3]
    A = [-13   4
           4  13
           8   5
         -15   4
           8 -13]
    b = [13, 181, 173, 15, 155]

    V1 = vrep(v)
    p = polyhedron(V1, lib)
    inequality_fulltest(p, A, b, BitSet([]))

    V2 = vrep(v[ind, :])
    p = polyhedron(V2, lib)
    inequality_fulltest(p, A, b, BitSet([]))
end


text:
module Lisp

using Tokenize
using ReplMaker: initrepl

export tokenize, read_sexp, Token

const OP = Tokens.OP

const LPAREN = Tokens.LPAREN
const RPAREN = Tokens.RPAREN

const LSQUARE = Tokens.LSQUARE  # [
const RSQUARE = Tokens.RSQUARE  # ]

const AT_SIGN = Tokens.AT_SIGN

const INTEGER = Tokens.INTEGER
const FLOAT = Tokens.FLOAT
const STRING = Tokens.STRING
const CHAR = Tokens.CHAR

const LITERAL_KINDS = Base.IdSet{Tokens.Kind}([
    INTEGER,
    FLOAT,
    STRING,
    CHAR,
])

isliteral(t) = Tokens.kind(t) ∈ LITERAL_KINDS

const WHITESPACE = Tokens.WHITESPACE
const EQ = Tokens.EQ # =
const PLUS_EQ = Tokens.PLUS_EQ # +=
const MINUS_EQ = Tokens.MINUS_EQ # -=
const STAR_EQ = Tokens.STAR_EQ # *=
const FWD_SLASH_EQ = Tokens.FWD_SLASH_EQ # /=
const FWDFWD_SLASH_EQ = Tokens.FWDFWD_SLASH_EQ # //=
const OR_EQ = Tokens.OR_EQ # |=
const CIRCUMFLEX_EQ = Tokens.CIRCUMFLEX_EQ # ^=
const DIVISION_EQ = Tokens.DIVISION_EQ # ÷=
const REM_EQ = Tokens.REM_EQ # %=
const LBITSHIFT_EQ = Tokens.LBITSHIFT_EQ # <<=
const RBITSHIFT_EQ = Tokens.RBITSHIFT_EQ # >>=
const UNSIGNED_BITSHIFT_EQ = Tokens.UNSIGNED_BITSHIFT_EQ # >>>=
const BACKSLASH_EQ = Tokens.BACKSLASH_EQ # \=
const AND_EQ = Tokens.AND_EQ # &=
const COLON_EQ = Tokens.COLON_EQ # :=
const APPROX = Tokens.APPROX # ~
const EX_OR_EQ = Tokens.EX_OR_EQ # $=
const XOR_EQ = Tokens.XOR_EQ # ⊻=
const LAZY_AND = Tokens.LAZY_AND # &&
const LAZY_OR = Tokens.LAZY_OR # ||

const OPEQUAL_EXACTKINDS = Base.IdSet{Tokens.Kind}([
    PLUS_EQ, # +=
    MINUS_EQ, # -=
    STAR_EQ, # *=
    FWD_SLASH_EQ, # /=
    FWDFWD_SLASH_EQ, # //=
    OR_EQ, # |=
    CIRCUMFLEX_EQ, # ^=
    DIVISION_EQ, # ÷=
    REM_EQ, # %=
    LBITSHIFT_EQ, # <<=
    RBITSHIFT_EQ, # >>=
    UNSIGNED_BITSHIFT_EQ, # >>>=
    BACKSLASH_EQ, # \=
    AND_EQ, # &=
    COLON_EQ, # :=
    EX_OR_EQ, # $=
    XOR_EQ, # ⊻=
])
struct Token
    value
    token
    function Token(t::Symbol)
        new(t, nothing)
    end
    function Token(t)
        t_str = untokenize(t)
        if isliteral(t)
            t_value = Meta.parse(t_str)
        else
            t_value = Symbol(t_str)
        end
        new(t_value, t)
    end
end

value(t::Token) = t.value
token(t::Token) = t.token
kind(t::Token) = Tokens.kind(t.token)
exactkind(t::Token) = Tokens.exactkind(t.token)
startpos(t::Token) = Tokens.startpos(t.token)
isopequal(t::Token) = exactkind(t) ∈ OPEQUAL_EXACTKINDS
islazy(t::Token) = exactkind(t) === LAZY_AND || exactkind(t) === LAZY_OR

Base.show(io::IO, t::Token) = show(io, value(t))

"""
Chunk string `s` into sections delimited by parentheses and brackets.
"""
function read_sexp(s)
    forms = []
    form_stack = []
    current_form = forms
    kind_stack = Tokens.Kind[]
    for t in tokenize(s)
        if Tokens.kind(t) === LPAREN
            push!(kind_stack, LPAREN)
            push!(current_form, [])
            push!(form_stack, current_form)
            current_form = current_form[end]
            continue
        elseif Tokens.kind(t) === RPAREN
            top_kind = pop!(kind_stack)
            if top_kind !== LPAREN
                (line, col) = Tokens.startpos(t)
                msg = """
                    Encountered unexpected '$(untokenize(t))' at
                    line $(line) and column $(col)"""
                throw(DomainError(s, msg))
            end
            current_form = pop!(form_stack)
        elseif Tokens.kind(t) === LSQUARE
            push!(kind_stack, LSQUARE)
            push!(current_form, Any[Token(:ref)])
            push!(form_stack, current_form)
            current_form = current_form[end]
            continue
        elseif Tokens.kind(t) === RSQUARE
            top_kind = pop!(kind_stack)
            if top_kind !== LSQUARE
                (line, col) = Tokens.startpos(t)
                msg = """
                    Encountered unexpected '$(untokenize(t))' at
                    line $(line) and column $(col)"""
                throw(DomainError(s, msg))
            end
            current_form = pop!(form_stack)
        elseif Tokens.kind(t) === WHITESPACE
            continue
        else
            push!(current_form, Token(t))
        end

    end
    # strip empty form from trailing ENDMARKER token
    forms[1:end-1]
end

################################################################
# forms to julia expr

export parse_sexp, parse_function, parse_op

function parse_sexp(form::Vector)
    @show "form" form
    if length(form) == 0
        return form
    elseif form[1] isa Vector
        return parse_call(form)
    end

    head = value(form[1])
    if head === :function
        return parse_function(form)
    elseif head === :using
        return parse_using(form)
    elseif head === :while
        return parse_while(form)
    elseif head === :for
        return parse_for(form)
    elseif head === :ref
        return parse_ref(form)
    elseif head === :return
        return parse_return(form)
    elseif head === :module || head === :baremodule
        return parse_module(form)
    elseif head === :struct || head === :mutable
        return parse_struct(form)
    elseif head === :do
        return parse_do(form)
    elseif head === :.
        return parse_dotcall(form)
    elseif head === :(=)
        return parse_equal(form)
    elseif head === :(:)
        return parse_range(form)
    elseif isopequal(form[1])
        return parse_opequal(form)
    elseif islazy(form[1])
        return parse_lazyop(form)
    elseif kind(form[1]) === OP
        return parse_op(form)
    elseif kind(form[1]) === AT_SIGN
        return parse_macrocall(form)
    else  # must be a function call
        return parse_call(form)
    end
end

function parse_sexp(form::Token)
    if value(form) === :break
        return Expr(:break)
    else
        return value(form)
    end
end

"""
Parse

    (function (f <formals>) body<...>)

into

    function f(<formals>)
        body
        <...>
    end

"""
function parse_function(form::Vector)
    @assert value(form[1]) === :function
    @assert form[2] isa Vector
    @assert length(form) > 2
    _make_function(form)
end

function _make_function(form::Vector)
    Expr(
        value(form[1]),
        Expr(:call, parse_sexp.(form[2])...),
        Expr(:block, parse_sexp.(form[3:end])...))
end

function parse_using(form::Vector)
    @assert value(form[1]) === :using
    @assert length(form) > 1

    # TODO support `using Mod: name1`
    # TODO support `using ..Mod` for submodules
    make_module(sym) = Expr(:., sym)

    modules = make_module.(value.(form[2:end]))
    Expr(value(form[1]), modules...)
end

"""
Parse

    (f <args...>)

into

    f(<args...>)

"""
function parse_call(form::Vector)
    head = parse_sexp(form[1])
    args = [parse_sexp(f) for f in form[2:end]]
    Expr(:call, head, args...)
end

"""
Parse

    (@f <args...>)

into

    @f(<args...>)

"""
function parse_macrocall(form::Vector)
    @assert kind(form[1]) === AT_SIGN
    @assert length(form) > 1
    head = Symbol(string(value.(form[1:2])...))
    args = parse_sexp.(form[3:end])
    # LineNumberNode is required for correct parsing
    e = Expr(:macrocall, head, LineNumberNode(1, :none), args...)
    @show head args e
    e
end

"""
Parse

    [f <indices...>]

into

    f[<indices...>]

and

    [. x y <attributes...>]

into

    x.y<.<attributes...>>

and

    [:: x y <types...>]

into

    x::y<::<types...>>

"""
function parse_ref(form::Vector)
    head = parse_sexp(form[2])
    if head === :.
        function make_getfield(x, y)
            return Expr(:., x, QuoteNode(y))
        end
        return foldl(make_getfield, parse_sexp.(form[3:end]))
    elseif head === :(::)
        function make_annotation(x, y)
            return Expr(:(::), x, y)
        end
        return foldl(make_annotation, parse_sexp.(form[3:end]))
    else  # array indexing
        indices = [parse_sexp(f) for f in form[3:end]]
        return Expr(:ref, head, indices...)
    end
end

"""
Parse

    (. f <args...>)

into

    f.(<args...>)

"""
function parse_dotcall(form::Vector)
    @assert length(form) > 1
    args = [parse_sexp(f) for f in form[3:end]]
    Expr(value(form[1]), value(form[2]), Expr(:tuple, args...))
end

function parse_op(form::Vector)
    @assert kind(form[1]) === OP
    @assert length(form) > 1

    call_op(xs...) = Expr(:call, value(form[1]), xs...)

    if length(form) == 2
        # unary
        arg = parse_sexp(form[2])
        return call_op(parse_sexp(form[2]))
    else
        # TODO handle ternary operator
        # TODO handle right associative
        # binary, so fold
        args = [parse_sexp(f) for f in form[2:end]]
        return foldr(call_op, args)
    end
end

function parse_equal(form)
    @assert value(form[1]) === :(=)
    @assert length(form) > 2
    @show "equal" form

    if (form[2] isa Vector) && (value(form[2][1]) !== :ref)
        # special case of function definition
        _make_function(form)
    else
        @assert length(form) == 3
        destination = parse_sexp(form[2])
        assignment_value = parse_sexp(form[3])
        return Expr(value(form[1]), destination, assignment_value)
    end
end

parse_range(form::Vector) = parse_call(form)

function parse_opequal(form::Vector)
    @assert isopequal(form[1])
    @assert length(form) == 3
    op = value(form[1])
    lhs = value(form[2])
    rhs = parse_sexp(form[3])
    Expr(op, lhs, rhs)
end

function parse_lazyop(form::Vector)
    @assert islazy(form[1])
    @assert length(form) > 2
    op = value(form[1])
    call_op(lhs, rhs) = Expr(op, lhs, rhs)
    args = [parse_sexp(f) for f in form[2:end]]
    return foldr(call_op, args)
end

function parse_while(form::Vector)
    @assert value(form[1]) === :while
    @assert length(form) > 2
    condition = parse_sexp(form[2])
    body = [parse_sexp(f) for f in form[3:end]]
    Expr(:while, condition, Expr(:block, body...))
end

"""
Convert

    (for (i itr) body)

into

    for i = itr
        body
    end

and

    (for ((i itr1) (j itr2)) body)

into

    for i = itr1, j = itr2
        body
    end

"""
function parse_for(form::Vector)
    @assert value(form[1]) === :for
    @assert length(form) > 2

    make_itr(form) = Expr(:(=), value(form[1]), parse_sexp(form[2]))

    head = value(form[1])
    body = parse_sexp.(form[3:end])
    if form[2][1] isa Vector
        # multiple iterators
        itrs = make_itr.(form[2])
        return Expr(head, Expr(:block, itrs...), Expr(:block, body...))
    else
        # single iterator
        itr = make_itr(form[2])
        return Expr(head, itr, Expr(:block, body...))
    end
end

function parse_return(form::Vector)
    @assert value(form[1]) === :return
    if length(form) == 2
        ret = parse_sexp(form[2])
        Expr(value(form[1]), ret)
    else  # length(form) > 2
        ret = Expr(:tuple, parse_sexp.(form[2:end])...)
        Expr(value(form[1]), ret)
    end
end

function parse_module(form::Vector)
    @assert value(form[1]) === :module || value(form[1]) === :baremodule
    @assert length(form) > 1
    function make_module(form, import_base)
        body = parse_sexp.(form[3:end])
        Expr(:module, import_base, value(form[2]), Expr(:block, body...))
    end

    if value(form[1]) === :module
        make_module(form, true)
    else  # baremodule
        make_module(form, false)
    end
end

function parse_struct(form::Vector)
    head = value(form[1])
    @assert (
        head === :struct || (head === :mutable && value(form[2]) === :struct))
    @assert length(form) > 1

    function make_struct(form, is_mutable, name, body)
        Expr(:struct, is_mutable, name, Expr(:block, body...))
    end

    is_mutable = head === :mutable
    name_index = is_mutable ? 3 : 2
    name = value(form[name_index])
    body = parse_sexp.(form[(name_index + 1):end])
    return make_struct(form, is_mutable, name, body)
end

"""
Parse

    (do (f <args...>) (<formals...>) <body...>)

into

    do f(<args...>) <formals...>
        <body...>
    end

"""
function parse_do(form::Vector)
    head = value(form[1])
    @assert head === :do
    @assert length(form) > 3

    func = parse_sexp(form[2])
    formals = Expr(:tuple, value.(form[3])...)
    body = parse_sexp.(form[4:end])
    Expr(head, func, Expr(:->, formals, Expr(:block, body...)))
end

################################################################
# loading code
export include_stringl, includel
function include_stringl(mod, str)
    for s in read_sexp(str)
        expr = parse_sexp(s)
        @show expr
        Base.eval(mod, expr)
    end
end

function includel(mod, filename)
    open(filename, "r") do source_file
        source = read(source_file, String)
        include_stringl(mod, source)
    end
end

################################################################
# utilities for testing
export dump_sexp, parse_sexp_string, enable_repl

dump_sexp(sexp_str) = dump(parse_sexp(read_sexp(sexp_str)[1]))
parse_sexp_string(sexp_str) = parse_sexp(read_sexp(sexp_str)[1])
function enable_repl()
    initrepl(parse_sexp_string;
        prompt_text="jlisp> ",
        prompt_color=:blue,
        start_key=')',
        mode_name="JLisp_mode",
    )
end

end # module


text:
export add_materials_to_physical_groups_by_color

"""
    add_materials_to_physical_groups_by_color(materials::Vector{Material})

Assign entities with RGBA color matching the color in a Material X to a physical group
named "Material: X".
"""
function add_materials_to_physical_groups_by_color(materials::Vector{Material})
    color_to_ent = get_entities_by_color()
    for mat in sort(materials, by = m -> m.name)
        r = Int32(mat.color.r.i)
        g = Int32(mat.color.g.i)
        b = Int32(mat.color.b.i)
        a = Int32(mat.color.alpha.i)
        rgba = (r,g,b,a)
        ents = color_to_ent[rgba]
        safe_add_physical_group("Material: "*mat.name, ents)
    end
    return nothing
end

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
cxx"""
#include <opencv2/videoio/videoio_c.h>
"""

for name in [
    :CAP_ANY,
    :CAP_VFW,
    :CAP_V4L,
    :CAP_V4L2,
    :CAP_FIREWARE,
    :CAP_FIREWIRE,
    :CAP_IEEE1394,
    :CAP_DC1394,
    :CAP_CMU1394,
    :CAP_QT,
    :CAP_UNICAP,
    :CAP_DSHOW,
    :CAP_PVAPI,
    :CAP_OPENNI,
    :CAP_OPENNI_ASUS,
    :CAP_ANDROID,
    :CAP_XIAPI,
    :CAP_AVFOUNDATION,
    :CAP_GIGANETIX,
    :CAP_MSMF,
    :CAP_WINRT,
    :CAP_INTELPERC,
    :CAP_OPENNI2,
    :CAP_OPENNI2_ASUS,
    :CAP_GPHOTO2,
    :CAP_GSTREAMER,
    :CAP_FFMPEG,
    :CAP_IMAGES,

    :CAP_PROP_POS_MSEC,
    :CAP_PROP_POS_FRAMES,
    :CAP_PROP_POS_AVI_RATIO,
    :CAP_PROP_FRAME_WIDTH,
    :CAP_PROP_FRAME_HEIGHT,
    :CAP_PROP_FPS,
    :CAP_PROP_FOURCC,
    :CAP_PROP_FRAME_COUNT,
    :CAP_PROP_FORMAT,
    :CAP_PROP_MODE,
    :CAP_PROP_BRIGHTNESS,
    :CAP_PROP_CONTRAST,
    :CAP_PROP_SATURATION,
    :CAP_PROP_HUE,
    :CAP_PROP_GAIN,
    :CAP_PROP_EXPOSURE,
    :CAP_PROP_CONVERT_RGB,
    :CAP_PROP_WHITE_BALANCE_BLUE_U,
    :CAP_PROP_RECTIFICATION,
    :CAP_PROP_MONOCHROME,
    :CAP_PROP_SHARPNESS,
    :CAP_PROP_AUTO_EXPOSURE,
    :CAP_PROP_GAMMA,
    :CAP_PROP_TEMPERATURE,
    :CAP_PROP_TRIGGER,
    :CAP_PROP_TRIGGER_DELAY,
    :CAP_PROP_WHITE_BALANCE_RED_V,
    :CAP_PROP_ZOOM,
    :CAP_PROP_FOCUS,
    :CAP_PROP_GUID,
    :CAP_PROP_ISO_SPEED,
    :CAP_PROP_BACKLIGHT,
    :CAP_PROP_PAN,
    :CAP_PROP_TILT,
    :CAP_PROP_ROLL,
    :CAP_PROP_IRIS,
    :CAP_PROP_SETTINGS,
    :CAP_PROP_BUFFERSIZE,
    :CAP_PROP_AUTOFOCUS
    ]
    ex = Expr(:macrocall, Symbol("@icxx_str"), string("cv::", name, ";"))
    @eval begin
        global const $name = $ex
        export $name
    end
end


text:
# simply whatever experiment is currently slated to run on calmip
include("test/controller.jl")

plot_all()


text:
"""
Utilities to evaluate players against one another.

Typically, between each training iteration, different players
that possibly depend on the current neural network
compete against a set of baselines.
"""
module Benchmark

using ..AlphaZero

using ProgressMeter
using Statistics: mean
using Base: @kwdef

"""
    Benchmark.Player

Abstract type to specify a player that can be featured in a benchmark duel.

Subtypes must implement the following functions:
  - `Benchmark.instantiate(player, nn)`: instantiate the player specification
      into an [`AbstractPlayer`](@ref) given a neural network
  - `Benchmark.name(player)`: return a `String` describing the player
"""
abstract type Player end

# instantiate(::Player, gspec, nn)
function instantiate end

# name(::Player) :: String
function name end

"""
    Evaluation

Abstract type for a benchmark item specification.
"""
abstract type Evaluation end

"""
    Single <: Evaluation

Evaluating a single player in a one-player game.
"""
@kwdef struct Single <: Evaluation
  player :: Player
  sim :: SimParams
end

"""
    Duel <: Evaluation

Evaluating a player by pitting it against a baseline player in a two-player game.
"""
@kwdef struct Duel <: Evaluation
  player :: Player
  baseline :: Player
  sim :: SimParams
end

name(s::Single) = name(s.player)

name(d::Duel) = "$(name(d.player)) / $(name(d.baseline))"

"""
    Benchmark.run(env::Env, duel::Benchmark.Evaluation, progress=nothing)

Run a benchmark duel and return a [`Report.Evaluation`](@ref).

If a `progress` is provided, `next!(progress)` is called
after each simulated game.
"""
function run end

function run(env::Env, eval::Evaluation, progress=nothing)
  net() = Network.copy(env.bestnn, on_gpu=eval.sim.use_gpu, test_mode=true)
  if isa(eval, Single)
    simulator = Simulator(net, record_trace) do net
      instantiate(eval.player, env.gspec, net)
    end
  else
    @assert isa(eval, Duel)
    simulator = Simulator(net, record_trace) do net
      player = instantiate(eval.player, env.gspec, net)
      baseline = instantiate(eval.baseline, env.gspec, net)
      return TwoPlayers(player, baseline)
    end
  end
  samples, elapsed = @timed simulate(
    simulator, env.gspec, eval.sim,
    game_simulated=(() -> next!(progress)))
  gamma = env.params.self_play.mcts.gamma
  rewards, redundancy = rewards_and_redundancy(samples, gamma=gamma)
  return Report.Evaluation(
    name(eval), mean(rewards), redundancy, rewards, nothing, elapsed)
end

#####
##### Statistics for games with ternary rewards
#####

struct TernaryOutcomeStatistics
  num_won  :: Int
  num_draw :: Int
  num_lost :: Int
end

function TernaryOutcomeStatistics(rewards::AbstractVector{<:Number})
  num_won  = count(==(1), rewards)
  num_draw = count(==(0), rewards)
  num_lost = count(==(-1), rewards)
  @assert num_won + num_draw + num_lost == length(rewards)
  return TernaryOutcomeStatistics(num_won, num_draw, num_lost)
end

function TernaryOutcomeStatistics(report::Report.Evaluation)
  return TernaryOutcomeStatistics(report.rewards)
end

#####
##### Standard players
#####

"""
    Benchmark.MctsRollouts(params) <: Benchmark.Player

Pure MCTS baseline that uses rollouts to evaluate new positions.

Argument `params` has type [`MctsParams`](@ref).
"""
struct MctsRollouts <: Player
  params :: MctsParams
end

name(p::MctsRollouts) = "MCTS ($(p.params.num_iters_per_turn) rollouts)"

function instantiate(p::MctsRollouts, gspec::AbstractGameSpec, nn)
  return MctsPlayer(gspec, MCTS.RolloutOracle(gspec), p.params)
end

"""
    Benchmark.Full(params) <: Benchmark.Player

Full AlphaZero player that combines MCTS with the learnt network.

Argument `params` has type [`MctsParams`](@ref).
"""
struct Full <: Player
  params :: MctsParams
end

name(::Full) = "AlphaZero"

function instantiate(p::Full, gspec::AbstractGameSpec, nn)
  return MctsPlayer(gspec, nn, p.params)
end

"""
    Benchmark.NetworkOnly(;τ=1.0) <: Benchmark.Player

Player that uses the policy output by the learnt network directly,
instead of relying on MCTS.
"""
@kwdef struct NetworkOnly <: Player
  τ :: Float64 = 1.0
end

name(::NetworkOnly) = "Network Only"

function instantiate(p::NetworkOnly, ::AbstractGameSpec, nn)
  player = NetworkPlayer(nn)
  return PlayerWithTemperature(player, ConstSchedule(p.τ))
end

"""
    Benchmark.MinMaxTS(;depth, τ=0.) <: Benchmark.Player

Minmax baseline, which relies on [`MinMax.Player`](@ref).
"""
@kwdef struct MinMaxTS <: Player
  depth :: Int
  amplify_rewards :: Bool
  τ :: Float64 = 0.
end

name(p::MinMaxTS) = "MinMax (depth $(p.depth))"

function instantiate(p::MinMaxTS, ::AbstractGameSpec, nn)
  return MinMax.Player(
    depth=p.depth, amplify_rewards=p.amplify_rewards, τ=p.τ)
end

end


text:
"""Importance sampling examples.

TODOs:
- Try other f's, distributions (see Distributions.jl).
- Try a simple adaptive importance sampling algorithm.
"""

using Distributions
using PyPlot
using Random

# Plots f, p, q, and fp in the domain [lb, ub].
#
# Args:
# - f_x: a Function.
# - p_x: a Distribution.
# - q_x: a Function, a Distribution, or a Vector of the types.
# - lb/ub: Floats.
function plot_fs(f_x, p_x, q_x, lb, ub; lw=0.75)
    xl = [x for x in lb:0.05:ub]

    pdf_p_xl = pdf.(p_x, xl)
    plt.plot(xl, pdf_p_xl, label=raw"$p$", lw=lw)

    q_x_l = (isa(q_x, Vector) ? q_x : [q_x])
    for (idx, q_x_i) in enumerate(q_x_l)
        pdf_q_xl = (isa(q_x_i, Function) ? q_x_i.(xl) : pdf.(q_x_i, xl))
        label = string(raw"$q$", (idx == 1 ? "" : string(idx)))
        plt.plot(xl, pdf_q_xl, label=label, lw=lw)
    end

    f_xl = f_x.(xl)
    plt.plot(xl, f_xl, label=raw"$f$", lw=lw)
    plt.plot(xl, f_xl .* pdf_p_xl, label=raw"$f \times p$", lw=lw)
    plt.legend();
end

# Plots incremental sampling estimates with a confidence region.
#
# Args:
# - n_estimates: a Vector of estimates of length n.
# - delta_n: compute incremental estimates every this many steps.
# - label: a String label to use for the plot.
function plot_estimates(n_estimates::Vector, delta_n::Integer, label::String)

    function _x_range(l)
        total_n = length(l)
        ret = delta_n:delta_n:total_n
        if (last(ret) != total_n); push!(ret, total_n); end
        return ret
    end

    function _n_mean(n_estimates, xl)
        ret = []
        # TODO(kykim): Do a cumsum type of optimization.
        for x in xl
            x_n_estimates = [mean(estimates[1:x]) for estimates in n_estimates]
            push!(ret, x_n_estimates)
        end
        return ret
    end

    xl = _x_range(first(n_estimates))
    n_y = _n_mean(n_estimates, xl)
    min_y, max_y = minimum.(n_y), maximum.(n_y)
    mid_y = (max_y + min_y) ./ 2.0

    # For Plots.jl:
    #  plot(x_mc, mid_y_mc, ribbon=(max_y_mc - mid_y_mc), fillalpha=0.15,
    #       label="MC", lw=2, xlabel="no. of samples", ylabel="estimates")
    #  plot(x_is, mid_y_is, ribbon=(max_y_is - mid_y_is), fillalpha=0.15, label="IS", lw=2)
    plt.plot(xl, mid_y, label=label, lw=0.5)
    plt.fill_between(xl, min_y, max_y, alpha=0.1)
end

# Computes standard weights.
function is_standard_weights(p_x, q_x, samples)
    weights = pdf.(p_x, samples) ./ pdf.(q_x, samples)
    return weights
end

# Computes mixture weights based on all proposals.
function is_mixture_weights(p_x, q_x_l, samples)
    q_x_l = (isa(q_x_l, Vector) ? q_x_l : [q_x_l])
    denom = .+([pdf.(q_x_i, samples) for q_x_i in q_x_l]...) / length(q_x_l)
    weights = pdf.(p_x, samples) ./ denom
    return weights
end

# Simple experiment comparing Gaussian nominal and propsal distributions.
#
# Importance sampling with the proposal used leads to slightly smaller variance
# as it is better configured w.r.t to f * p.
function simple_exp(mc_n, is_n, delta_n, n_trials)
    Random.seed!(0)

    f_x = (x) -> (1. / (1. + exp(-(x - 4.5))))
    p_x = Normal(2.0, 1.0)
    # Note. Using a smaller σ (e.g., 0.5) leads to higher variance as q now has
    # lighter tails than p. This case of using a Gaussian proposal when the
    # nominal is also Gaussian is discussed in Owen 9.1.
    q_x = Normal(3.0, 0.50)  # Compare σ=1.0 vs. σ=0.5

    # Plot f, p, fp, q.
    figure(figsize=(9.0, 6.0))
    plot_fs(f_x, p_x, q_x, -3.0, 9.0)
    plt.tight_layout()
    # plt.savefig("distribs.png", dpi=500)

    # Compute n_trials many estimates.
    n_mc_estimates, n_is_estimates = [], []
    for _ in 1:n_trials
        mc_estimates = f_x.(rand(p_x, mc_n))
        push!(n_mc_estimates, mc_estimates)

        is_samples = rand(q_x, is_n)
        pdf_p, pdf_q = pdf.(p_x, is_samples), pdf.(q_x, is_samples)
        is_estimates = f_x.(is_samples) .* (pdf_p ./ pdf_q)
        push!(n_is_estimates, is_estimates)
    end

    # Plot MC and IS estimates w/ confidence regions.
    figure(figsize=(9.0, 6.0))
    plot_estimates(n_mc_estimates, delta_n, "MC")
    plot_estimates(n_is_estimates, delta_n, "IS")
    plt.xlabel("samples"); plt.ylabel("estimates"); plt.legend();
    plt.tight_layout()
    # plt.savefig("estimates.png", dpi=500)
end

# This is an example given in the chapter 6 importance sampling note.
#
# The exponential proposal distribution is so well-suited for the given f * p
# that importance sampling done with the distribution leads to extremely small
# variance.
#
# An alternate importance sampling with a "better" Gaussian also works well.
function normal_exponential_exp(mc_n, is_n, delta_n, n_trials)
    Random.seed!(0)

    f_x = (x) -> (x >= 4 ? 1.0 : 0.0)
    p_x = Normal(2.5, 1.0)

    # Using this to sample from an exponential shifted to the right by 4.
    q_x_tmp = Exponential(1.0)
    shift = 4.0
    pdf_q_fn = (x) -> pdf(q_x_tmp, x - shift)
    is_samples1 = rand(q_x_tmp, is_n) .+ shift

    q_x2 = Normal(3.5, 1.0)
    is_samples2 = rand(q_x2, is_n)

    # Plot f, p, fp, q.
    plt.figure(figsize=(9.0, 6.0))
    plot_fs(f_x, p_x, [pdf_q_fn, q_x2], -3.0, 9.0)

    # Compute n_trials many estimates.
    n_mc_estimates, n_is_estimates1, n_is_estimates2 = [], [], []
    for _ in 1:n_trials
        mc_estimates = f_x.(rand(p_x, mc_n))
        push!(n_mc_estimates, mc_estimates)

        pdf_p1, pdf_q1 = pdf.(p_x, is_samples1), pdf_q_fn.(is_samples1)
        is_estimates1 = f_x.(is_samples1) .* (pdf_p1 ./ pdf_q1)
        push!(n_is_estimates1, is_estimates1)

        pdf_p2, pdf_q2 = pdf.(p_x, is_samples2), pdf.(q_x2, is_samples2)
        is_estimates2 = f_x.(is_samples2) .* (pdf_p2 ./ pdf_q2)
        push!(n_is_estimates2, is_estimates2)
    end

    # Plot MC and IS estimates w/ confidence regions.
    plt.figure(figsize=(9.0, 6.0))
    plot_estimates(n_mc_estimates, delta_n, "MC")
    for (idx, n_is_estimates) in enumerate([n_is_estimates1, n_is_estimates2])
        label = string("IS", (idx == 1 ? "" : string(idx)))
        plot_estimates(n_is_estimates, delta_n, label)
    end
    plt.xlabel("samples"); plt.ylabel("estimates"); plt.legend();
end

# Experiment to demonstrate a simple multiple importance sampling use case.
#
# The f and p in the experiment are such that when a Gaussian proposal with a
# smaller variance is used for importance sampling, it can lead to a higher
# variance in estimate than the Monte Carlo estimate. This experiment
# demonstrates that using multiple importance sampling can be an option in such
# cases.
function exp_gaussian_mis(mc_n, is_n, delta_n, n_trials, weight_type="mixture")
    Random.seed!(0)

    f_x = (x) -> (1. / (1. + exp(-(x - 4.5))))
    p_x = Normal(2.0, 1.0)
    q_x1 = Normal(3.0, 0.75)  # Compare σ=0.75 vs. σ=0.5
    q_x2 = Normal(2.25, 0.75)
    q_x3 = Normal(3.75, 0.75)
    q_x_l = [q_x1, q_x2, q_x3]
    is_n_i = is_n ÷ length(q_x_l)  # Draw an equal no. of samples for each.

    # Plot f, p, fp, q.
    plt.figure(figsize=(9.0, 6.0))
    plot_fs(f_x, p_x, [q_x1, q_x2, q_x3], -3.0, 9.0)

    # Compute n_trials many estimates.
    n_mc_estimates, n_is_estimates = [], []
    for _ in 1:n_trials
        mc_estimates = f_x.(rand(p_x, mc_n))
        push!(n_mc_estimates, mc_estimates)

        is_estimates = []
        for q_x_i in q_x_l
            is_samples = rand(q_x_i, is_n_i)
            if weight_type == "standard"
                weights = is_standard_weights(p_x, q_x_i, is_samples)
            elseif weight_type == "mixture"
                weights = is_mixture_weights(p_x, q_x_l, is_samples)
            end
            append!(is_estimates, f_x.(is_samples) .* weights)
        end
        shuffle!(is_estimates)
        push!(n_is_estimates, is_estimates)
    end

    # Plot MC and IS estimates w/ confidence regions.
    plt.figure(figsize=(9.0, 6.0))
    plot_estimates(n_mc_estimates, delta_n, "MC")
    plot_estimates(n_is_estimates, delta_n, "IS")
    plt.xlabel("samples"); plt.ylabel("estimates"); plt.legend();
    plt.tight_layout()
    # plt.savefig("myplot.png", dpi=500)
end

# Experiment to compare a single Gaussian proposal and multiple proposals.
#
# The nominal distribution is a multimodal distribution designed in such a way
# that also leads to a multimodal f * p. Importance sampling with a single
# Gaussian can work reasonably well when configured properly to cover the
# regions of f * p, but the multimodal nature makes this somewhat tricky.
# Instead, importance sampling done with multiple Gaussians makes it easier to
# achieve the desired "coverage" and leads to lower variance.
function exp_mixture_gaussian_mis(mc_n, is_n, delta_n, n_trials,
                                  weight_type="mixture")
    Random.seed!(0)

    f_x = (x) -> (1. / (1. + exp(-(x - 1.5))))
    p_x = MixtureModel(
        Normal[Normal(-2.0, 1.0), Normal(1.0, 0.50), Normal(3.0, 0.5)],
        [0.6, 0.25, 0.15])

    # Compare a single Guassian proposal vs. multiple proposals.
    q_x = Normal(0.0, 1.5)
    q_x_mis1 = Normal(3.0, 0.75)
    q_x_mis2 = Normal(1.0, 1.0)
    q_x_mis3 = Normal(-1.0, 1.5)
    q_x_misl = [q_x_mis1, q_x_mis2, q_x_mis3]
    is_n_i = is_n ÷ length(q_x_misl)  # Draw an equal no. of samples for each.

    # Plot f, p, fp, q.
    plt.figure(figsize=(9.0, 6.0))
    plot_fs(f_x, p_x, [q_x, q_x_misl...], -6.0, 7.0)

    # Compute n_trials many estimates.
    n_mc_estimates, n_is_estimates, n_mis_estimates = [], [], []
    for _ in 1:n_trials
        mc_estimates = f_x.(rand(p_x, mc_n))
        push!(n_mc_estimates, mc_estimates)

        is_samples = rand(q_x, is_n)
        pdf_p, pdf_q = pdf.(p_x, is_samples), pdf.(q_x, is_samples)
        is_estimates = f_x.(is_samples) .* (pdf_p ./ pdf_q)
        push!(n_is_estimates, is_estimates)

        mis_estimates = []
        for q_x_i in q_x_misl
            mis_samples = rand(q_x_i, is_n_i)
            if weight_type == "standard"
                weights = is_standard_weights(p_x, q_x_i, mis_samples)
            elseif weight_type == "mixture"
                weights = is_mixture_weights(p_x, q_x_misl, mis_samples)
            end
            append!(mis_estimates, f_x.(mis_samples) .* weights)
        end
        shuffle!(mis_estimates)
        push!(n_mis_estimates, mis_estimates)
    end

    # Plot MC and IS estimates w/ confidence regions.
    plt.figure(figsize=(9.0, 6.0))
    plot_estimates(n_mc_estimates, delta_n, "MC")
    plot_estimates(n_is_estimates, delta_n, "IS")
    plot_estimates(n_mis_estimates, delta_n, "MIS")
    plt.xlabel("no. samples"); plt.ylabel("estimates"); plt.legend();
    plt.tight_layout()
    # plt.savefig("myplot.png", dpi=500)
end

simple_exp(10_000, 10_000, 100, 5)

# normal_exponential_exp(10_000, 10_000, 100, 5)

# Compare the "standard" and "mixture" weightings. Mixture weights does
# noticeably better presumably because it is more robust against small q which
# can significantly increase variance.
# exp_gaussian_mis(9000, 9000, 100, 5, "mixture")

# exp_mixture_gaussian_mis(9000, 9000, 100, 5, "mixture")


text:
WandererSymbol(x...) = Symbol(join(["Wanderer", x...], "."))

ARG = WandererSymbol("ARG") # we just need limited mangled symbols here.
TYPE_ROOT = WandererSymbol("TYPE_ROOT")

IN_TYPES = WandererSymbol("IN", "TYPES")
IN_FIELDS = WandererSymbol("IN", "FIELDS")
IN_SOURCE = WandererSymbol("IN", "SOURCE")

OUT_TYPES = WandererSymbol("OUT", "TYPES")
OUT_FIELDS = WandererSymbol("Wanderer", "OUT.FIELDS")
OUT_SOURCE = WandererSymbol("Wanderer", "OUT.SOURCE")
RECORD = WandererSymbol("RECORD")

N = WandererSymbol("N")
GROUPS = WandererSymbol("GROUPS")
GROUP_KEY = WandererSymbol("GROUP_KEY")

FN = WandererSymbol("FN")
FN_RETURN_TYPES = WandererSymbol("FN", "RETURN_TYPES")
FN_OUT_FIELDS = WandererSymbol("FN", "OUT_FIELDS")

AGG = WandererSymbol("AGG")
AGG_TYPES = WandererSymbol("AGG", "TYPES")

_gen_sym_count = 0
function gen_sym()
    global _gen_sym_count
    let sym = WandererSymbol("TMP", _gen_sym_count)
        _gen_sym_count = _gen_sym_count  + 1
        sym
    end
end

function gen_sym(a :: Union{Symbol, Int, String})
    global _gen_sym_count
    WandererSymbol("Symbol", a)
end


text:
# FASTA Writer
# ============

"""
    FASTA.Writer(output::IO; width=70)

Create a data writer of the FASTA file format.

# Arguments
* `output`: data sink
* `width=70`: wrapping width of sequence characters
"""
struct Writer{S <: TranscodingStream} <: BioGenerics.IO.AbstractWriter
    output::S
    # maximum sequence width (no limit when width ≤ 0)
    width::Int
end

function BioGenerics.IO.stream(writer::Writer)
    return writer.output
end

# This constructor should not be here - previously, constructing a Writer
# unintentionally used the default constructor instead of the manually crafted one.
# This allowed a signature that should not have been allowed, so we keep it here.
Writer(output::IO; width::Integer=70) = Writer(output, width)

function Writer(output::IO, width::Integer)
    if output isa TranscodingStream
        return Writer{typeof(output)}(output, width)
    else
        stream = TranscodingStreams.NoopStream(output)
        return Writer{typeof(stream)}(stream, width)
    end
end

function Base.write(writer::Writer, record::Record)
    checkfilled(record)
    output = writer.output
    n::Int = 0
    if writer.width ≤ 0
        n += write(output, record, '\n')
    else
        headerlen = hasdescription(record) ? last(record.description) : last(record.identifier)
        n += unsafe_write(output, pointer(record.data), headerlen)
        n += write(output, '\n')
        p = pointer(record.data, first(record.sequence))
        p_end = pointer(record.data, last(record.sequence))
        while p ≤ p_end
            w = min(writer.width, p_end - p + 1)
            n += unsafe_write(output, p, w)
            n += write(output, '\n')
            p += w
        end
    end
    return n
end


text:
# using Random, Distributions
# using Parameters
# using Statistics
# using Plots
# using LinearAlgebra
# using Latexify
# using PrettyTables
# using Interpolations
# using Dierckx
# using ForwardDiff
# using Optim
#      using Optim: converged, maximum, maximizer, minimizer, iterations
# using Roots
# include("Scaled_Interpolation_Functions.jl")

cd("/Volumes/SSD Hans/Github/MacroFall2020/AdvMacro/Assignment5")
# mkpath("Figures")
using Plots
# using LateXStrings # Pkg.add("LaTeXStrings") # https://github.com/stevengj/LaTeXStrings.jl
using Dierckx # Pkg.add("Dierckx") # https://github.com/kbarbary/Dierckx.jl
using Interpolations # Pkg.add("Interpolations") # https://github.com/JuliaMath/Interpolations.jl
using ForwardDiff # Pkg.add("ForwardDiff") # https://github.com/JuliaDiff/ForwardDiff.jl
using Optim # Pkg.add("Optim") # https://julianlsolvers.github.io/Optim.jl/stable/
    using Optim: converged, maximum, maximizer, minimizer, iterations
using Roots # Pkg.add("Roots") # https://github.com/JuliaMath/Roots.jl
using Parameters # Pkg.add("Parameters") # https://github.com/mauro3/Parameters.jl
using Distributions #Pkg.add("Distributions")
using QuadGK # Pkg.add("QuadGK") # https://juliamath.github.io/QuadGK.jl/latest/
using LinearAlgebra
using Random
using Statistics
using Latexify
Pkg.add("StatsPlots")
using StatsPlots
# Call Scaled Interpolation Functions
    include("../Scaled_Interpolation_Functions.jl")

# Set seed
Random.seed!(0021)

# Parameters
    # Generate structure for parameters using Parameters module
    # Set default values for parameters
    @with_kw struct Par
        # Model Parameters
        z_bar::Float64 = 1; # Reference level for productivity
        α::Float64 = 1/3  ; # Production function
        β::Float64 = 0.98 ; # Discount factor
        σ::Float64 = 2 ; # consumption elasticity of subsitution
        η::Float64 = 1 ; # labor/leisure elasticity of substitution
        δ::Float64 = 0.1 ; # Depreciation rate of capital
        ρ::Float64 = 0.8 ; # Persistance of AR(1) productivity process: log(z')=ρlog(z)+η
        σ_η::Float64 = 0.1 ; # Variance of innovations for productivity where η ∼ N(0,σ_η)
        # VFI Paramters
        max_iter::Int64   = 2000  ; # Maximum number of iterations
        dist_tol::Float64 = 1E-9  ; # Tolerance for distance between current and previous value functions
        # Policy functions
        H_tol::Float64    = 1E-9  ; # Tolerance for policy function iteration
        N_H::Int64        = 20    ; # Maximum number of policy iterations
        # Minimum consumption for numerical optimization
        c_min::Float64    = 1E-16
    end
# Allocate parameters to object p for future calling
p = Par()
global gam = ((p.α*p.z_bar*p.β)/(1-(1-p.δ)*p.β))^(1/(1-p.α)) # Some constant

# Steady state values
function SS_values(p::Par)
    # This function takes in parameters and provides steady state values
    # Parameters: productivity (z), returns to scale (α), discount factor (β), rate of capital depreciation (δ)
    #             consumption elasticity of substitution (σ), labor/leisure elasticity of subsitution (η)
    # Output: values for capital, labor, production, consumption, rental rate, wage
    @unpack z_bar,α,β,δ = p
    l_ss = 0.4 # Labor
    k_ss = gam*l_ss # Capital
    y_ss = z_bar*(k_ss^α)*(l_ss^(1-α)) # Output
    c_ss = y_ss-δ*k_ss # Consumption
    w_ss = (1-α)*z_bar*(k_ss^α)*(l_ss^(-α)) # Wage = marginal product of labor
    r_ss = α*z_bar*(k_ss^(α-1))*(l_ss^(1-α)) # Rental rate of capital = marginal product of capital
    return k_ss,y_ss,c_ss,r_ss,w_ss,l_ss
end
# Test steady state function
k_ss,y_ss,c_ss,r_ss,w_ss,l_ss = SS_values(p)
println(" ")
println("------------------------")
println(" Steady State variables")
println("   Quantities: k = $k_ss; y = $y_ss; c = $c_ss;")
println("   Prices:     r = $r_ss; w = $w_ss;")
println("------------------------")
println(" ")
# Get χ such that steady state labor = 0.4
function get_chi(p::Par,l_ss,c_ss,k_ss)
    @unpack z_bar, α, β, δ, σ, η = p
    chi = (c_ss^(-σ))*z_bar*(1-α)*(k_ss^α)*(l_ss^(-α-η))
    return chi
end
global χ = get_chi(p,l_ss,c_ss,k_ss)

# Function to make grid for capital (first state var)
function Make_K_Grid(n_k,θ_k,p::Par)
    # Get SS
    k_ss,y_ss,c_ss,r_ss,w_ss,l_ss = SS_values(p)
    # Lower and upper bounds
    lb = 1E-5
    ub = 2*k_ss
    # Get k_grid
    if θ_k≠1
        k_grid = PolyRange(lb,ub;θ=θ_k,N=n_k)
    else
    k_grid = range(lb,ub,length=n_k)
    end
    # Return
    return k_grid
end

# Function that returns the percentage error in the euler equation
function Euler_Error(k,z,kp,kpp,l,lp,p::Par)
    # Return percentage error in Euler equation
    @unpack α, β, σ, δ = p
    LHS = (z.*(k.^α).*(l.^(1-α)).+(1-δ).*k.-kp).^(-σ)
    RHS = β.*(α.*z.*((lp./kp).^(1-α)).+(1-δ)).*((z.*(kp.^α).*(lp.^(1-α)).+(1-δ).*kp.-kpp).^(-σ))
    return real((RHS./LHS.-1).*100)
end

# Period utility function (no penalty)
function utility(k,z,kp,l,p::Par)
    @unpack α,δ,σ,η,c_min = p
    u_c = 0; # intialize utility from consumption
    # Utility of consumption
    c = z*(k^α)*(l^(1-α))+(1-δ)*k-kp # Consumption from resource constraint
    if c<=c_min
        u_c = (c_min^(1-σ))/(1-σ)
    else
        u_c = (c^(1-σ))/(1-σ)
    end
    # Disutility of labor
    u_l = χ*((l^(1+η))/(1+η))
    return u_c-u_l
end
# Derivative of utility function wrt labor l
function d_utility_l(k,z,kp,l,p::Par)
    @unpack α,δ,σ,η,c_min = p
    c = z*(k^α)*(l^(1-α))+(1-δ)*k-kp
    d_u = 0
    if c>c_min
        d_u = c^(-σ)
    else
        d_u = c_min^(-σ)
    end
    return d_u*z*(k^α)*(1-α)*(l^(-α))-χ*(l^η)
end
# Derivative of utility function wrt capital k'
function d_utility_kp(k,z,kp,l,p::Par)
    @unpack α,δ,σ,η,c_min = p
    c = z*(k^α)*(l^(1-α))+(1-δ)*k-kp
    d_u = 0
    if c>c_min
        d_u = c^(-σ)
    else
        d_u = c_min^(-σ)
    end
    return -d_u
end

# Function to distretize AR(1) markov process with Rouwenhorst (1995)
function Rouwenhorst95(N,p::Par)
    @unpack ρ,σ_η=p
    # INPUTS:
        # ρ: persistence of unerlying AR(1) process where log(z') = ρlog(z)+η
        # σ_z: Std dev of inovation η in AR(1) process where η∼N(0,σ^2)
        # N: Size of grid for discrete process
    # OUTPUT:
        # z: All possible values of discretized AR(1) process, equally spaced grid of size N
        # Π: Matrix of transition probabilities
        # PDF_z: Stationary PDF of z
    #---------------------------------------------------------------------------
    Π = zeros(N,N)
    Π_Nm = zeros(N-1,N-1)
    P = (1+ρ)/2
    ϕ = σ_η*(sqrt((N-1)/(1-ρ^2)))
    z = range(-ϕ,ϕ;length=N)
    if N==2
        Π = [P 1-P;1-P P]
    else
        Π_Nm = Rouwenhorst95(N-1,p)[2]
        o = zeros(N-1)
        Π = P*[Π_Nm o; o' 0] + (1-P)*[o Π_Nm; 0 o'] + (1-P)*[o' 0; Π_Nm o] + P*[0 o';o Π_Nm]
        Π = Π./repeat(sum(Π,dims=2),1,N)
    end
    PDF_z = pdf.(Binomial(N-1,0.5),(0:N-1))
    return (z,Π,PDF_z)
end

# Generate structure of model objects
    @with_kw struct Model
        # Parameters
        p::Par = Par() # Model paramters
        # Grids
        θ_k::Float64    = 1     # Default Curvature of k_grid
        n_k::Int64      = 20    # Default Size of k_grid
        n_k_fine::Int64 = 1000  # Default Size of fine grid for interpolation
        n_z::Int64      = 10    # Default size of discretized grid for productivity as a markov process
        scale_type::Int64 = 1   # Default grid type (polynomial)
        k_grid          = Make_K_Grid(n_k,θ_k,p)    # k_grid for model solution
        k_grid_fine     = Make_K_Grid(n_k_fine,1,p) # Fine grid for interpolation
        # Value and policy functions
        V         = Array{Float64}(undef,n_k,n_z)   # Value Function
        G_kp      = Array{Float64}(undef,n_k,n_z)       # Policy Function for capital k'
        G_c       = Array{Float64}(undef,n_k,n_z)       # Policy Function for consumption c
        G_l       = Array{Float64}(undef,n_k,n_z)       # Policy Function for labor l
        V_fine    = Array{Float64}(undef,n_k_fine,n_z)  # Interpolated Value Function
        G_kp_fine = Array{Float64}(undef,n_k_fine,n_z)  # Interpolated Policy Function for capial k'
        G_c_fine  = Array{Float64}(undef,n_k_fine,n_z)  # Interpolated Policy Function for consumption c
        G_l_fine  = Array{Float64}(undef,n_k_fine,n_z)  # Interpolated Policy Function for labor l
        Euler     = Array{Float64}(undef,n_k_fine,n_z)  # Errors in Euler equation
        z_grid    = Array{Float64}(undef,n_z)       # discretized grid for productivity
        Π         = Array{Float64}(undef,n_z,n_z)   # Probability transition matrix for productivity
    end
    M = Model()

# Function that finds the fixed point of the value function, then interpolates between grid points
function VFI_fixedpoint(T::Function,M::Model)
            ### T : Bellman operator (interior loop) ###
            ### M : Model structure                  ###
    # Unpack model structure
    @unpack p,n_k,n_k_fine,θ_k,k_grid,k_grid_fine,n_z,V_fine,G_kp_fine,G_c_fine,G_l_fine,Euler,z_grid,Π = M
    # VFI paramters
    @unpack max_iter, dist_tol = p
    # Initialize variables for iteration
    V_old = zeros(n_k,n_z) # Value function
    V_dist = 1 # Distance between old and new value function
    iter = 1
    println(" ")
    println("------------------------")
    println("VFI - n_k=$n_k - grid curvature θ_k=$θ_k - n_z=$n_z")
    # Start VFI
    while iter <= max_iter
        # Update value function and policy functions
        a=1.0
        V_new = a*T(Model(M,V=copy(V_old)))[1] + (1-a)*V_old # Call Bellman operator which returns a new value function at each capital grid point
        # Update value function and distance between previous and current iteration
        V_dist = maximum(abs.(V_new./V_old.-1))
        V_old = V_new
        # Report progress
        # println("   VFI Loop: iter=$iter, dist=",100*V_dist," %")
        # Report progress every 100 iterations
        if mod(iter,100)==0
           println("   VFI Loop: iter=$iter, dist=",100*V_dist,"%")
        end
        # Check if converged
        if V_dist <= dist_tol
            V_new, G_kp, G_c, G_l = T(Model(M,V=copy(V_new)))
            println("VFI - n_k=$n_k - θ_k=$θ_k")
            println("Converged after $iter iterations with a distance of ",100*V_dist," %")
            println("------------------------")
            println(" ")
            # Interpolate to fine grid on capital using natural cubic spline if it converged
            for i in 1:n_z
                V_ip = ScaledInterpolations(M.k_grid,V_new[:,i], FritschButlandMonotonicInterpolation()) # Monotonic spline because I know that the value function is always increasing in capital
                    V_fine[:,i] = V_ip.(collect(M.k_grid_fine))
                G_kp_ip = ScaledInterpolations(M.k_grid,G_kp[:,i], BSpline(Cubic(Line(OnGrid()))))
                    G_kp_fine[:,i] = G_kp_ip.(collect(M.k_grid_fine))
                G_c_ip = ScaledInterpolations(M.k_grid,G_c[:,i], BSpline(Cubic(Line(OnGrid()))))
                    G_c_fine[:,i] = G_c_ip.(collect(M.k_grid_fine))
                G_l_ip = ScaledInterpolations(M.k_grid,G_l[:,i], BSpline(Cubic(Line(OnGrid()))))
                    G_l_fine[:,i] = G_l_ip.(collect(M.k_grid_fine))
                # Percent Euler Error on fine grid
                #Euler[:,i] = Euler_Error(M.k_grid_fine,G_kp_fine,G_kp_ip.(collect(G_kp_fine)),G_l_fine,G_l_ip.(collect(G_kp_fine)),p)
            end
            # Update model with solution
            M = Model(M; V=V_new,G_kp=G_kp,G_c=G_c,G_l=G_l,V_fine=V_fine,G_kp_fine=G_kp_fine,G_c_fine=G_c_fine,G_l_fine=G_l_fine,Euler=Euler)
            return M
        end
        # If it didn't converge, go to next iteration
        iter += 1
    end
    # If it didn't converge, return error
    error("Error in VFI - Solution not found")
end

# Bellman operator for the nested continuous choice of labor and capital tomorrow with capital and productivity as state variables
function T_nested_max(M::Model)
    @unpack p,n_k,k_grid,n_z,V,G_kp,G_c,G_l,z_grid,Π,z_grid = M
    @unpack β,α,δ,η,σ,c_min = p
    get_kp_max(k,l) = z*(k^α)*(l^(1-α)) + (1-δ)*k - c_min # Function because the max k' depends on the value of capital today k and labor l
    # Define boundaries on labor l
    l_min = 1E-16
    l_max = 1.0
    # Define boundaries for k'
    kp_min = 1.001*k_grid[1]
    get_kp_max(k,l,z) = z*(k^α)*(l^(1-α)) + (1-δ)*k - c_min # Function because the max k' depends on (k,l,z)
    # Function to get V_old
    Vp = [x->ScaledInterpolations(k_grid,x[:,i], BSpline(Cubic(Line(OnGrid())))) for i in 1:n_z]
    # Function that returns objective function for a given (z,k,k',l)
    function Obj_fn(k,z,kp,l,Π_z::Vector,p::Par)
        # Π_z: Vector of conditional probabilites for productivity next period given z today
        Emax = sum(Π_z[x]*Vp[x](V).(kp) for x in 1:n_z) # sum of z_j probability when starting at z_i : ∑_j π_ij V(kp,z_j)
        return -utility(k,z,kp,l,p) - β*Emax
    end
    # Function to get derivative of Emax wrt capital k'
    dVp(x,Π_z::Vector) = sum(Π_z[i]*ForwardDiff.derivative(Vp[i](V),x) for i in 1:n_z)
    # Function that returns derivative of objective function wrt k'
    d_Obj_fn_kp(k,z,kp,l,Π_z::Vector,p::Par) = d_utility_kp(k,z,kp,l,p) + β*dVp(kp,Π_z)
    # Derivative of objective function wrt labor l
    d_Obj_fn_l(k,z,kp,l,p::Par) = d_utility_l(k,z,kp,l,p)
    # Define function that finds optimal labor l given (k,z,k') and returns objective function conditional on optimal labor
    function Obj_fn_condl(k,z,kp,Π_z::Vector,p::Par)
        # Check for corner solutions on labor
        dobj_min = d_utility_l(k,z,kp,l_min,p)
        dobj_max = d_utility_l(k,z,kp,l_max,p)
        if dobj_min <= 0
            return Obj_fn(k,z,kp,l_min,Π_z,p),l_min
        elseif dobj_max >= 0
            return Obj_fn(k,z,kp,l_max,Π_z,p),l_max
        else
        # if no corner solutions, find interior solution
            min_result = optimize(x->d_utility_l(k,z,kp,x,p).^2,l_min,l_max,Brent())
            l = min_result.minimizer
            return Obj_fn(k,z,kp,l,Π_z,p),l
        end
    end
    # Outer loop for all possible values of productivity today
    for j in 1:n_z
        # Inner loop for each capital level in the grid
        for i in 1:n_k
            kp_max = min(get_kp_max(k_grid[i],1.0,z_grid[j]),k_grid[end])
            # Check for corner solutions on capital
            l_kp_min = Obj_fn_condl(k_grid[i],z_grid[j],kp_min,Π[j,:],p)[2]
            l_kp_max = Obj_fn_condl(k_grid[i],z_grid[j],kp_max,Π[j,:],p)[2]
            dobj_min = d_Obj_fn_kp(k_grid[i],z_grid[j],kp_min,l_kp_min,Π[j,:],p)
            dobj_max = d_Obj_fn_kp(k_grid[i],z_grid[j],kp_max,l_kp_max,Π[j,:],p)
            if dobj_min <= 0.0
                G_kp[i,j] = kp_min
                G_l[i,j] = l_kp_min
                V[i,j] = -Obj_fn(k_grid[i],z_grid[j],kp_min,l_kp_min,Π[j,:],p)
            elseif dobj_max >= 0.0
                G_kp[i,j] = kp_max
                G_l[i,j] = l_kp_max
                V[i,j] = -Obj_fn(k_grid[i],z_grid[j],kp_max,l_kp_max,Π[j,:],p)
            else
            # If no corner solution, find interior solution
                min_result = optimize(x->Obj_fn_condl(k_grid[i],z_grid[j],x,Π[j,:],p)[1],kp_min,kp_max,Brent())
                # Check result
                #converged(min_result) || error("Failed to solve Bellman max for capital =" k_grid[i]" in $(iterations(min_result)) iterations")
                # Record results
                V[i,j] = -min_result.minimum
                G_kp[i,j] = min_result.minimizer
                G_l[i,j] = Obj_fn_condl(k_grid[i],z_grid[j],G_kp[i],Π[j,:],p)[2]
            end
        end
        # Fill in policy for consumption
        G_c[:,j] = z_grid[j].*(collect(k_grid).^α).*(G_l[:,j].^(1-α)) .- G_kp[:,j]
    end
    # Return results
    return V, G_kp, G_c, G_l
end

                ### Solve the problem for θ_k=2,n_z=10,n_k=20 ###
# Get discrete grid for productivity and transition matrix
(log_z,Π) = Rouwenhorst95(10,p)[1:2]
z = exp.(log_z)
# Get solution
@time Mc  = VFI_fixedpoint(T_nested_max,Model(n_k=20,θ_k=2,n_z=10,z_grid=z,Π=Π))
# Interpolate the value function along the z dimension for 3d plot
z_grid = range(z[1],z[end],length=size(z)[1])
z_grid_fine = range(z[1],z[end],length=Mc.n_k_fine)
V_fine_3d = zeros(Mc.n_k_fine,Mc.n_k_fine)
V_fine_33d = [ScaledInterpolations(z_grid,Mc.V_fine[i,:], BSpline(Cubic(Line(OnGrid())))).(collect(z_grid_fine)) for i in 1:Mc.n_k_fine]
for i in 1:Mc.n_k_fine
    V_fine_3d[i,:] = V_fine_33d[i]
end
# Surface and contour plot of the value function
gr()
x= Mc.k_grid_fine; y=z_grid_fine; f=V_fine_3d;
plot(x,y,f, st=:surface,xlabel="Capital",ylabel="Productivity") # Surface plot
title!("Surface 3D plot")
savefig("./Figures/surface_vf")

plot(x,y,f, st=:contour,xlabel="Capital",ylabel="Productivity",nlevels=100, width=2, size=[800,480]) # Contour plot
title!("Level curves")
savefig("./Figures/contour_vf")

max_stars_repo_name: py-r-hans/AdvMacro

max_issues_repo_name: py-r-hans/AdvMacro

max_forks_repo_name: py-r-hans/AdvMacro