text:
filename = @__FILE__()
file_dir = @__DIR__()
file_contents = read(filename, String)

using ArgParse
s = ArgParseSettings()
@add_arg_table s begin
    "--step", "-s"
        default = 10
        arg_type = Int
        help = "big step for problem definition"
    "--obs", "-o"
        help = "observation noise added to distance dependent noise"
        arg_type = Float64
        default = 0.0001
    "-N"
        help = "number of simulations per solver"
        arg_type = Int
        default = 1000
    "--test"
        help = "flag for using only one solver and N=1"
        action = :store_true
    "--solver"
        help = "select one solver to run by name"
        arg_type = String
end
parsed_args = parse_args(ARGS, s)

using DataFrames
using CSV
using Printf
using Random
using Dates

using POMDPs
using ParticleFilters
using MCTS
using POMCPOW
using QMDP
using DiscreteValueIteration
using POMDPSimulators
using POMDPPolicies

using SunbergTypes
using IPFT
using Scenarios

seed = Random.GLOBAL_RNG.seed[1]
file_contents = "rng seed $seed \n" * repr("text/plain", parsed_args) * "\n" * file_contents

@show max_time = 1.0
@show max_depth = 20
N = parsed_args["test"] ? 1 : parsed_args["N"]
@show N

pomdp = SimpleLightDark(big_step=parsed_args["step"], min_obs_noise=parsed_args["obs"])
@show pomdp

solvers = Dict{String, Union{Solver,Policy}}(

    "pomcpow" => begin
        rng = MersenneTwister(seed)
        ro = ValueIterationSolver()
        solver = POMCPOWSolver(tree_queries=10_000_000,
                               criterion=MaxUCB(90.0),
                               max_depth=max_depth,
                               max_time=max_time,
                               enable_action_pw=false,
                               k_observation=5.,
                               alpha_observation=1/15.0,
                               estimate_value=FOValue(ro),
                               check_repeat_obs=false,
                               tree_in_info=false,
                               rng=rng
                              )
        solver
    end,

    "pft" => begin
        rng = MersenneTwister(seed)
        m = 20
        ro = solve(QMDPSolver(), pomdp)
        solver = DPWSolver(n_iterations=typemax(Int),
                           exploration_constant=100.0,
                           depth=max_depth,
                           max_time=max_time,
                           k_state=4.0,
                           alpha_state=1/10,
                           check_repeat_state=false,
                           estimate_value=RolloutEstimator(ro),
                           enable_action_pw=false,
                           tree_in_info=false,
                           rng=rng
                          )
        belief_mdp = MeanRewardBeliefMDP(pomdp,
                                         LowVarianceResampler(m),
                                         0.05
                                        )

        solve(solver, belief_mdp)
    end,

    "ipft-1" => begin
        rng = MersenneTwister(seed)
        m = 20
        ro = solve(QMDPSolver(), pomdp)
        ifm = DiscreteEntropy()
        discount_information_gain = false
        solver = IPFTSolver(n_iterations=typemax(Int),
                        criterion=MultiObjectiveUCB([1.0,50.0],100.0),
                        depth=max_depth,
                        max_time=max_time,
                        check_repeat_action=true,
                        reuse_states=true,
                        k_observation = 5.0,
                        alpha_observation = 1/20,
                        estimate_value=RolloutEstimator(ro),
                        enable_action_pw=false,
                        tree_in_info=true,
                        rng=rng
                        )
        belief_mdp = InformationRewardBeliefMDP(pomdp,
                                         LowVarianceResampler(m),
                                         ifm,
                                         discount_information_gain,
                                         0.05
                                        )
        solve(solver, belief_mdp)
    end,

    "ipft-gamma" => begin
        rng = MersenneTwister(seed)
        m = 20
        ro = solve(QMDPSolver(), pomdp)
        ifm = DiscreteEntropy()
        discount_information_gain = true
        solver = IPFTSolver(n_iterations=typemax(Int),
                        criterion=MultiObjectiveUCB([1.0,50.0],100.0),
                        depth=max_depth,
                        max_time=max_time,
                        check_repeat_action=true,
                        reuse_states=true,
                        k_observation = 5.0,
                        alpha_observation = 1/20,
                        estimate_value=RolloutEstimator(ro),
                        enable_action_pw=false,
                        tree_in_info=true,
                        rng=rng
                        )
        belief_mdp = InformationRewardBeliefMDP(pomdp,
                                         LowVarianceResampler(m),
                                         ifm,
                                         discount_information_gain,
                                         0.05
                                        )
        solve(solver, belief_mdp)
    end,

)

# Consider not all solvers, depending on arguments
if parsed_args["test"]
    solvers = Dict(rand(solvers))
elseif !(parsed_args["solver"] isa Nothing)
    name = parsed_args["solver"]
    solvers = Dict(name=>solvers[name])
end

alldata = DataFrame()
for (k, solver) in solvers
    global alldata
    @show k
    if isa(solver, Solver)
        planner = solve(solver, pomdp)
    else
        planner = solver
    end
    sims = []

    be_updater = planner isa InformationMCTSPlanner

    for i in 1:N
        Random.seed!(planner, seed+i+50_000)

        # up = SIRParticleFilter(pomdp, 10_000, MersenneTwister(seed+i+90_000))
        up = ObsAdaptiveParticleFilter(pomdp,
                                           LowVarianceResampler(10_000),
                                           0.05, MersenneTwister(seed+i+90_000))
        if be_updater
            up = InformationFilter(up, planner.mdp.im)
        end

        sim = Sim(pomdp,
                  planner,
                  up,
                  rng=MersenneTwister(seed+i+70_000),
                  max_steps=100,
                  metadata=Dict(:solver=>k, :i=>i)
                 )

        push!(sims, sim)
    end

    data = run_parallel(sims)
    # data = run(sims)

    rs = data[:reward]
    println(@sprintf("reward: %6.3f ± %6.3f", mean(rs), std(rs)/sqrt(length(rs))))
    if isempty(alldata)
        alldata = data
    else
        alldata = vcat(alldata, data)
    end
end

# Filename without path
filename = split(filename,'/')[end]
# Filename without extension
filename = split(filename,'.')[1]
# Get rid of the _table suffix
filename = join(split(filename,'_')[1:max(1,end-1)], '_')
# Append problem parameters to filename
filename = filename * repr(parsed_args["step"])

result_dir = joinpath(file_dir, "results")
if !isdir(result_dir)
    mkdir(result_dir)
end

datestring = Dates.format(now(), "yyyy-mm-dd_HHMMSS")
copyname = joinpath(result_dir, filename * "_table_$(datestring).jl")
write(copyname, file_contents)
csvname = joinpath(result_dir, filename * "_$(datestring).csv")
println("saving to $csvname...")
CSV.write(csvname, alldata)
println("done.")


text:
"""
HTTP パッケージを使った最小限の http サーバ
"""

using HTTP

println("Serving on http://127.0.0.1:8081")

HTTP.listen("127.0.0.1", UInt16(8081)) do http::HTTP.Stream
    # debug
    @show http.message
    @show HTTP.header(http, "Content-Type")
    while !eof(http)
        println("body data: ", String(readavailable(http)))
    end

    # / => 200 OK: "Hello, world"
    if http.message.target === "/"
        HTTP.setstatus(http, 200)
        HTTP.setheader(http, "Content-Type" => "text/html; charset=utf-8")
        startwrite(http)
        write(http, "Hello, world")
        return
    end

    # 静的ファイル配信
    filename = HTTP.unescapeuri(http.message.target[2:end]) # ファイル名は先頭の'/'を抜いた部分
    if isfile(filename)
        HTTP.setstatus(http, 200)
        HTTP.setheader(http, "Content-Type" => "text/plain")
        startwrite(http)
        write(http, read(filename))
        return
    end
    HTTP.setstatus(http, 404)
end

ext: jl

max_stars_repo_name: amenoyoya/wsgi-api-server

max_issues_repo_name: amenoyoya/wsgi-api-server

max_forks_repo_name: amenoyoya/wsgi-api-server


text:
"""
make a list of blog posts for inclusion on home page
"""
function hfun_blogposts()
    post_section = "
\n"
    list = readdir("pages")
    for post_file in reverse(list)
      info = get_post_info(post_file)
      post_link = "
$(info.title) $(info.date)
"
      post_div = "
\n $post_link\n
\n"
      post_section *= post_div
    end
    post_section *= "
\n"
    return post_section
end

"""
    {{ addcomments }}

Add a comment widget, managed by utterances <https://utteranc.es>.
Set up according to https://franklinjl.org/extras/utterances/
"""
function hfun_addcomments()
    html_str = """

    """
    return html_str
end

function hfun_bar(vname)
  val = Meta.parse(vname[1])
  return round(sqrt(val), digits=2)
end

function hfun_m1fill(vname)
  var = vname[1]
  return pagevar("index", var)
end

function lx_baz(com, _)
  # keep this first line
  brace_content = Franklin.content(com.braces[1]) # input string
  # do whatever you want here
  return uppercase(brace_content)
end

struct PostInfo
  title::String
  pagename::String
  date::String
end

function get_post_info(post_file)
  file_path = joinpath("pages", post_file)
  title = open(file_path) do f
      r = read(f, String)
      m = match(r"@def title = \"(.*?)\"", r)
      return string(first(m.captures))
  end
  #@info " .... processing page $title"
  pagename = first(splitext(post_file))
  postdate = pagename[1:10]
  return PostInfo(title, pagename, postdate)
end

max_stars_repo_name: matthijscox/Blog

max_issues_repo_name: matthijscox/Blog

max_forks_repo_name: matthijscox/Blog


text:
using Optim
using PyPlot
using StatsBase
using ForwardDiff
using celerite
include("regress.jl")

function test_co2_grid()

function log_like_derivative_wrapper(p0,p,x,t,y)
  n = length(t)
  nex = (4(p-p0)+2p0+1)*(n-1)+1
  m1 = 2(p-p0)+p0+2
  if p0 == p
    m1 = p0+ 1
  end
  width = 2m1+1
  indx::Vector{Int64} = collect(1:nex)
#  aex = zeros(Number,width,nex)
#  al_small= zeros(Number,m1,nex)
#  bex = zeros(Number,nex)
  nx = length(x)
  x0 = copy(x)
  xvary=x[1:nx-1]
  function log_like(xvary::Vector)
    x = [xvary;x0[nx]]
    w0 = exp(x[1])
    alpha = exp(x[2:p+1])
    beta_real = exp(x[p+2:2p+1])
    beta_imag = zeros(eltype(x),p)
    if p > p0
      beta_imag[p0+1:p] = exp(x[2p+2:3p-p0+1])
    end
    aex = zeros(eltype(x),width,nex)
    al_small= zeros(eltype(x),m1,nex)
    bex = zeros(eltype(x),nex)
    logdeta = compile_matrix_symm(alpha,beta_real,beta_imag,w0,t,nex,aex,al_small,indx)
    log_like = compute_likelihood(p,p0,y,aex,al_small,indx,logdeta,bex)
#    println("alpha: ",alpha," beta_real: ",beta_real," white: ",w0," log_like: ",log_like)
    return -log_like
  end

## Optimize the likelihood:
  xdiff = copy(xvary)
## Now perturb x:
#  dlogdx = zeros(eltype(x),length(x))
#  log_like0=log_like(x)
#  println(log_like0,x)
#  log_like0=log_like(x)
#  x0 = copy(x)
#  for i=1:length(x)
#    x    = copy(x0)
#    if (x[i] == 0.) then
#      h = 1e-6
#    else
#      h = 1e-6*x[i]
#    end
#    x[i] = x[i]+h
#    println(i," ",x,size(x))
#    log_like_h = log_like(x)
#    dlogdx[i]=(log_like_h-log_like0)/h
#  end
#  result=ForwardDiff.gradient(log_like,xdiff)
#  result=ForwardDiff.gradient(log_like,xdiff)
#  println("Gradient AutoDiff:    ",result)
#  println("Gradient Finite diff: ",dlogdx)

  result = optimize(log_like, xdiff, BFGS(), OptimizationOptions(autodiff = true))
#  result = optimize(log_like, xdiff, BFGS())
#  result = optimize(log_like, x0, BFGS())
#  result = optimize(log_like, xdiff)
  println(result)
  xopt = x0
  xopt[1:nx-1]=exp(Optim.minimizer(result))
  return xopt,Optim.minimum(result)
#  return result
end

# Read in CO2 data:

data = readdlm("data/CO2_data.csv.txt",',')
co2 = vec(data[:,2])
time = vec(data[:,1] - data[1,1])
# Carry out cubic regression:
nt = length(time)
fn = zeros(4,nt)
tnorm = (time-time[1])/(time[nt]-time[1])
fn[1,:]=1.0
fn[2,:]= tnorm
fn[3,:]= tnorm.*tnorm
fn[4,:]= tnorm.*tnorm.*tnorm
sig = ones(nt)
coeff,cov = regress(fn,co2,sig)
println("Coeff: ",coeff)
co2_trend = zeros(nt)
co2_sub = copy(co2)
for i=1:nt
  for j=1:4
    co2_trend[i] += coeff[j]*fn[j,i]
  end
  co2_sub[i] -= co2_trend[i]
end

plot(time,co2)
plot(time,co2_trend)
plot(time,co2_sub)
#read(STDIN,Char)

# Compute ACF:
nt = length(time)
dt = median(time[2:nt]-time[1:nt-1])
lags = collect(0:1:200)
println(typeof(lags))
acf = autocov(co2_sub,lags)
lags *=dt
clf()
# Now, model covariance:
acf_model = 3.5.*exp(-lags./50.).*cos(2.*pi.*lags)+0.5*cos(2.*pi*lags/0.5)
clf()
plot(lags,acf-acf_model)
#read(STDIN,Char)

# Initialize model with guessed parameters:

nkermax = 5
nperiod = 200
alpha   = [1.0]
beta_real = [0.01]
beta_imag = Float64[]
period = logspace(-1,1,nperiod)
log_like_period = zeros(nperiod)
xbest = Float64[]
ibest = -1
for nkernel=2:nkermax
  alpha = [alpha;1.0]
  beta_real = [beta_real;0.01]
  log_like_best = Inf
  for iperiod=1:nperiod
# Now, insert into the parameter vector:
    x0 = [log(1.0);log(alpha);log(beta_real);log(beta_imag);log(2pi/period[iperiod])]
    nx = length(x0)
    p = nkernel-1
    p0 = 1
# Vary all parameters except the last one (frequency):
# x,best_log_like = log_like_derivative_wrapper(p0,p,x0,time,co2_sub)
    xopt,log_like_opt = log_like_derivative_wrapper(p0,p,x0,time,co2_sub)
    log_like_period[iperiod]=log_like_opt
    if log_like_opt < log_like_best
       log_like_best = copy(log_like_opt)
       xbest = copy(xopt)
       ibest = iperiod
    end
  end
  clf()
  semilogx(period,log_like_period)
  print(ibest," ",xbest," ",log_like_best)
  beta_imag = [beta_imag;2pi/period[ibest]]
  read(STDIN,Char)
end
#println(x0,x,best_log_like)
println(xbest)
# Plot the optimized ACF:
acf_model = xbest[2].*exp(-lags.*xbest[nkernel+2])
for j=2:nkernel
  acf_model += xbest[j+1]*exp(-lags.*xbest[nkernel+j+1]).*cos(lags.*xbest[2*nkernel+j])
end
#+xbest[4]*exp(-lags.*xbest[8]).*cos(lags.*xbest[11])+xbest[5]*exp(-lags.*xbest[9]).*cos(lags.*xbest[12])
acf_model[1] += xbest[1]
clf()
plot(lags,acf_model)
plot(lags,acf)
plot(lags,acf-acf_model)
read(STDIN,Char)
# Plot the power spectrum:
nfreq = 1000
omega = collect(logspace(log10(0.001),log10(1000.),nfreq))
psd = zeros(nfreq)
for i=1:nfreq
  psd[i] = 2.0*xbest[2]*xbest[6]/(omega[i]^2+xbest[6]^2)
  for j=2:nkernel
    psd[i] += 2.0*xbest[j+1]*xbest[nkernel+j+1]/((omega[i]-xbest[2*nkernel+j])^2+xbest[nkernel+j+1]^2)
  end
end
clf()
loglog(omega,psd)
read(STDIN,Char)

return
end


text:
# SPDX-License-Identifier: MIT

# NOTE: depending on the font the variable 'ν' (Greek letter nu) may appear
# identical to 'v' (the latin letter v).

struct Wishart{T1<:AbstractMatrix, T2<:Real} <: ExponentialFamilyDistribution
    W::T1
	ν::T2
end

function η(p::Wishart)
	D = size(p.W, 1)
	W⁻¹ = inv(Symmetric(p.W))
	vcat(-(1/2)*vec(W⁻¹), (p.ν - D - 1)/2)
end

function ξ(p::Wishart, η)
	D = size(p.W, 1)
	W = inv(-2*reshape(η[1:end-1], D, D))
	ν = 2*η[end] + D + 1
	vcat(vec(W), ν)
end

function unpack(p::Wishart, μ)
	D = size(p.W, 1)
	X = reshape(μ[1:end-1], D, D)
	logdetX = μ[end]
	(X=X, logdetX=logdetX)
end

function A(p::Wishart, η)
	D = size(p.W, 1)

	# We don't use `p.W` or `p.ν` to be able to calculate the gradient w.r.t. η.
	W⁻¹ = -2*reshape(η[1:end-1], D, D)
	ν = 2*η[end] + D + 1

	(ν/2)*(-logdet(W⁻¹) + D*log(2)) + sum(loggamma.((ν + 1 .- (1:D))/2))
end

sample(p::Wishart) = rand(Distributions.Wishart(ν, W))

max_stars_repo_name: BUTSpeechFIT/PPCA

max_issues_repo_name: BUTSpeechFIT/PPCA

max_forks_repo_name: BUTSpeechFIT/PPCA


text:
for i in 1:5
    for j in 1:i
        print("*")
    end
    println()
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
function cox_f(S, fs, ls, ξ , X, β, λ, alive, afterΘ) # preprocessed already
    Xβ = X*β
    Θ = exp.(Xβ)
    after!(afterΘ, Θ)

    #compute loglikelihood
    y = 0.
    for i in 1:length(fs)
        for j in (fs[i]):(ls[i])
            ρ = (alive[j]-alive[fs[i]])/(alive[fs[i]]-alive[ls[i]+1])
            ϕ = afterΘ[fs[i]]-ρ*(afterΘ[fs[i]]-afterΘ[ls[i]+1])
            y -= Xβ[j] -log(ϕ)
        end
    end
    y += λ*(β' * β)
    return y
end

# preprocessed already:
# fs = index first deaths, ls = index last deaths,
# X is covariates, ξ is covariate covariate transpose
function cox_h!(grad,hes, S, fs, ls, ξ , X, β, λ, alive, afterΘ, afterXΘ, afterξΘ)
    #compute relevant quantities for loglikelihood, score, fischer_info

    Xβ = X*β
    Θ = exp.(Xβ)
    after!(afterΘ,Θ)
    after!(afterXΘ, X.*Θ)
    after!(afterξΘ,ξ.*Θ)

    #compute loglikelihood, score, fischer_info
    #From v0.6 remember to use . notation and do it inplace, possibly using views!
    y = 0.
    grad[:] = 0.
    hes[:] = 0.

    # preallocate
    Z = zeros(size(X,2))
    Ξ = zeros(size(X,2),size(X,2))

    for i in 1:length(fs)
        for j in (fs[i]):(ls[i])
            ρ = (alive[j]-alive[fs[i]])/(alive[fs[i]]-alive[ls[i]+1])
            ϕ = afterΘ[fs[i]]-ρ*(afterΘ[fs[i]]-afterΘ[ls[i]+1])
            for k in eachindex(Z)
                @inbounds Z[k] = afterXΘ[fs[i],k]-ρ*(afterXΘ[fs[i],k]-afterXΘ[ls[i]+1,k])
            end
            for k2 in 1:size(Ξ,2)
                for k1 in 1:size(Ξ,1)
                    @inbounds Ξ[k1,k2] = afterξΘ[fs[i],k1,k2]-ρ*(afterξΘ[fs[i],k1,k2]-afterξΘ[ls[i]+1,k1,k2])
                end
            end
            y -= Xβ[j] -log(ϕ)
            for k2 in 1:size(Ξ,2)
                @inbounds grad[k2] -= X[j,k2]
                @inbounds grad[k2] += Z[k2]/ϕ
                for k1 in 1:size(Ξ,1)
                    @inbounds hes[k1, k2] += Ξ[k1,k2]/ϕ - Z[k1]*Z[k2]/ϕ^2
                end
            end
        end
    end
    y += λ*(β' * β)

    for k1 in 1:size(Ξ,1)
        grad[k1] +=  2*λ*β[k1]
        hes[k1,k1] +=  2*λ
    end
    return y
end

function coxph(S::AbstractVector,X::AbstractArray; l2_cost = 0., kwargs...)
    ξ = zeros(size(X,1),size(X,2),size(X,2))
    for k2 in 1:size(ξ,3)
        for k1 in 1:size(ξ,2)
            for i in 1:size(ξ,1)
                @inbounds ξ[i,k1,k2] = X[i,k1]*X[i,k2]
            end
        end
    end

    # compute first and last!

    fs = find(firsts(S))
    ls = find(lasts(S))

    # do optimization

    alive = after(ones(Int64, length(S)))
    afterΘ = init_after(zeros(size(X,1)))
    afterXΘ = init_after(X.*zeros(size(X,1)))
    afterξΘ = init_after(ξ.*zeros(size(X,1)))

    f1 = (β) -> cox_f(S, fs, ls, ξ , X, β, l2_cost, alive, afterΘ)
    h1! = (β,grad,hes) -> cox_h!(grad,hes, S, fs, ls, ξ , X, β, l2_cost,alive, afterΘ, afterXΘ, afterξΘ)
    return newton_raphson(f1,h1!, zeros(size(X,2)); kwargs...)
end

function coxph(formula::Formula, data::DataFrame; l2_cost = 0., kwargs...)
    sorted_data = sort(data, cols = formula.lhs)
    M = DataFrames.ModelFrame(formula,sorted_data)
    S = convert(Array, M.df[:,1])
    model_matrix = DataFrames.ModelMatrix(M)
    X = model_matrix.m[:,2:size(model_matrix.m,2)]
    β, neg_ll,grad, hes =  coxph(S, X; l2_cost = l2_cost, kwargs...)
    rownms = coefnames(M)[2:end]
    se = sqrt.(diag(pinv(hes)))
    z_score = β./se
    pvalues = 2*cdf(Normal(),-abs.(z_score))
    coefmat = CoefTable(hcat([β, se, z_score, pvalues]...),
    ["Estimate", "Std.Error", "z value", "Pr(>|z|)"], rownms, 4)
    CoxModel("Cox; ", formula, coefmat, M, -neg_ll, -grad, hes)
end


text:
module ShellChron

# Export names
# export run_model

# Load requirements
begin
    using CSV
    using DataFrames
end

# Define types and structs
struct Temperature
    av::Number
    amp::Number
    pha::Number
    per::Number
end

struct Growth
    av::Number
    amp::Number
    pha::Number
    per::Number
    skw::Number
end

struct Model
    t::Temperature
    g::Growth
end

# Include functions in module
include("gui.jl")
include("data_import.jl")
include("age_corr.jl")
include("sinreg.jl")

# Experiment with functionality
T = Temperature(4.2, 5, 2, 3)
G = Growth(4.2, 5, 2, 3, 1)
mod = Model(T, G)

data_import("data-raw/test.csv")

end


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    acceleration_functions = gather_accelerations_for_potentials(simulation)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


text:
using PyPlot

function McCall_ddp(beta, b, w, rho)

    N = length(w)
    REJECT = 1
    ACCEPT = 2

    #  Construct reward function f(s,a)
    f = zeros(N, 2)
    f[:, REJECT] = b
    f[:, ACCEPT] = w

    #  Construct state transition probability matrix P(s'|s,x)
    P = zeros(2, N, N)
    P[REJECT, :, :] = (rho-(1-rho)/(N-1))*eye(N,N) + (1-rho)/(N-1)*ones(N,N)
    P[ACCEPT, :, :] = eye(N)

    #  Pack model structure
    model = Dict()
    model["reward"] = f
    model["transprob"] = P
    model["discount"] = beta

    #  Solve model
    v, a = ddpsolve(model)

    return v
end

N =  100
beta = .9
w_min = 0
w_max = 200
b =  [0 20 50]
rho=  [1/N .2 .5]
w = linspace(w_min, w_max, N)

Vb = zeros(N, 3)
Vrho = zeros(N, 3)
for i = 1:3
    Vb[:, i] = McCall_ddp(beta, b[i], w, 1/N)
    Vrho[:, i] = McCall_ddp(beta, 0, w, rho[i])
end

fig, (ax1, ax2) = subplots(2, 1)
ax1[:plot](w, Vb[:, 1], "k-", w, Vb[:, 2], "r--", w, Vb[:, 3], "b-.")
ax1[:set_ylabel]("V(wage)")
ax1[:set_xlabel]("wage")
ax1[:set_title](L"Varying $b$ fixing $\rho=1/N$")
ax1[:legend](["b=$i" for i in b])

ax2[:plot](w, Vrho[:, 1], "k-", w, Vrho[:, 2], "r--", w, Vrho[:, 3], "b-.")
ax2[:set_ylabel]("V(wage)")
ax2[:set_xlabel]("wage")
ax2[:set_title](L"Varying $\rho$ fixing $b=0$")
ax2[:legend]([L"$\rho$ = "* "$i" for i in rho])


text:
#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 1
#:::::::::::::::::::::::::::::::::::::::::::::::::::
using DataFrames
using CSV
using HTTP
using Optim
using LinearAlgebra # For diag function

url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS3-gev/nlsw88w.csv"

df = CSV.read(HTTP.get(url).body)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
         df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occupation

function mlogit_with_Z(theta, X, Z, y)

    # Define alpha as all elements of theta except the coeficient of Z
    alpha = theta[1:end-1]
    # Define gamma as the coefficient on Z
    gamma = theta[end]
    # Define K as the second dimension of X
    K = size(X, 2)
    # Define J as the number of alternatives to choose from
    J = length(unique(y))
    # Define N as the number of observations in the dependent variable
    N = length(y)
    # Define bigY as an array of Zeros of size NxJ
    bigY = zeros(N, J)

    # Fill each column of bigY with indicator = 1 if col_i = choice_i
    for j=1:J
        bigY[:,j] = y.==j
    end
    #=
    Create Kx(J-1) array with an additional column of zeros (i.e. a KxJ array).
    Note this takes alpha (a vector) and transforms it into a KxJ array.
    =#
    bigAlpha = [reshape(alpha, K, J-1) zeros(K)]

    # Store the promoted type of X and alpha in variable T
    T = promote_type(eltype(X), eltype(theta))
    # Initialize numerator (array of zeros of type T and size NxJ)
    num = zeros(T, N, J)
    # Initialize denominator
    dem = zeros(T,N)
    #Create loop
    for j=1:J
        #= Each column i of num refers to the numerator of the choice probability for choice i
        =#
        num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)
        # The denominator is the sum of each of the choice probability numerators
        dem .+= num[:,j]
    end

    #= Create matrix in which every column i is a choice probability for alternative i. Need to divide each column of num by dem which is a vector, so repeat dem J times to allow element-wise division.
    =#
    P = num./repeat(dem,1,J)

    loglike = -sum( bigY.*log.(P) )

    return loglike
end

# Create a 22 element array of starting values
startvals = [2*rand(7*size(X, 2)).-1; .1]
# Declare function is twice differentiable
td = TwiceDifferentiable(theta -> mlogit_with_Z(theta, X, Z, y), startvals; autodiff = :forward)

# Run the optimizer (I'll use Newton's method instead of LBFGS)
theta_hat_optim_ad = optimize(td, startvals, Newton(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace = true, show_every=50))
theta_hat_mle_ad = theta_hat_optim_ad.minimizer

# Evaluate the Hessian at the estimates
H = Optim.hessian!(td, theta_hat_mle_ad)
theta_hat_mle_ad_se = sqrt.(diag(inv(H)))
println([theta_hat_mle_ad theta_hat_mle_ad_se]) # these standard errors match Stata

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 2
#:::::::::::::::::::::::::::::::::::::::::::::::::::

#= The coefficient gamma represents the change in utility with a 1-unit change in log wages. More properly, gamma/100 is the change in utility with a 1% increase in expected wage
=#

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 3
#:::::::::::::::::::::::::::::::::::::::::::::::::::

function nested_logit_with_Z(theta, X, Z, y, nesting_stucture)

    #= Unwrap parameter vector. Alpha contains the coefficients on the x's, lambda contains the measures of the degree of independence in unobserved utility among the alternatives within each nest, and gamma is the coefficient on the difference in Z's
    =#
    alpha = theta[1:end-3]
    lambda = theta[end-2:end-1]
    gamma = theta[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j = 1:J
        bigY[:,j] = y.==j
    end


text:
grid_size = 25
α_vals = linspace(0.05, 0.5, grid_size)
w_bar_vals = similar(α_vals)

mcm = McCallModel()

for (i, α) in enumerate(α_vals)
    mcm.α = α
    w_bar = compute_reservation_wage(mcm)
    w_bar_vals[i] = w_bar
end

plot(α_vals,
     w_bar_vals,
     lw=2,
     α=0.7,
     xlabel="job separation rate",
     ylabel="reservation wage",
     label=L"$\bar w$ as a function of $\α$")

max_stars_repo_name: yywingliang/lecture-source-jl

max_issues_repo_name: yywingliang/lecture-source-jl

max_forks_repo_name: yywingliang/lecture-source-jl


text:
# Implementation of the amb operator using shift/reset
# http://community.schemewiki.org/?amb

include("continuations.jl");

struct Backtrack end

function require(x)
  x || throw(Backtrack())
  return
end

unwrap(e) = e
unwrap(e::CapturedException) = e.ex

function amb(iter)
  shift() do k
    for x in iter
      try
        return k(x)
      catch e
        unwrap(e) isa Backtrack || rethrow()
      end
    end
    throw(Backtrack())
  end
end

function ambrun(f)
  try
    @reset f()
  catch e
    e isa Backtrack || rethrow()
    error("No possible combination found.")
  end
end

ambrun() do
  x = amb([1, 2, 3])
  y = amb([1, 2, 3])
  require(x^2 + y == 7)
  (x, y)
end

ambrun() do
  N = 20
  i = amb(1:N)
  j = amb(i:N)
  k = amb(j:N)
  require(i*i + j*j == k*k)
  (i, j, k)
end


text:
@testcase "active patterns" begin
    @testset "regular case" begin
        @lift @active LessThan0(x) begin
            if x > 0
                nothing
            else
                Some(x)
            end
        end

        b = (@match 15 begin
            LessThan0(_) => :a
            _ => :b
        end)

        @test b === :b

        @test (@match -15 begin
            LessThan0(a) => a
            _ => 0
        end) == -15
    end

    @testset "parametric case" begin

        @lift @active Re{r :: Regex}(x) begin
            ret = match(r, x)
            ret === nothing || return Some(ret)
        end

        @test (@match "123" begin
            Re{r"\d+"}(x) => x.match
            _ => @error ""
        end) == "123"

        @test_skip @match "abc" begin
            Re{r"\d+"}(x) => x
            _ => @error ""
        end
    end

    @testset "custom pattern for given structs" begin
        @eval struct Interval end

        @lift @active internal Interval{a, b}(arg) begin
            a <= arg <= b
        end

        @use Enum

        @lift @active visible in (@__MODULE__) IsEven(a) begin
            a % 2 === 0
        end

        @lift MLStyle.is_enum(::Type{IsEven}) = true
        function parity(x)
            @match x begin
                IsEven => :even
                _ => :odd
            end
        end
        @test :even === parity(4)
        @test :odd === parity(3)

        @test 2 == @match 3 begin
            Interval{1, 2} => 1
            Interval{3, 4} => 2
            Interval{5, 6} => 3
            Interval{7, 8} => 4
        end

    end

    @testset "drop Union{T, Nothing}" begin
        @lift @active F(x) begin
            x
        end
        @test_throws Any @match 1 begin
            F(1) => 1
        end
    end

    @lift using MLStyle.Sugars
    @testset "eye-candy 'and' patterns" begin
        @test @match (1, 2) begin
            And[
                (1, _),
                (_, 2)
            ] => true
            _ => false
        end
    end
end


text:
function print_logo(io = stdout)
    c = string(Crayon(foreground = :light_red))
    r = string(Crayon(foreground = :light_green))
    a = string(Crayon(foreground = :light_yellow))
    y = string(Crayon(foreground = :light_blue))
    o = string(Crayon(foreground = :light_magenta))
    n = string(Crayon(foreground = :light_cyan))
    s = string(Crayon(foreground = :dark_gray))
    res = string(Crayon(reset = true))

    str = """
       $(c)██████╗$(r)██████╗  $(a)█████╗ $(y)██╗   ██╗ $(o)██████╗ $(n)███╗   ██╗$(s)███████╗
      $(c)██╔════╝$(r)██╔══██╗$(a)██╔══██╗$(y)╚██╗ ██╔╝$(o)██╔═══██╗$(n)████╗  ██║$(s)██╔════╝
      $(c)██║     $(r)██████╔╝$(a)███████║ $(y)╚████╔╝ $(o)██║   ██║$(n)██╔██╗ ██║$(s)███████╗
      $(c)██║     $(r)██╔══██╗$(a)██╔══██║  $(y)╚██╔╝  $(o)██║   ██║$(n)██║╚██╗██║$(s)╚════██║
      $(c)╚██████╗$(r)██║  ██║$(a)██║  ██║   $(y)██║   $(o)╚██████╔╝$(n)██║ ╚████║$(s)███████║
       $(c)╚═════╝$(r)╚═╝  ╚═╝$(a)╚═╝  ╚═╝   $(y)╚═╝    $(o)╚═════╝ $(n)╚═╝  ╚═══╝$(s)╚══════╝$(res)
    """

    print(io, "\n\n", str, "\n")
end


text:
using VarianceComponentsHDFE
using Documenter

makedocs(;
    modules=[VarianceComponentsHDFE],
    authors="Various Collaborators",
    repo="https://github.com/HighDimensionalEconLab/VarianceComponentsHDFE.jl/blob/{commit}{path}#L{line}",
    sitename="VarianceComponentsHDFE.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://HighDimensionalEconLab.github.io/VarianceComponentsHDFE.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/HighDimensionalEconLab/VarianceComponentsHDFE.jl",
)


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

D-optimal experiment design maximizes the determinant of the information matrix
adapted from Boyd and Vandenberghe, "Convex Optimization", section 7.5
maximize    F(V × Diagonal(x) × V')
subject to  sum(np) == n
            0 .<= np .<= n_max
where np is a vector of variables representing the number of experiment to run,
and the columns of V are the vectors representing each experiment

if logdet_obj or rootdet_obj is true, F is the logdet or rootdet function
if geomean_obj is true, we use a formulation from
https://picos-api.gitlab.io/picos/optdes.html that finds an equivalent minimizer
=#

struct DOptimalDesignJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    q::Int
    p::Int
    n::Int
    n_max::Int
    logdet_obj::Bool # use formulation with logdet objective
    rootdet_obj::Bool # use formulation with rootdet objective
    geomean_obj::Bool # use formulation with geomean objective
end

function build(inst::DOptimalDesignJuMP{T}) where {T <: Float64}
    (q, p, n, n_max) = (inst.q, inst.p, inst.n, inst.n_max)
    @assert (p > q) && (n > q) && (n_max <= n)
    @assert +(inst.logdet_obj, inst.geomean_obj, inst.rootdet_obj) == 1
    V = randn(q, p)

    model = JuMP.Model()
    JuMP.@variable(model, np[1:p])
    JuMP.@constraint(model, sum(np) == n)
    mid = n_max / 2
    JuMP.@constraint(model, vcat(mid, np .- mid) in MOI.NormInfinityCone(p + 1))
    JuMP.@variable(model, hypo)
    JuMP.@objective(model, Max, hypo)

    if inst.logdet_obj || inst.rootdet_obj
        # information matrix lower triangle
        v1 = [JuMP.@expression(model,
            sum(V[i, k] * np[k] * V[j, k] for k in 1:p)) for i in 1:q for j in 1:i]
        if inst.logdet_obj
            JuMP.@constraint(model, vcat(hypo, 1, v1) in MOI.LogDetConeTriangle(q))
        else
            JuMP.@constraint(model, vcat(hypo, v1) in MOI.RootDetConeTriangle(q))
        end
    else
        # hypogeomean + epinormeucl formulation
        JuMP.@variable(model, L[i in 1:q, j in 1:i])
        JuMP.@variable(model, W[1:p, 1:q])
        VW = V * W
        JuMP.@constraints(model, begin
            [i in 1:q, j in 1:i], VW[i, j] == L[i, j]
            [i in 1:q, j in (i + 1):q], VW[i, j] == 0
            vcat(hypo, [L[i, i] for i in 1:q]) in MOI.GeometricMeanCone(1 + q)
            [i in 1:p], vcat(sqrt(q) * np[i], W[i, :]) in JuMP.SecondOrderCone()
        end)
    end

    return model
end


text:
#!/usr/bin/env julia

function parse_num(v::Vector{Bool})
  lb = 0
  ub = 2^length(v) - 1

  for b in v
    mid = (lb + ub) ÷ 2
    if b
      lb = mid + 1
    else
      ub = mid
    end
  end

  return lb
end

function str_to_bool_vec(str::String, falsy::Char, truthy::Char)
  map(collect(str)) do c
    if c == falsy
      return false
    end
    if c == truthy
      return true
    end

    throw("Unknown type of char")
  end
end

function parse_seat(str::String)
  row = str_to_bool_vec(str[1:7], 'F', 'B')
  col = str_to_bool_vec(str[8:10], 'L', 'R')

  parse_num(row), parse_num(col)
end

function find_max_id()
  max_idx = -1
  for line in eachline("./input")
    row, col = parse_seat(line)

    max_idx = max(max_idx, row * 8 + col)
  end
  println(max_idx)
end

function find_missing_id()
  row_col_seats = [parse_seat(l) for l in eachline("./input")]
  seats = Set([x[1] * 8 + x[2] for x in row_col_seats])
  for res in minimum(seats):maximum(seats)
    if res ∉ seats
      println(res)
    end
  end
end

find_missing_id()


text:
export iauPfw06
"""
Precession angles, IAU 2006 (Fukushima-Williams 4-angle formulation).

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  canonical model.

Given:
	date1,date2  double   TT as a 2-part Julian Date (Note 1)

Returned:
	gamb         double   F-W angle gamma_bar (radians)
	phib         double   F-W angle phi_bar (radians)
	psib         double   F-W angle psi_bar (radians)
	epsa         double   F-W angle epsilon_A (radians)

Notes:

	1. The TT date date1+date2 is a Julian Date, apportioned in any
	convenient way between the two arguments.  For example,
	JD(TT)=2450123.7 could be expressed in any of these ways,
	among others:

		date1          date2

		2450123.7           0.0       (JD method)
		2451545.0       -1421.3       (J2000 method)
		2400000.5       50123.2       (MJD method)
		2450123.5           0.2       (date & time method)

	The JD method is the most natural and convenient to use in
	cases where the loss of several decimal digits of resolution
	is acceptable.  The J2000 method is best matched to the way
	the argument is handled internally and will deliver the
	optimum resolution.  The MJD method and the date & time methods
	are both good compromises between resolution and convenience.

	2. Naming the following points:

		e = J2000.0 ecliptic pole,
		p = GCRS pole,
		E = mean ecliptic pole of date,
	and   P = mean pole of date,

	the four Fukushima-Williams angles are as follows:

	gamb = gamma_bar = epE
	phib = phi_bar = pE
	psib = psi_bar = pEP
	epsa = epsilon_A = EP

	3. The matrix representing the combined effects of frame bias and
	precession is:

	PxB = R_1(-epsa).R_3(-psib).R_1(phib).R_3(gamb)

	4. The matrix representing the combined effects of frame bias,
	precession and nutation is simply:

	NxPxB = R_1(-epsa-dE).R_3(-psib-dP).R_1(phib).R_3(gamb)

	where dP and dE are the nutation components with respect to the
	ecliptic of date.

Reference:

	Hilton, J. et al., 2006, Celest.Mech.Dyn.Astron. 94, 351

Called:
	iauObl06     mean obliquity, IAU 2006

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauPfw06(double date1, double date2,
#               double *gamb, double *phib, double *psib, double *epsa)
function iauPfw06(date1::Real, date2::Real)
	ref_gamb = Ref{Float64}(0.0)
	ref_phib = Ref{Float64}(0.0)
	ref_psib = Ref{Float64}(0.0)
	ref_epsa = Ref{Float64}(0.0)

	ccall((:iauPfw06, libsofa_c), Cvoid,
		(Cdouble, Cdouble,
		Ref{Cdouble}, Ref{Cdouble},
		Ref{Cdouble}, Ref{Cdouble}),
		convert(Float64, date1),
		convert(Float64, date2),
		ref_gamb, ref_phib, ref_psib, ref_epsa)

	return ref_gamb[], ref_phib[], ref_psib[], ref_epsa[]
 end


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

using Latte
using FactCheck

function interpolation_forward(input, resize_factor, kernel, pad, output)
    width, height, channels, num = size(input)
    interpolated_height, interpolated_width, _, _ = size(output)

    for n = 1:num
        for c = 1:channels
            for ih = 1:interpolated_height
                r_f = ih * resize_factor
                delta_r = r_f - floor(r_f)
                for iw = 1:interpolated_width
                    c_f = iw * resize_factor
                    delta_c = c_f - floor(c_f)

                    hstart = floor(Int, (ih-1) / resize_factor) - pad + 1
                    wstart = floor(Int, (iw-1) / resize_factor) - pad + 1
                    hend = min(hstart + kernel - 1, height)
                    wend = min(wstart + kernel - 1, width)
                    hstart = max(1, hstart)
                    wstart = max(1, wstart)

                    region =  sub(input, wstart:wend, hstart:hend, c, n)

                    r = 1
                    c = 1
                    r_plus_1 = size(region, 1) > 1 ? 2 : 1
                    c_plus_1 = size(region, 2) > 1 ? 2 : 1
                    output[iw, ih, c, n] = region[r, c] * (1-delta_r) * (1-delta_c)
                                        +  region[r_plus_1, c] * delta_r * (1-delta_c)
                                        +  region[r, c_plus_1]  * (1-delta_r) * delta_c
                                        +  region[r_plus_1, c_plus_1] * delta_r * delta_c
                end
            end
        end
    end
end

function interpolation_backward(input, resize_factor, kernel, pad, output)
    width, height, channels, num = size(input)
    interpolated_height, interpolated_width, _, _ = size(output)

    for n = 1:num
        for c = 1:channels
            for ih = 1:interpolated_height
                r_f = ih * resize_factor
                delta_r = r_f - floor(r_f)
                for iw = 1:interpolated_width
                    c_f = iw * resize_factor
                    delta_c = c_f - floor(c_f)

                    hstart = floor(Int, (ih-1) / resize_factor) - pad + 1
                    wstart = floor(Int, (iw-1) / resize_factor) - pad + 1
                    hend = min(hstart + kernel - 1, height)
                    wend = min(wstart + kernel - 1, width)
                    hstart = max(1, hstart)
                    wstart = max(1, wstart)

                    region =  sub(input, wstart:wend, hstart:hend, c, n)
                    r = 1
                    c = 1
                    r_plus_1 = size(region, 1) > 1 ? 2 : 1
                    c_plus_1 = size(region, 2) > 1 ? 2 : 1

                    region[r,c] += (1-delta_r) * (1-delta_c) * output[iw, ih, c, n]
                    region[r_plus_1, c] += delta_r * (1-delta_c) * output[iw, ih, c, n]
                    region[r, c_plus_1] += (1-delta_r) * delta_c * output[iw, ih, c, n]
                    region[r_plus_1, c_plus_1] += delta_r * delta_c * output[iw, ih, c, n]
                end
            end
        end
    end
end

facts("Testing Bilinear Interpolation Layer (Enlarge)") do

    net = Net(1)
    resize_factor = Float32(8.0)
    data, data_value = MemoryDataLayer(net, :data, (64, 64, 1))
    data_value[:] = rand(Float32, size(data_value)...)
    interp1 = InterpolationLayer(:interp1, net, data, 2, 0, resize_factor)

    init(net)

    params = SolverParameters(
        lr_policy    = LRPolicy.Decay(.01f0, 5.0f-7),
        mom_policy   = MomPolicy.Fixed(0.9),
        max_epoch    = 300,
        regu_coef    = .0005)
    sgd = SGD(params)

    context("Forward") do
        forward(net; solver=sgd)
        input = get_buffer(net, :datavalue)
        expected = zeros(get_buffer(net, :interp1value))
        interpolation_forward(input, resize_factor, 2, 0, expected)
        @fact expected --> roughly(get_buffer(net, :interp1value))
    end

    context("Backward") do
        top_diff = get_buffer(net, :interp1∇)
            ∇input = get_buffer(net, :data∇)
            ∇input_expected = deepcopy(∇input)

        backward(net; solver=sgd)

        interpolation_backward(∇input_expected, resize_factor, 2, 0, top_diff)
        @fact ∇input --> roughly(∇input_expected)
    end

end

facts("Testing Bilinear Interpolation Layer (Shrink)") do

    net = Net(1)
    resize_factor = Float32(1/8.0)
    data, data_value = MemoryDataLayer(net, :data, (64, 64, 1))
    data_value[:] = rand(Float32, size(data_value)...)
    interp1 = InterpolationLayer(:interp1, net, data, 2, 0, resize_factor)

    init(net)

    params = SolverParameters(
        lr_policy    = LRPolicy.Decay(.01f0, 5.0f-7),
        mom_policy   = MomPolicy.Fixed(0.9),
        max_epoch    = 300,
        regu_coef    = .0005)
    sgd = SGD(params)

    context("Forward") do
        forward(net; solver=sgd)
        input = get_buffer(net, :datavalue)
        expected = zeros(get_buffer(net, :interp1value))
        interpolation_forward(input, resize_factor, 2, 0, expected)
        @fact expected --> roughly(get_buffer(net, :interp1value))
    end

    context("Backward") do
        top_diff = get_buffer(net, :interp1∇)
            ∇input = get_buffer(net, :data∇)
            ∇input_expected = deepcopy(∇input)

        backward(net; solver=sgd)

        interpolation_backward(∇input_expected, resize_factor, 2, 0, top_diff)
        @fact ∇input --> roughly(∇input_expected)
    end

end

FactCheck.exitstatus()


text:
using Plots
using LaTeXStrings
using Printf

function rotate_text(TextHeight::Float64, Vec::Vector{T}) where {T<:Real}
    modVec = hypot(Vec[1], Vec[2])
    Vec[1] >= 0 && (angVec = asind(Vec[2]/modVec))
    Vec[1] < 0 && (angVec = 180 - asind(Vec[2]/modVec))
    rotM = [cosd(angVec) sind(angVec)
           -sind(angVec) cosd(angVec)]
    TextPos = [modVec/2 TextHeight]
    return TextPos*rotM
end

default(grid = true, markersize = 8, framestyle = :none, linewidth = 3, legend = :false,
    aspect_ratio = :equal, gridalpha = 1)

#definindo número de vetores a serem somados
nvec = 8

#criando nvec vetores aleatórios
vecs = rand(-5:1:5, (nvec,2))

#resultante dos nvec vetores
res = [sum(vecs[:,1]), sum(vecs[:,2])]

#calculando ângulo que a resultante faz com a direção x
ang_res = atand(res[2]/res[1])

#string em latex com os valores de cada vetor.
vecstr = String("")
for i = 1:nvec
    global vecstr = string(vecstr, string("\\vec{a}_{", i,"} = (", vecs[i,1], ",", vecs[i,2], ")\\, \\newline"))
end

latexstring(vecstr)

#calculando limites na direção x e na direção y para a criação da imagem
xmax = max(vecs[1,1],0)
ymax = max(vecs[1,2],0)

xmin = min(vecs[1,1],0)
ymin = min(vecs[1,2],0)

#gerando vetor com os pontos de origem de cada um dos vetores que serão somados
origins = [0 0]

#calculando limites e origem dos vetores
for i = 2: nvec
    global origins = vcat(origins, [(origins[i-1,1] + vecs[i-1,1]) (origins[i-1,2] + vecs[i-1,2])])
    global xmax = max(xmax, origins[i,1], vecs[i,1])
    global ymax = max(ymax, origins[i,2], vecs[i,2])
    global xmin = min(xmin, origins[i,1], vecs[i,1])
    global ymin = min(ymin, origins[i,2], vecs[i,2])
end

xmax = max(xmax, res[1])
xmin = min(xmin, res[1])

ymax = max(ymax, res[2])
ymin = min(ymin, res[2])

xlimits = (xmin - 1, xmax + 1)
ylimits = (ymin - 1, ymax + 1)

#gerando texto que será inserido na imagem
res_pos = rotate_text(0.35, res)

resstr = latexstring("{\\rm Resultante:\\;} \\vec{S}  = (", res[1],",", res[2],")")

ann_res = [(res_pos[1,1], res_pos[1,2], text(resstr, halign = :center, valign = :center, pointsize = 24, rotation = ang_res))]

emailstr = "ghralima at cefetmg.br"
ann_autor = [((xmax +0.5),(ymax+ymin)/2, text(emailstr, halign = :center, valign = :center,
    pointsize = 18, rotation = 90, color = :dodgerblue))]

npoints = 30
a2_ani_x = vcat(range(0,origins[2,1], length = npoints), origins[2,1].+zeros(nvec*npoints))
a2_ani_y = vcat(range(0,origins[2,2], length = npoints), origins[2,2].+zeros(nvec*npoints))

origens_ani = [[a2_ani_x, a2_ani_y]]

line_somax = vcat(zeros((nvec-1)*npoints), range(0, res[1], length = npoints), res[1].+zeros(npoints))
line_somay = vcat(zeros((nvec-1)*npoints), range(0, res[2], length = npoints), res[2].+zeros(npoints))

for i = 3: nvec
    ai_ani_x = vcat(zeros((i-2)*npoints), range(0,origins[i,1], length = npoints), origins[i,1].+zeros((nvec+2-i)*npoints))
    ai_ani_y = vcat(zeros((i-2)*npoints), range(0,origins[i,2], length = npoints), origins[i,2].+zeros((nvec+2-i)*npoints))
    push!(origens_ani, [ai_ani_x, ai_ani_y])
end

anim = @animate for i = 1: (nvec+1)*npoints
    quiver([0],[0], quiver = ([vecs[1,1]],[vecs[1,2]]), lw = 2, alpha = 0.5, framestyle = :grid,
        xaxis = ("direção x", xlimits, xlimits[1]:1:xlimits[2]),
        yaxis = ("direção y", ylimits, ylimits[1]:1:ylimits[2]), ann = ann_autor, size = (800,800))
     for j = 2:nvec
        quiver!([origens_ani[j-1][1][i]],[origens_ani[j-1][2][i]], quiver = ([vecs[j,1]],[vecs[j,2]]),
            lw = 2, alpha = 0.5)
     end
     quiver!([0],[0], quiver = ([line_somax[i]],[line_somay[i]]), lw = 4)

     if i > nvec*30
        annotate!(ann_res)
     end
end

gif(anim, "soma_vetores_ani.gif", fps = 25)

max_stars_repo_name: ghralima/mecanica

max_issues_repo_name: ghralima/mecanica

max_forks_repo_name: ghralima/mecanica


text:
include("test-sorting-networks.jl")
include("test-transpose.jl")
include("test-bitonic-merge.jl")
include("test-data-buffers.jl")
include("test-chipsort.jl")


text:
abstract type MultiFieldStyle end

struct ConsecutiveMultiFieldStyle <: MultiFieldStyle end

struct StridedMultiFieldStyle <: MultiFieldStyle end

"""
    struct MultiFieldFESpace{S<:MultiFieldStyle,B} <: FESpace
      spaces::Vector{<:SingleFieldFESpace}
      multi_field_style::S
      constraint_style::Val{B}
    end
"""
struct MultiFieldFESpace{S<:MultiFieldStyle,B} <: FESpace
  spaces::Vector{<:SingleFieldFESpace}
  multi_field_style::S
  constraint_style::Val{B}
  function MultiFieldFESpace(spaces::Vector{<:SingleFieldFESpace},mfs::MultiFieldStyle)
    S = typeof(mfs)
    B = any( map(has_constraints,spaces) )
    cs = Val{B}()
    new{S,B}(spaces,mfs,cs)
  end
end

"""
    MultiFieldFESpace(spaces::Vector{<:SingleFieldFESpace})
"""
function MultiFieldFESpace(spaces::Vector{<:SingleFieldFESpace})
  MultiFieldFESpace(spaces,ConsecutiveMultiFieldStyle())
end

MultiFieldStyle(::Type{MultiFieldFESpace{S,B}}) where {S,B} = S()

MultiFieldStyle(f::MultiFieldFESpace) = MultiFieldStyle(typeof(f))

# Implementation of FESpace

FESpaces.constraint_style(::Type{MultiFieldFESpace{S,B}}) where {S,B} = Val{B}()

function FESpaces.num_free_dofs(f::MultiFieldFESpace)
  n = 0
  for U in f.spaces
    n += num_free_dofs(U)
  end
  n
end

function FESpaces.get_cell_axes(f::MultiFieldFESpace)
  all_axes = map(get_cell_axes,f.spaces)
  create_array_of_blocked_axes(all_axes...)
end

function FESpaces.get_cell_axes_with_constraints(f::MultiFieldFESpace)
  all_axes = map(get_cell_axes_with_constraints,f.spaces)
  create_array_of_blocked_axes(all_axes...)
end

function FESpaces.get_cell_basis(f::MultiFieldFESpace)
  all_axes = map(i->get_cell_axes(get_cell_basis(i)),f.spaces)
  all_bases = [
    _insert_cell_basis_in_block(i,get_cell_basis(f.spaces[i]),all_axes)
    for i in 1:length(f.spaces) ]
  MultiFieldCellField(all_bases)
end

function _insert_cell_basis_in_block(i::Int,cb,all_axes)
  array =  insert_array_of_bases_in_block(i,get_array(cb),all_axes...)
  similar_object(cb,array,array.axes,MetaSizeStyle(cb))
end

function FESpaces.FEFunction(fe::MultiFieldFESpace, free_values)
  blocks = SingleFieldFEFunction[]
  for (field, U) in enumerate(fe.spaces)
    free_values_i = restrict_to_field(fe,free_values,field)
    uhi = FEFunction(U,free_values_i)
    push!(blocks,uhi)
  end
  MultiFieldFEFunction(free_values,fe,blocks)
end

function FESpaces.EvaluationFunction(fe::MultiFieldFESpace, free_values)
  blocks = SingleFieldFEFunction[]
  for (field, U) in enumerate(fe.spaces)
    free_values_i = restrict_to_field(fe,free_values,field)
    uhi = EvaluationFunction(U,free_values_i)
    push!(blocks,uhi)
  end
  MultiFieldFEFunction(free_values,fe,blocks)
end

function CellData.CellField(fe::MultiFieldFESpace,cell_values)
  single_fields = CellField[]
  for field in 1:num_fields(fe)
    cell_values_field = apply(a->a[Block(field)],cell_values)
    cf = CellField(fe.spaces[field],cell_values_field)
    push!(single_fields,cf)
  end
  MultiFieldCellField(single_fields)
end

function CellData.CellField(fe::MultiFieldFESpace,cell_values::VectorOfBlockArrayCoo)
  single_fields = CellField[]
  for field in 1:num_fields(fe)
    cf = CellField(fe.spaces[field],cell_values[Block(field)])
    push!(single_fields,cf)
  end
  MultiFieldCellField(single_fields)
end

"""
    restrict_to_field(f::MultiFieldFESpace,free_values::AbstractVector,field::Integer)
"""
function restrict_to_field(f::MultiFieldFESpace,free_values::AbstractVector,field::Integer)
  _restrict_to_field(f,MultiFieldStyle(f),free_values,field)
end

function  _restrict_to_field(f,::MultiFieldStyle,free_values,field)
  @notimplemented
end

function  _restrict_to_field(f,::ConsecutiveMultiFieldStyle,free_values,field)
  offsets = compute_field_offsets(f)
  U = f.spaces
  pini = offsets[field] + 1
  pend = offsets[field] + num_free_dofs(U[field])
  SubVector(free_values,pini,pend)
end

"""
    compute_field_offsets(f::MultiFieldFESpace)
"""
function compute_field_offsets(f::MultiFieldFESpace)
  @assert MultiFieldStyle(f) == ConsecutiveMultiFieldStyle()
  U = f.spaces
  n = length(U)
  offsets = zeros(Int,n)
  for i in 1:(n-1)
    Ui = U[i]
    offsets[i+1] = offsets[i] + num_free_dofs(Ui)
  end
  offsets
end

function FESpaces.zero_free_values(fs::MultiFieldFESpace)
  zeros(num_free_dofs(fs))
end

function FESpaces.get_cell_isconstrained(f::MultiFieldFESpace)
  data = map(get_cell_isconstrained,f.spaces)
  apply( (args...) -> +(args...)>0,  data...)
end

function FESpaces.get_cell_constraints(f::MultiFieldFESpace)
  blocks = Tuple(map(get_cell_constraints,f.spaces))
  blockids = [ (i,i) for i in 1:length(f.spaces)]
  all_axs_rows = map(get_cell_axes_with_constraints,f.spaces)
  all_axs_cols = map(get_cell_axes,f.spaces)
  axs_rows = create_array_of_blocked_axes(all_axs_rows...)
  axs_cols = create_array_of_blocked_axes(all_axs_cols...)
  axs = apply((r,c) -> (r[1],c[1]),axs_rows,axs_cols)
  VectorOfBlockArrayCoo(blocks,blockids,axs)
end

function FESpaces.get_cell_dofs(f::MultiFieldFESpace)
  _get_cell_dofs(f,MultiFieldStyle(f))
end

function _get_cell_dofs(f,::MultiFieldStyle)
  @notimplemented
end

function _get_cell_dofs(f,::ConsecutiveMultiFieldStyle)
  offsets = compute_field_offsets(f)
  spaces = f.spaces
  function fun(i,space)
    cell_dofs = get_cell_dofs(space)
    if i == 1
      return cell_dofs
    end
    offset = offsets[i]
    o = Fill(offset,length(cell_dofs))
    apply(elem(_sum_if_first_positive),cell_dofs,o)
  end
  blocks = Tuple([ fun(i,space) for (i,space) in enumerate(spaces) ])
  blockids = [ (i,) for i in 1:length(spaces)]
  axs = get_cell_axes(f)
  VectorOfBlockArrayCoo(blocks,blockids,axs)
end

function _sum_if_first_positive(a,b)
  if a>0
    return a+b
  else
    return a
  end
end

# API for multi field case

"""
    num_fields(f::MultiFieldFESpace)
"""
function num_fields(f::MultiFieldFESpace)
  length(f.spaces)
end

Base.iterate(m::MultiFieldFESpace) = iterate(m.spaces)

Base.iterate(m::MultiFieldFESpace,state) = iterate(m.spaces,state)

Base.getindex(m::MultiFieldFESpace,field_id::Integer) = m.spaces[field_id]

# API for the ConsecutiveMultiFieldStyle
import Gridap.FESpaces: interpolate
import Gridap.FESpaces: interpolate_everywhere
import Gridap.FESpaces: interpolate_dirichlet

@deprecate(
  interpolate(fs::MultiFieldFESpace, object),
  interpolate(object, fs::MultiFieldFESpace)
)

@deprecate(
  interpolate_everywhere(fs::MultiFieldFESpace, object),
  interpolate_everywhere(object, fs::MultiFieldFESpace)
)

@deprecate(
  interpolate_dirichlet(fs::MultiFieldFESpace, object),
  interpolate_dirichlet(object, fs::MultiFieldFESpace)
)

"""
The resulting MultiFieldFEFunction is in the space (in particular it fulfills Dirichlet BCs
even in the case that the given cell field does not fulfill them)
"""
function FESpaces.interpolate(objects, fe::MultiFieldFESpace)
  free_values = zero_free_values(fe)
  blocks = SingleFieldFEFunction[]
  for (field, (U,object)) in enumerate(zip(fe.spaces,objects))
    free_values_i = restrict_to_field(fe,free_values,field)
    uhi = interpolate!(object, free_values_i,U)
    push!(blocks,uhi)
  end
  MultiFieldFEFunction(free_values,fe,blocks)
end

"""
like interpolate, but also compute new degrees of freedom for the dirichlet component.
The resulting MultiFieldFEFunction does not necessary belongs to the underlying space
"""
function FESpaces.interpolate_everywhere(objects, fe::MultiFieldFESpace)
  free_values = zero_free_values(fe)
  blocks = SingleFieldFEFunction[]
  for (field, (U,object)) in enumerate(zip(fe.spaces,objects))
    free_values_i = restrict_to_field(fe,free_values,field)
    dirichlet_values_i = zero_dirichlet_values(U)
    uhi = interpolate_everywhere!(object, free_values_i,dirichlet_values_i,U)
    push!(blocks,uhi)
  end
  MultiFieldFEFunction(free_values,fe,blocks)
end

"""
"""
function FESpaces.interpolate_dirichlet(objects, fe::MultiFieldFESpace)
  free_values = zero_free_values(fe)
  blocks = SingleFieldFEFunction[]
  for (field, (U,object)) in enumerate(zip(fe.spaces,objects))
    free_values_i = restrict_to_field(fe,free_values,field)
    dirichlet_values_i = zero_dirichlet_values(U)
    uhi = interpolate_dirichlet!(object, free_values_i,dirichlet_values_i,U)
    push!(blocks,uhi)
  end
  MultiFieldFEFunction(free_values,fe,blocks)
end


text:
export BrickletColorColor
struct BrickletColorColor
    r::Integer
    g::Integer
    b::Integer
    c::Integer
end

export BrickletColorColorCallbackThreshold
struct BrickletColorColorCallbackThreshold
    option::Char
    min_r::Integer
    max_r::Integer
    min_g::Integer
    max_g::Integer
    min_b::Integer
    max_b::Integer
    min_c::Integer
    max_c::Integer
end

export BrickletColorConfig
struct BrickletColorConfig
    gain::Integer
    integration_time::Integer
end

export BrickletColorIdentity
struct BrickletColorIdentity
    uid::String
    connected_uid::String
    position::Char
    hardware_version::Vector{Integer}
    firmware_version::Vector{Integer}
    device_identifier::Integer
end

export BrickletColor
"""
Measures color (RGB value), illuminance and color temperature
"""
mutable struct BrickletColor <: TinkerforgeDevice
    ipcon::IPConnection
    deviceInternal::PyObject

    """
    Creates an object with the unique device ID *uid* and adds it to
    the IP Connection *ipcon*.
    """
    function BrickletColor(uid::String, ipcon::IPConnection)
        package = pyimport("tinkerforge.bricklet_color")
        deviceInternal = package.BrickletColor(uid, ipcon.ipconInternal)

        return new(ipcon, deviceInternal)
    end
end

export get_color
"""
    $(SIGNATURES)

Returns the measured color of the sensor.

The red (r), green (g), blue (b) and clear (c) colors are measured
with four different photodiodes that are responsive at different
wavelengths:

.. image:: /Images/Bricklets/bricklet_color_wavelength_chart_600.jpg
   :scale: 100 %
   :alt: Chart Responsivity / Wavelength
   :align: center
   :target: ../../_images/Bricklets/bricklet_color_wavelength_chart_600.jpg

If you want to get the color periodically, it is recommended
to use the :cb:`Color` callback and set the period with
:func:`Set Color Callback Period`.
"""
function get_color(device::BrickletColor)
    return device.deviceInternal.get_color()
end

export set_color_callback_period
"""
    $(SIGNATURES)

Sets the period with which the :cb:`Color` callback is triggered
periodically. A value of 0 turns the callback off.

The :cb:`Color` callback is only triggered if the color has changed since the
last triggering.
"""
function set_color_callback_period(device::BrickletColor, period)
    device.deviceInternal.set_color_callback_period(period)
end

export get_color_callback_period
"""
    $(SIGNATURES)

Returns the period as set by :func:`Set Color Callback Period`.
"""
function get_color_callback_period(device::BrickletColor)
    return device.deviceInternal.get_color_callback_period()
end

export set_color_callback_threshold
"""
    $(SIGNATURES)

Sets the thresholds for the :cb:`Color Reached` callback.

The following options are possible:

.. csv-table::
 :header: "Option", "Description"
 :widths: 10, 100

 "'x'",    "Callback is turned off"
 "'o'",    "Callback is triggered when the temperature is *outside* the min and max values"
 "'i'",    "Callback is triggered when the temperature is *inside* the min and max values"
 "'<'",    "Callback is triggered when the temperature is smaller than the min value (max is ignored)"
 "'>'",    "Callback is triggered when the temperature is greater than the min value (max is ignored)"
"""
function set_color_callback_threshold(device::BrickletColor, option, min_r, max_r, min_g, max_g, min_b, max_b, min_c, max_c)
    device.deviceInternal.set_color_callback_threshold(option, min_r, max_r, min_g, max_g, min_b, max_b, min_c, max_c)
end

export get_color_callback_threshold
"""
    $(SIGNATURES)

Returns the threshold as set by :func:`Set Color Callback Threshold`.
"""
function get_color_callback_threshold(device::BrickletColor)
    return device.deviceInternal.get_color_callback_threshold()
end

export set_debounce_period
"""
    $(SIGNATURES)

Sets the period with which the threshold callback

* :cb:`Color Reached`

is triggered, if the threshold

* :func:`Set Color Callback Threshold`

keeps being reached.
"""
function set_debounce_period(device::BrickletColor, debounce)
    device.deviceInternal.set_debounce_period(debounce)
end

export get_debounce_period
"""
    $(SIGNATURES)

Returns the debounce period as set by :func:`Set Debounce Period`.
"""
function get_debounce_period(device::BrickletColor)
    return device.deviceInternal.get_debounce_period()
end

export light_on
"""
    $(SIGNATURES)

Turns the LED on.
"""
function light_on(device::BrickletColor)
    device.deviceInternal.light_on()
end

export light_off
"""
    $(SIGNATURES)

Turns the LED off.
"""
function light_off(device::BrickletColor)
    device.deviceInternal.light_off()
end

export is_light_on
"""
    $(SIGNATURES)

Returns the state of the LED. Possible values are:

* 0: On
* 1: Off
"""
function is_light_on(device::BrickletColor)
    return device.deviceInternal.is_light_on()
end

export set_config
"""
    $(SIGNATURES)

Sets the configuration of the sensor. Gain and integration time
can be configured in this way.

For configuring the gain:

* 0: 1x Gain
* 1: 4x Gain
* 2: 16x Gain
* 3: 60x Gain

For configuring the integration time:

* 0: 2.4ms
* 1: 24ms
* 2: 101ms
* 3: 154ms
* 4: 700ms

Increasing the gain enables the sensor to detect a
color from a higher distance.

The integration time provides a trade-off between conversion time
and accuracy. With a longer integration time the values read will
be more accurate but it will take longer time to get the conversion
results.
"""
function set_config(device::BrickletColor, gain, integration_time)
    device.deviceInternal.set_config(gain, integration_time)
end

export get_config
"""
    $(SIGNATURES)

Returns the configuration as set by :func:`Set Config`.
"""
function get_config(device::BrickletColor)
    return device.deviceInternal.get_config()
end

export get_illuminance
"""
    $(SIGNATURES)

Returns the illuminance affected by the gain and integration time as
set by :func:`Set Config`. To get the illuminance in Lux apply this formula::

 lux = illuminance * 700 / gain / integration_time

To get a correct illuminance measurement make sure that the color
values themselves are not saturated. The color value (R, G or B)
is saturated if it is equal to the maximum value of 65535.
In that case you have to reduce the gain, see :func:`Set Config`.
"""
function get_illuminance(device::BrickletColor)
    return device.deviceInternal.get_illuminance()
end

export get_color_temperature
"""
    $(SIGNATURES)

Returns the color temperature.

To get a correct color temperature measurement make sure that the color
values themselves are not saturated. The color value (R, G or B)
is saturated if it is equal to the maximum value of 65535.
In that case you have to reduce the gain, see :func:`Set Config`.
"""
function get_color_temperature(device::BrickletColor)
    return device.deviceInternal.get_color_temperature()
end

export set_illuminance_callback_period
"""
    $(SIGNATURES)

Sets the period with which the :cb:`Illuminance` callback is triggered
periodically. A value of 0 turns the callback off.

The :cb:`Illuminance` callback is only triggered if the illuminance has changed
since the last triggering.
"""
function set_illuminance_callback_period(device::BrickletColor, period)
    device.deviceInternal.set_illuminance_callback_period(period)
end

export get_illuminance_callback_period
"""
    $(SIGNATURES)

Returns the period as set by :func:`Set Illuminance Callback Period`.
"""
function get_illuminance_callback_period(device::BrickletColor)
    return device.deviceInternal.get_illuminance_callback_period()
end

export set_color_temperature_callback_period
"""
    $(SIGNATURES)

Sets the period with which the :cb:`Color Temperature` callback is
triggered periodically. A value of 0 turns the callback off.

The :cb:`Color Temperature` callback is only triggered if the color temperature
has changed since the last triggering.
"""
function set_color_temperature_callback_period(device::BrickletColor, period)
    device.deviceInternal.set_color_temperature_callback_period(period)
end

export get_color_temperature_callback_period
"""
    $(SIGNATURES)

Returns the period as set by :func:`Set Color Temperature Callback Period`.
"""
function get_color_temperature_callback_period(device::BrickletColor)
    return device.deviceInternal.get_color_temperature_callback_period()
end

export get_identity
"""
    $(SIGNATURES)

Returns the UID, the UID where the Bricklet is connected to,
the position, the hardware and firmware version as well as the
device identifier.

The position can be 'a', 'b', 'c', 'd', 'e', 'f', 'g' or 'h' (Bricklet Port).
A Bricklet connected to an :ref:`Isolator Bricklet <isolator_bricklet>` is always at
position 'z'.

The device identifier numbers can be found :ref:`here <device_identifier>`.
|device_identifier_constant|
"""
function get_identity(device::BrickletColor)
    return device.deviceInternal.get_identity()
end

export register_callback
"""
Registers the given *function* with the given *callback_id*.
"""
function register_callback(device::BrickletColor, callback_id, function_)
    if isnothing(function_)
        device.registered_callbacks.pop(callback_id, None)
    else
        device.registered_callbacks[callback_id] = function_
    end
end