text:
# A recreation of <https://gpflow.readthedocs.io/en/master/notebooks/advanced/gps_for_big_data.html>

# # Regression: Sparse Variational Gaussian Process for Stochastic Optimisation with Flux.jl

# In this example, we show how to construct and train the stochastic variational
# Gaussian process (SVGP) model for efficient inference in large scale datasets.
# For a basic introduction to the functionality of this library, please refer to
# the [User Guide](@ref).
#
# ## Setup

using ApproximateGPs
using Distributions
using LinearAlgebra

using Plots
## We output plots as PNG to reduce their size
default(; fmt=:png, palette=:seaborn_colorblind, legend=:outertopright, size=(700, 400))
#nb ## Regardless of the `fmt` setting, by default Literate saves also the SVG output
#nb ## in notebooks: https://github.com/fredrikekre/Literate.jl/issues/61
#nb ## We use DisplayAs to ensure that only the PNG output is saved.
#nb using DisplayAs

using Random
Random.seed!(1234);

# ## Generate some training data
#
# We define a data-generating function `g`:

g(x) = sin(3π * x) + 0.3 * cos(9π * x) + 0.5 * sin(7π * x)

N = 10000 # Number of training points
x = rand(Uniform(-1, 1), N)
y = g.(x) + 0.3 * randn(N)

scatter(x, y; xlabel="x", ylabel="y", markershape=:xcross, markeralpha=0.1, legend=false)
#nb DisplayAs.PNG(current())

# ## Set up a Flux model
#
# We shall use the excellent framework provided by [Flux.jl](https://fluxml.ai/)
# to perform stochastic optimisation. The SVGP approximation has three sets of
# parameters to optimise - the inducing input locations, the mean and covariance
# of the variational distribution `q` and the parameters of the
# kernel.
#
# First, we define a helper function to construct the kernel from its
# parameters (often called kernel hyperparameters), and pick some
# initial values `k_init`. Since Flux performs optimisation on
# unconstrained parameters, we need to use softplus to ensure that the
# kernel parameters are positive.

using StatsFuns: softplus, invsoftplus

function make_kernel(k_params)
    variance = softplus(k_params[1])
    lengthscale = softplus(k_params[2])
    return variance * with_lengthscale(SqExponentialKernel(), lengthscale)
end

init_variance = 1.3
init_lengthscale = 0.3
k_init = [invsoftplus(init_variance), invsoftplus(init_lengthscale)];

# Then, we select some inducing input locations `z_init`. In this
# case, we simply choose the first `M` data inputs.

M = 20 # number of inducing points
z_init = x[1:M];

# Given a set of parameters, we now define a Flux 'layer' which forms
# the basis of our model.

using Flux

struct SVGPModel
    k  # kernel parameters
    z  # inducing points
    m  # variational mean
    A  # square-root of variational covariance
end

Flux.@functor SVGPModel (k, z, m, A);

# Set the observation noise for our model, along with a `jitter` term
# to help with numerical stability.

lik_noise = 0.3
jitter = 1e-5;

# Next, we define some useful functions on the model - creating the prior GP
# under the model, as well as the `SparseVariationalApproximation` struct
# needed to create the posterior approximation and to compute the ELBO.

function prior(m::SVGPModel)
    kernel = make_kernel(m.k)
    return GP(kernel)
end

# The variational distribution is given by `q(u)` where ``u ~ f(z)`` are the
# pseudo-points. We parameterise the covariance matrix of ``q`` as ``S = A A^T``
# since this guarantees that ``S`` is positive definite. We also only use the
# lower triangular part of `A`, to ensure the minimum number of free parameters.

using PDMats: PDMat

function make_approx(m::SVGPModel, prior)
    ## Efficiently constructs S as A*Aᵀ
    S = PDMat(Cholesky(LowerTriangular(m.A)))
    q = MvNormal(m.m, S)
    fz = prior(m.z, jitter)
    return SparseVariationalApproximation(fz, q)
end;

# Create the approximate posterior GP under the model.

function model_posterior(m::SVGPModel)
    svgp = make_approx(m, prior(m))
    return posterior(svgp)
end;

# Define a predictive function for the model - in this case the prediction is
# the joint distribution of the approximate posterior GP at some test inputs `x`
# (defined by an `AbstractGPs.FiniteGP`).

function (m::SVGPModel)(x)
    post = model_posterior(m)
    return post(x)
end;

# Return the loss given data - for the SVGP model as constructed using `SparseVariationalApproximation`, the
# loss used is the negative ELBO (also known as the Variational Free Energy).
# `num_data` is required for minibatching used below.

function loss(m::SVGPModel, x, y; num_data=length(y))
    f = prior(m)
    fx = f(x, lik_noise)
    svgp = make_approx(m, f)
    return -elbo(svgp, fx, y; num_data)
end;

# Finally, we choose some initial parameters and instantiate our model.

m_init = zeros(M)
A_init = Matrix{Float64}(I, M, M)

model = SVGPModel(k_init, z_init, m_init, A_init);

# Taking a look at the model posterior under these initial parameters shows a
# very poor fit to the data, as expected:

init_post = model_posterior(model)
scatter(
    x,
    y;
    xlabel="x",
    ylabel="y",
    markershape=:xcross,
    markeralpha=0.1,
    label="Training Data",
)
plot!(-1:0.001:1, init_post; label="Initial Posterior", color=4)
#nb DisplayAs.PNG(current())

# ## Training the model
#
# Training the model now simply proceeds with the usual `Flux.jl` training loop.

opt = ADAM(0.001)  # Define the optimiser
params = Flux.params(model);  # Extract the model parameters

# One of the major advantages of the SVGP model is that it allows stochastic
# estimation of the ELBO by using minibatching of the training data. This is
# very straightforward to achieve with `Flux.jl`'s utilities:

b = 100 # minibatch size
data_loader = Flux.Data.DataLoader((x, y); batchsize=b)

# The loss (negative ELBO) before training

loss(model, x, y)

# Train the model. N.B. when using minibatching, the length of the
# full dataset `num_data` must be passed to the loss.

using IterTools: ncycle

Flux.train!(
    (x, y) -> loss(model, x, y; num_data=N),
    params,
    ncycle(data_loader, 300), # Train for 300 epochs
    opt,
);

# Negative ELBO after training

loss(model, x, y)

# Finally, we plot the optimised approximate posterior to see the
# results.

post = model_posterior(model)

scatter(
    x,
    y;
    markershape=:xcross,
    markeralpha=0.1,
    xlim=(-1, 1),
    xlabel="x",
    ylabel="y",
    title="posterior (VI with sparse grid)",
    label="Training Data",
    color=1,
)
plot!(-1:0.001:1, post; label="Posterior", color=4)
sticks!(model.z, fill(0.13, M); label="Pseudo-points", linewidth=1.5, color=5)
#nb DisplayAs.PNG(current())


text:
function prefix_sum(arr::Vector{T}) where {T<:Number}
    pre = []
    preans = zero(T)
    for i in arr
        preans += i
        push!(pre, preans)
    end
    return pre
end

ext: jl

max_issues_repo_name: AugustoCL/Julia

max_forks_repo_name: AugustoCL/Julia


text:
using Test
using LogicCircuits
using LogicCircuits.LoadSave: parse_lc_file, save_lines, parse_psdd_file, parse_sdd_file, parse_vtree_file, zoo_lc_file, zoo_psdd_file, zoo_sdd_file, zoo_vtree_file

@testset "Parse-save-parse test" begin
    mktempdir() do tmp

        lines = parse_lc_file(zoo_lc_file("little_4var.circuit"))
        save_lines("$tmp/temp.circuit", lines)
        lines2 = parse_lc_file("$tmp/temp.circuit")
        @test length(lines) == length(lines2)

        lines = parse_psdd_file(zoo_psdd_file("little_4var.psdd"))
        save_lines("$tmp/temp.psdd", lines)
        lines2 = parse_psdd_file("$tmp/temp.psdd")
        @test length(lines) == length(lines2)

        lines = parse_sdd_file(zoo_sdd_file("random.sdd"))
        save_lines("$tmp/temp.sdd", lines)
        lines2 = parse_sdd_file("$tmp/temp.sdd")
        @test length(lines) == length(lines2)

        lines = parse_vtree_file(zoo_vtree_file("little_4var.vtree"))
        # TODO split up vtree saver into decompilation phase and file writing phase so we can test like above
        @test !isempty(lines)

    end
end


text:
struct PRParam <: EoSParam
    a::PairParam{Float64}
    b::PairParam{Float64}
    acentricfactor::SingleParam{Float64}
    Tc::SingleParam{Float64}
    Pc::SingleParam{Float64}
    Mw::SingleParam{Float64}
end

abstract type PRModel <: ABCubicModel end
@newmodel PR PRModel PRParam

export PR
function PR(components::Array{String,1}; userlocations::Array{String,1}=String[], idealmodel = BasicIdeal,verbose=false)
    params = getparams(components, ["properties/critical.csv", "properties/molarmass.csv","SAFT/PCSAFT/PCSAFT_unlike.csv"]; userlocations=userlocations, verbose=verbose)
    Mw = params["Mw"]
    k  = params["k"]
    _pc = params["pc"]
    pc = _pc.values
    Tc = params["Tc"]
    Tc_ = Tc.values
    acentricfactor = params["w"]
    Ωa,Ωb = ab_consts(PRModel)
    a = epsilon_LorentzBerthelot(SingleParam(params["pc"], @. Ωa*R̄^2*Tc_^2/pc), k)
    b = sigma_LorentzBerthelot(SingleParam(params["pc"], @. Ωb*R̄*Tc_/pc))

    packagedparams = PRParam(a, b, acentricfactor, Tc,_pc,Mw)
    model = PR(packagedparams,idealmodel)
    return model
end

function ab_consts(::Type{<:PRModel})
    return 0.457235,0.077796
end

function cubic_ab(model::PRModel,T,z=SA[1.0],n=sum(z))
    a = model.params.a.values
    b = model.params.b.values
    ω = model.params.acentricfactor.values
    Tc = model.params.Tc.values
    invn = one(n)/n
    αx = @. min((1+(0.37464+1.54226*ω-0.26992*ω^2)*(1-√(T/Tc)))^2,one(T)) * z * invn
    āᾱ = dot(αx,Symmetric(a),αx)
    b̄ = dot(z,Symmetric(b),z) * invn*invn
    return āᾱ ,b̄
end

function cubic_abp(model::PRModel, V, T, z)
    n = sum(z)
    v = V/n
    āᾱ ,b̄ = cubic_ab(model,T,z,n)
    _1 = one(b̄)
    denom = evalpoly(v,(-b̄*b̄,2*b̄,_1))
    p = R̄*T/(v-b̄) - āᾱ /denom
    return āᾱ, b̄,p
end

function cubic_poly(model::PRModel,p,T,z)
    a,b = cubic_ab(model,T,z)
    RT⁻¹ = 1/(R̄*T)
    A = a*p*RT⁻¹*RT⁻¹
    B = b*p*RT⁻¹
    k₀ = B*(B*(B+1.0)-A)
    k₁ = -B*(3*B+2.0) + A
    k₂ = B-1.0
    k₃ = 1.0
    return [k₀,k₁,k₂,k₃]
end

#=
 (-B2-2(B2+B)+A)
 (-B2-2B2-2B+A)
 (-3B2-2B+A)
=#
function a_res(model::PRModel, V, T, z)
    n = sum(z)
    a,b = cubic_ab(model,T,z,n)
    Δ1 = 1+√2
    Δ2 = 1-√2
    ΔPRΔ = 2*√2
    RT⁻¹ = 1/(R̄*T)
    ρ = n/V
    return -log(1-b*ρ) - a*RT⁻¹*log((Δ1*b*ρ+1)/(Δ2*b*ρ+1))/(ΔPRΔ*b)

    #return -log(V-n*b̄) + āᾱ/(R̄*T*b̄*2^(3/2)) * log((2*V-2^(3/2)*b̄*n+2*b̄*n)/(2*V+2^(3/2)*b̄*n+2*b̄*n))
end

cubic_zc(::PRModel) = 0.3074


text:
using Test
using KernelAbstractions
using ClimateMachine.Thermodynamics
using ClimateMachine.TemperatureProfiles
using UnPack
using CUDA
using Random
using RootSolvers
const TD = Thermodynamics

using LinearAlgebra
using CLIMAParameters
using CLIMAParameters.Planet

struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

using ClimateMachine
ClimateMachine.init()
ArrayType = ClimateMachine.array_type()

@show ArrayType

device(::Type{T}) where {T <: Array} = CPU()
device(::Type{T}) where {T <: CuArray} = CUDADevice()

include("profiles.jl")

@kernel function test_thermo_kernel!(
    param_set,
    dst::AbstractArray{FT},
    e_int,
    ρ,
    p,
    q_tot,
) where {FT}
    i = @index(Group, Linear)
    @inbounds begin

        ts = PhaseEquil(param_set, FT(e_int[i]), FT(ρ[i]), FT(q_tot[i]))
        dst[1, i] = air_temperature(ts)

        ts_ρpq = PhaseEquil_ρpq(
            param_set,
            FT(ρ[i]),
            FT(p[i]),
            FT(q_tot[i]),
            true,
            100,
            RegulaFalsiMethod,
        )
        dst[2, i] = air_temperature(ts_ρpq)
    end
end

@testset "Thermodynamics - kernels" begin
    FT = Float32
    dev = device(ArrayType)
    profiles = PhaseEquilProfiles(param_set, Array)
    slice = Colon()
    profiles = convert_profile_set(profiles, ArrayType, slice)

    n_profiles = length(profiles.z)
    n_vars = length(propertynames(profiles))
    d_dst = ArrayType(Array{FT}(undef, 2, n_profiles))
    fill!(d_dst, 0)

    @unpack e_int, ρ, p, q_tot = profiles

    work_groups = (1,)
    ndrange = (n_profiles,)
    kernel! = test_thermo_kernel!(dev, work_groups)
    event = kernel!(param_set, d_dst, e_int, ρ, p, q_tot, ndrange = ndrange)
    wait(dev, event)

    ts_correct = PhaseEquil.(param_set, Array(e_int), Array(ρ), Array(q_tot))
    @test all(Array(d_dst)[1, :] .≈ air_temperature.(ts_correct))

    ts_correct =
        PhaseEquil_ρpq.(
            param_set,
            Array(ρ),
            Array(p),
            Array(q_tot),
            true,
            100,
            RegulaFalsiMethod,
        )
    @test all(Array(d_dst)[2, :] .≈ air_temperature.(ts_correct))

end

max_stars_repo_name: jm-c/CLIMA

max_issues_repo_name: jm-c/CLIMA

max_forks_repo_name: jm-c/CLIMA


text:
struct IndexToLocation{T}
    offsets::Vector{T}
end

function Base.parse(io::IO, ::Type{IndexToLocation{T}}, maxp::MaximumProfile) where {T}
    offsets = map(0:maxp.nglyphs) do i
        read(io, T)
    end
    IndexToLocation{T}(offsets)
end

function Base.parse(io::IO, ::Type{IndexToLocation}, maxp::MaximumProfile, head::FontHeader)
    head.index_to_loc_format in (0, 1) || error("Index to location format must be either 0 or 1.")
    T = head.index_to_loc_format == 0 ? UInt16 : UInt32
    parse(io, IndexToLocation{T}, maxp)
end

function glyph_ranges(loca::IndexToLocation)
    [start:finish for (start, finish) in zip(loca.offsets[1:end-1], loca.offsets[2:end])]
end


text:
using DelimitedFiles
using LinearAlgebra
using FastGaussQuadrature
include("Quadrature.jl")
include("standardQuadrature.jl")

# gauss
for n=2:2:100
	p,w = Quadrature(n,1)
	p = [p; p[:,1] p[:,2] -p[:,3]]
	w = [w;w]

	m = [p w]
	display(n)
	name = string(n)*"_gauss.txt"
	writedlm(name,m)
end
#
## octa
#for n=2:2:50
#	p,w = Quadrature(n,2)
#	m = [p w]
#	display(n)
#	name = string(n)*"_s_octa.txt"
#	writedlm(name,m)
#end
#
## ico
#for n=2:2:50
#	p,w = Quadrature(n,3)
#	m = [p w]
#	display(n)
#	name = string(n)*"_s_ico.txt"
#	writedlm(name,m)
#end

max_stars_repo_name: vavrines/CSN

max_issues_repo_name: vavrines/CSN

max_forks_repo_name: vavrines/CSN


text:
module ExplorerHat
import JuliaBerry.Motor
import JuliaBerry.LED
import JuliaBerry.OutputPin
import JuliaBerry.InputPin

# Onboard LEDs above 1, 2, 3, 4
const LED1 = 4
const LED2 = 17
const LED3 = 27
const LED4 = 5

# Outputs via ULN2003A
const OUT1 = 6
const OUT2 = 12
const OUT3 = 13
const OUT4 = 16

# 5v Tolerant Inputs
const IN1 = 23
const IN2 = 22
const IN3 = 24
const IN4 = 25

# Motor, via DRV8833PWP Dual H-Bridge
const M1B = 19
const M1F = 20
const M2B = 21
const M2F = 26

# Number of times to update
# pulsing LEDs per second
const PULSE_FPS = 50
const PULSE_FREQUENCY = 1000

const DEBOUNCE_TIME = 20

const CAP_PRODUCT_ID = 107

function __init__()
    global  motor = [
        Motor(M1F, M1B),
        Motor(M2F, M2B)
    ]

    global  led = [
        LED(LED1),
        LED(LED2),
        LED(LED3),
        LED(LED4)
    ]

    global  output = [
        OutputPin(OUT1),
        OutputPin(OUT2),
        OutputPin(OUT3),
        OutputPin(OUT4)
    ]

    global  input = [
        InputPin(OUT1),
        InputPin(OUT2),
        InputPin(OUT3),
        InputPin(OUT4)
    ]

end
end


text:
sparse_matrix(g::SparseMatrixCSC{Complex{Float64},Int}) = g

"""
    Circuit{N}

Quantum circuit consisting of a unitary gate chain and measurement operators with `N` qubits.
"""
struct Circuit{N}
    moments::Vector{Moment}
    meas::Vector{MeasurementOperator}

    @doc """
        Circuit{N}(mops::Union{Vector{<:MeasurementOperator},Nothing}=nothing)

    Chain of quantum circuit gates in a circuit of `N` qubits. Constructed from vector of `CircuitGate` objects.
    """
    function Circuit{N}(mops::Vector{<:MeasurementOperator}=MeasurementOperator[]) where {N}
        if isempty(mops)
            return new{N}(Moment[], mops)
        end
        mops_N = maximum(req_wires.(mops))
        mops_N <= N || error("Measurement operators affecting $mops_N wires provided for Circuit of size $N")
        check_commute(mops) || error("Measurement operators do not commute")
        return new{N}(Moment[], mops)
    end

    @doc """
        Circuit{N}(gates::Vector{<:CircuitGate}, mops::Union{Vector{<:MeasurementOperator},Nothing}=nothing)

    Chain of quantum circuit gates in a circuit of `N` qubits. Constructed from vector of `CircuitGate` objects.
    """
    function Circuit{N}(gates::Vector{<:CircuitGate}, mops::Vector{<:MeasurementOperator}=MeasurementOperator[]) where {N}
        j = 1
        iwires = falses(N)
        moments = Moment[]
        for (i, cg) in enumerate(gates)
            cgwires = collect(cg.iwire)
            all(cgwires .<= N) || error("Unable to add CircuitGate with `iwire`: $(cg.iwire), with maximum circuit size `N`: $N")
            if any(iwires[cgwires])
                push!(moments, Moment(gates[j:i-1]))
                j = i
                iwires[:] .= false
            end
            iwires[cgwires] .= true
        end
        push!(moments, Moment(gates[j:end]))

        if isempty(mops)
            return new{N}(moments, mops)
        end
        mops_N = maximum(size.(mops))
        mops_N <= N || error("Measurement operators affecting $mops_N wires provided for Circuit of size $N")
        check_commute(mops) || error("Measurement operators do not commute")
        return new{N}(moments, mops)
    end

    @doc """
        Circuit{N}(gate::AbstractCircuitGate, mops::Union{Vector{<:MeasurementOperator},Nothing}=nothing)

    Chain of quantum circuit gates in a circuit of `N` qubits. Constructed from vector of `AbstractCircuitGate` objects.
    """
    function Circuit{N}(gate::CircuitGate, mops::Vector{<:MeasurementOperator}=MeasurementOperator[]) where {N}
        if isempty(mops)
            return new{N}([Moment(gate)], mops)
        end

        mops_N = maximum(size.(mops))
        mops_N <= N || error("Measurement operators affecting $mops_N wires provided for Circuit of size $N")
        check_commute(mops) || error("Measurement operators do not commute")
        new{N}([Moment(gate)], mops)
    end

    @doc """
    Circuit{N}(moments::Vector{Moment}, mops::Union{Vector{<:MeasurementOperator},Nothing}=nothing) where {N}

    Chain of quantum circuit gates in a circuit of `N` qubits. Constructed from vector of `AbstractCircuitGate` objects.
    """
    function Circuit{N}(moments::Vector{Moment}, mops::Vector{<:MeasurementOperator}=MeasurementOperator[]) where {N}
        if isempty(mops)
            return new{N}(moments, mops)
        end

        mops_N = maximum(size.(mops))
        mops_N <= N || error("Measurement operators affecting $mops_N wires provided for Circuit of size $N")
        check_commute(mops) || error("Measurement operators do not commute")

        new{N}(moments, mops)
    end
end

sparse_matrix(c::Circuit{N}) where {N} = sparse_matrix(c.moments, N)

function add_measurement!(c::Circuit{N}, mops::Vector{<:MeasurementOperator}) where {N}
    if isempty(mops)
        return c
    end
    meas_N = maximum(req_wires.(mops))
    meas_N <= N || error("Measurement operators affecting $meas_N wires provided for circuit of size $N")

    check_commute([c.meas; mops]) || error("Measurement operators do not commute")

    append!(c.meas, mops)
    return c
end

function add_measurement!(c::Circuit{N}, mop::MeasurementOperator) where {N}
    meas_N = maximum(req_wires(mop))
    meas_N <= N || error("Measurement operators affecting $meas_N wires provided for circuit of size $N")
    check_commute([c.meas; mop]) || error("Measurement operators do not commute")

    push!(c.meas, mop)
end

# make Circuit iterable and indexable
function Base.getindex(c::Circuit, i::Integer)
    1 <= i <= length(c.moments) || throw(BoundsError(c, i))
    return c.moments[i]
end

function Base.getindex(c::Circuit, ::Colon)
    return c.moments[:]
end

function Base.getindex(c::Circuit, ur::UnitRange{Int64})
    return c.moments[ur]
end

function Base.iterate(c::Circuit, state=1)
    return state > length(c.moments) ? nothing : (c.moments[state], state + 1)
end

# implement methods required for iteration
function Base.firstindex(::Circuit)
    return 1
end

function Base.lastindex(c::Circuit)
    return length(c.moments)
end

function distribution(ψ::AbstractVector, c::Circuit)
    return apply(ψ, c.moments)
end

function Base.append!(c::Circuit{N}, gate::CircuitGate) where {N}
    cg_N = maximum(gate.iwire)
    cg_N <= N || error("CircuitGate `cg` has maximum iwire `N`, $cg_N. However, Circuit object `c` has size of $N")
    push!(c.moments, Moment(gate))
end

function Base.append!(c::Circuit{N}, gates::Vector{<:CircuitGate}) where {N}
    j = 1
    iwires = falses(N)
    if !isempty(c)
        buffer = pop!(c).gates

        for cg in buffer
            cgwires = collect(cg.iwire)
            iwires[cgwires] .= true
        end
    else
        buffer = CircuitGate[]
    end

    for (i, cg) in enumerate(gates)
        cgwires = collect(cg.iwire)
        all(cgwires .<= N) || error("Unable to add gate with `iwire`: $(cg.iwire), with maximum circuit size `N`: $(N)")
        if any(iwires[cgwires])
            push!(c.moments, Moment(buffer))
            j = i
            iwires[:] .= false
            iwires[cgwires] .= true
            buffer = CircuitGate[cg]
        else
            iwires[cgwires] .= true
            push!(buffer, cg)
        end
    end
    push!(c.moments, Moment(buffer))
end

num_wires(::Circuit{N}) where {N} = N

Base.length(c::Circuit) = length(c.moments)

function Base.reverse(c::Circuit)
    c_new = deepcopy(c)
    reverse!(c_new.moments)
    reverse!.(c_new.moments)
    return c_new
end

function Base.reverse!(c::Circuit)
    reverse!(c.moments)
    reverse!.(c.moments)
    return c
end

function Base.pop!(c::Circuit)
    pop!(c.moments)
end

function Base.isempty(c::Circuit)
    isempty(c.moments)
end


text:
using TaylorIntegration, Test, DiffEqBase, OrdinaryDiffEq
using LinearAlgebra: norm

f(u,p,t) = u
@testset "Test integration of ODE with numbers in common interface" begin
    u0 = 0.5
    tspan = (0.0, 1.0)
    prob = ODEProblem(f, u0, tspan)
    sol = solve(prob, TaylorMethod(50), abstol=1e-20)

    @test sol[end] - 0.5*exp(1) < 1e-12
end

f!(du, u, p, t) = (du .= u)
@testset "Test integration of ODE with abstract arrays in common interface" begin
    u0 = rand(4, 2)
    tspan = (0.0, 1.0)
    prob = ODEProblem(f!, u0, tspan)
    sol = solve(prob, TaylorMethod(50), abstol=1e-20)

    @test norm(sol[end] - u0.*exp(1)) < 1e-12
end

u0 = 1.0
tspan = (0.0,5.0)
saveat_inputs = ([], 0:1:tspan[2], 0:1:(tspan[2]+5), 3:1:tspan[2], collect(0:1:tspan[2]))
@testset "Test saveat behavior with numbers in common interface" begin
    prob = ODEProblem(f, u0, tspan)
    for s ∈ saveat_inputs
        if isempty(s)
            sol = solve(prob, TaylorMethod(20),abstol=1e-20)
        else
            sol = solve(prob, Tsit5(),abstol=1e-20, saveat = s)
        end
        sol_taylor = solve(prob, TaylorMethod(20),abstol=1e-20, saveat = s)

        @test all(sol.t .== sol_taylor.t)
        @test length(sol_taylor.t) == length(sol_taylor.u)
    end
end

u0 = rand(2)
@testset "Test saveat behavior with abstract arrays in common interface" begin
    prob = ODEProblem(f!, u0, tspan)
    for s ∈ saveat_inputs
        if isempty(s)
            sol = solve(prob, TaylorMethod(20),abstol=1e-20)
        else
            sol = solve(prob, Tsit5(),abstol=1e-20, saveat = s)
        end
        sol_taylor = solve(prob, TaylorMethod(20),abstol=1e-20, saveat = s)

        @test all(sol.t .== sol_taylor.t)
        @test length(sol_taylor.t) == length(sol_taylor.u)
    end
end


text:
# version 0.1

module RobHood

using Vega
using Quandl

#Need to download only the new data.
#To do this, we need to store the current stock time series in *.jld files
#Then, for each new download of data of stock named ID, we should:
#-first: load any possible local file named ID.jld
#-then: to compare the time series (x) in the ID.jld file with the one we want to download (we need to calculate the time series before downloading it).
#-So we will only download the data we still don't have locally stored.

function readPortData(portfolio,nS)
    x = [];
    y = [];
    group = [];
    for(i in 1:nS)
       stockid = ASCIIString(portfolio[i])
#       stockid = portfolio[i]
#       println("type: ",typeof(stockid));
       dat = quandl(stockid,rows=365,format="DataFrame");
       println("Stock: $stockid");
       ndat = length(dat[1,])
       x = [x;collect(1:ndat)];
       y = [y;collect(dat[5,])];
       group = [group;[stockid for j in 1:ndat]];
    end
    nT = round(Int,length(y)/nS)

    return nT,x,y,group
end

function getTSPlot(x,y,group)
    myplot=lineplot(x = x, y = y, group = group)
    xlab!(myplot,title="Time (days)")
    ylab!(myplot,title="Closed")
    title!(myplot,title="Stocks")
    myplot.background = "white"
    return myplot
end

function tsplot(x,y,group)
    myplot = getTSPlot(x,y,group);
    myplot
end

###### From now on, there is the implementation of the Efficient Frontier algorithm

function getzbar(x,y,group,nS,nT)
    zbar=zeros(nS);
    for r = 1:nS
        zbar[r] = (y[nT,r] - y[1,r])/3;
    end
    return zbar;
end

# To activate this function we need to change the runtests.jl file in order to add an "epsilon" as a parameter
function removeLowValues!(x,y,group,epsilon)
    posLowValues = find(y.<= epsilon);
    deleteat!(y,);
    deleteat!(group,find(y.<= epsilon));
    deleteat!(x,find(y.<= epsilon));
    nS = length(unique(group));
    nT = round(Int,length(y)/nS)
end

"""
 Function to convert a vector containing stock data into a matrix

 Parameters:
 (x, y, group) =  (vector of time steps, vector of stock data, group: vector stock id's)

---
 Each different time step (x) corresponds to a row in the resulting matrix

 Each different stock id (group) corresponds to a column in the matrix

 The values of the y-vector are disposed in the matrix according to their respective x, and group.

---
 This function returns the number of stocks, the number of time steps, the vectors x, y, and group, and the resulting matrix ymat.
"""
function vecToMat(x,y,group)
#    removeLowValues!(x,y,group,epsilon)
    nS = length(unique(group));
    nT = round(Int,length(y)/nS)
    ymat = zeros(nT,nS);
    for i = 1:nT
        for j = 1:nS
            ymat[i,j] = y[(j-1)*nT + i];
        end
    end
    return nS,nT,x,y,group,ymat;
end

"""
   Function to calculate the Covariance matrix of the time series of stock data

   Parameters:

   -x: matrix of time (each column is a stock, each row is a time step)

   -y: matrix of time series of stock data (each column is a stock, each row is a time step)

   -group: the id of the stock

   -nS: number of stocks

   -nT: number of time steps

---
   This function returns:

   -M: covariance matrix

   -stdevs: standard deviation
"""
function getVarCovMatrix(x,y,group,nS,nT)
    CM=zeros(nS,nS);
    for c = 1:(nS-1)
       for c1 = (c+1):nS
           CM[c,c1] = cov(y[:,c],y[:,c1]);
       end
    end

    for c2 =1:nS
       CM[c2,c2] = var(y[:,c2]);
    end
    M = CM+CM';
    stdevs=sqrt(diag(M));
    return M,stdevs
end

"""
   Function to calculate the Efficient Frontier

   Parameters:

   -zbar:

   -M:

   -ns: Number of stocks

---
   This function returns:

   -mu: expected return

   -minvar: variance

   -minstd: standard deviation
"""
function getEffFrontier(zbar,M,nS)
    unity = ones(length(zbar));
    A = unity' * inv(M) * unity;
    B = unity' * inv(M) * zbar;
    C = zbar' * inv(M) * zbar;
    D = (A .* C) - (B.^2);
    mu = linspace(1,75,nS);

    minvar = ((A .* (mu.^2)) - ((2 .* B .* mu) .+ C)) ./ D;
    minstd = sqrt(minvar);
    return mu,minvar,minstd
end

"""
   Function to build the Vega.jl plotting object for the Efficient Frontier.

   Parameters:

   -stdevs: standard deviation

   -zbar: expected return

   -minstd: standard deviation

   -mu: expected return

---
   This function returns a Vega.jl plotting object
"""
function getEffFrontPlot(stdevs,zbar,mu,minstd,group,xlab_,ylab_)
    eff = layer(scatterplot(x = stdevs, y = zbar, group = unique(group)), lineplot(x = minstd, y = mu) );
#Make names unique in ef line

    xlim!(eff,min=0,max=maximum([stdevs minstd]));
    xlab!(eff,title=xlab_);
    ylab!(eff,title=ylab_);
    title!(eff,title="Efficient frontier Individual securities");
    eff.background = "white";
    return eff;
end

"""
   Function to build the plotting object for the Efficient Frontier.

   Parameters:

   -stdevs: standard deviation

   -zbar: expected return

   -minstd: standard deviation

   -mu: expected return

---
   This function will call getEffFrontPlot in order to build the plotting object
"""
function efffrontplot(stdevs,zbar,mu,minstd,group,xlab_,ylab_)
    myplot = getEffFrontPlot(stdevs,zbar,mu,minstd,group,xlab_,ylab_);
    myplot
end

#Write the functions here

end #module


text:
using Base.Test
using Klara

fnames = (
  :value,
  :loglikelihood,
  :logprior,
  :logtarget,
  :gradloglikelihood,
  :gradlogprior,
  :gradlogtarget,
  :tensorloglikelihood,
  :tensorlogprior,
  :tensorlogtarget,
  :dtensorloglikelihood,
  :dtensorlogprior,
  :dtensorlogtarget
)

println("    Testing ContUnvMarkovChain constructors and methods...")

nstaten = 4
nstate = ContUnvMarkovChain(nstaten)

@test eltype(nstate) == Float64
@test length(nstate.value) == nstaten
for i in 2:13
  @test length(getfield(nstate, fnames[i])) == 0
end
@test size(nstate.diagnosticvalues) == (0, 0)
@test nstate.n == nstaten
@test length(nstate.diagnostickeys) == 0

nstaten = 5
nstate = ContUnvMarkovChain(nstaten, [true; fill(false, 5); true; fill(false, 6)], [:accept], Float32)

@test eltype(nstate) == Float32
@test length(nstate.value) == nstaten
@test length(nstate.gradlogtarget) == nstaten
for i in [2:6; 8:13]
  @test length(getfield(nstate, fnames[i])) == 0
end
@test size(nstate.diagnosticvalues) == (1, nstaten)
@test nstate.n == nstaten
@test length(nstate.diagnostickeys) == 1

statev = Float32(3.)
stateglt = Float32(4.21)
state = BasicContUnvParameterState(statev, [:accept], nothing, nothing, [true])
state.gradlogtarget = stateglt
savei = 2

copy!(nstate, state, savei)
@test nstate.value[savei] == statev
@test nstate.gradlogtarget[savei] == stateglt
nstate.diagnosticvalues[savei] == true

nstaten = 10
nstate = ContUnvMarkovChain(nstaten, [:value, :logtarget], [:accept])

@test eltype(nstate) == Float64
@test length(nstate.value) == nstaten
@test length(nstate.logtarget) == nstaten
for i in [2:3; 5:13]
  @test length(getfield(nstate, fnames[i])) == 0
end
@test size(nstate.diagnosticvalues) == (1, nstaten)
@test nstate.n == nstaten
@test length(nstate.diagnostickeys) == 1

statev = Float64(1.25)
statelt = Float64(-1.12)
state = BasicContUnvParameterState(statev, [:accept], nothing, nothing, [false])
state.logtarget = statelt
savei = 7

copy!(nstate, state, savei)
@test nstate.value[savei] == statev
@test nstate.logtarget[savei] == statelt
nstate.diagnosticvalues[savei] == false

nstatev = Float64[4.39, -9,47]
z = ContUnvMarkovChain(length(nstatev))
z.value = copy(nstatev)
w = deepcopy(z)
@test isequal(z, w)

println("    Testing ContMuvMarkovChain constructors and methods...")

nstatesize = 2
nstaten = 4
nstate = ContMuvMarkovChain(nstatesize, nstaten)

@test eltype(nstate) == Float64
@test size(nstate.value) == (nstatesize, nstaten)
for i in (2, 3, 4)
  @test length(getfield(nstate, fnames[i])) == 0
end
for i in 5:7
  @test size(getfield(nstate, fnames[i])) == (0, 0)
end
for i in 8:10
  @test size(getfield(nstate, fnames[i])) == (0, 0, 0)
end
for i in 11:13
  @test size(getfield(nstate, fnames[i])) == (0, 0, 0, 0)
end
@test size(nstate.diagnosticvalues) == (0, 0)
@test nstate.size == nstatesize
@test nstate.n == nstaten
@test length(nstate.diagnostickeys) == 0

nstatesize = 2
nstaten = 5
nstate = ContMuvMarkovChain(nstatesize, nstaten, [true; fill(false, 3); true; fill(false, 8)], [:accept], Float32)

@test eltype(nstate) == Float32
@test size(nstate.value) == (nstatesize, nstaten)
@test size(nstate.gradloglikelihood) == (nstatesize, nstaten)
for i in (2, 3, 4)
  @test length(getfield(nstate, fnames[i])) == 0
end
for i in (6, 7)
  @test size(getfield(nstate, fnames[i])) == (0, 0)
end
for i in 8:10
  @test size(getfield(nstate, fnames[i])) == (0, 0, 0)
end
for i in 11:13
  @test size(getfield(nstate, fnames[i])) == (0, 0, 0, 0)
end
@test size(nstate.diagnosticvalues) == (1, nstaten)
@test nstate.size == nstatesize
@test nstate.n == nstaten
@test length(nstate.diagnostickeys) == 1

statev = Float32[0.17, 9.44]
stategll = Float32[-0.01, 4.7]
state = BasicContMuvParameterState(statev, [:gradloglikelihood], [:accept], nothing, nothing, [false])
state.gradloglikelihood = stategll
savei = 3

copy!(nstate, state, savei)
@test nstate.value[:, savei] == statev
@test nstate.gradloglikelihood[:, savei] == stategll
nstate.diagnosticvalues[savei] == false

nstatesize = 2
nstaten = 10
nstate = ContMuvMarkovChain(nstatesize, nstaten, [:value, :logtarget, :gradlogtarget], [:accept], Float16)

@test eltype(nstate) == Float16
@test size(nstate.value) == (nstatesize, nstaten)
@test length(nstate.logtarget) == nstaten
@test size(nstate.gradlogtarget) == (nstatesize, nstaten)
for i in (2, 3)
  @test length(getfield(nstate, fnames[i])) == 0
end
for i in (5, 6)
  @test size(getfield(nstate, fnames[i])) == (0, 0)
end
for i in 8:10
  @test size(getfield(nstate, fnames[i])) == (0, 0, 0)
end
for i in 11:13
  @test size(getfield(nstate, fnames[i])) == (0, 0, 0, 0)
end
@test size(nstate.diagnosticvalues) == (1, nstaten)
@test nstate.size == nstatesize
@test nstate.n == nstaten
@test length(nstate.diagnostickeys) == 1

statev = Float16[6.91, 0.42]
statelt = Float16(4.67)
stateglt = Float16[-0.01, 3.2]
state = BasicContMuvParameterState(statev, [:gradlogtarget], [:accept], nothing, nothing, [true])
state.logtarget = statelt
state.gradlogtarget = stateglt
savei = 7

copy!(nstate, state, savei)
@test nstate.value[:, savei] == statev
@test nstate.logtarget[savei] == statelt
@test nstate.gradlogtarget[:, savei] == stateglt
nstate.diagnosticvalues[savei] == true

nstatev = Float64[4.3 9.2 -7.44; -0.2 8.1 4.43]
z = ContMuvMarkovChain(size(nstatev)...)
z.value = copy(nstatev)
w = deepcopy(z)
@test isequal(z, w)


text:
using JLD2

expected = parse(Int, ARGS[1])

for input in ARGS[2:end]
    println(input)
    data = load(input)["results"];
    for noise in keys(data)
        for signal in keys(data[noise])
            shape = size(data[noise][signal])
            if shape[2] < expected
                println("$noise $signal ", shape)
            end
        end
    end
    println()
end


text:
c = InfoTheory.LibzCompressor()

str1 = "a" ^ 100
cstr1 = compress(c, str1)
@test typeof(cstr1) <: AbstractString
@test length(cstr1) < 100

str2 = "b" ^ 100
@test clen(c, str1) < clen(c, str1 * str2)


text:
include("setup.jl")

@testset "Downloads.jl" begin
    @testset "API coverage" begin
        value = "Julia is great!"
        base64 = "SnVsaWEgaXMgZ3JlYXQh"
        url = "$server/base64/$base64"
        headers = ["Foo" => "Bar"]
        # test with one argument
        path = Downloads.download(url)
        @test isfile(path)
        @test value == read(path, String)
        rm(path)
        # with headers
        path = Downloads.download(url, headers=headers)
        @test isfile(path)
        @test value == read(path, String)
        rm(path)
        # test with two arguments
        arg_writers() do path, output
            @arg_test output begin
                @test output == Downloads.download(url, output)
            end
            @test isfile(path)
            @test value == read(path, String)
            rm(path)
            # with headers
            @arg_test output begin
                @test output == Downloads.download(url, output, headers=headers)
            end
            @test isfile(path)
            @test value == read(path, String)
            rm(path)
        end
    end

    @testset "get request" begin
        url = "$server/get"
        data = download_json(url)
        @test "url" in keys(data)
        @test data["url"] == url
    end

    @testset "headers" begin
        url = "$server/headers"

        @testset "set headers" begin
            headers = ["Foo" => "123", "Header" => "VaLuE", "Empty" => ""]
            data = download_json(url, headers)
            @test "headers" in keys(data)
            headers′ = data["headers"]
            for (key, value) in headers
                @test header(headers′, key) == value
            end
            @test header(headers′, "Accept") == "*/*"
        end

        @testset "override default header" begin
            headers = ["Accept" => "application/tar"]
            data = download_json(url, headers)
            @test "headers" in keys(data)
            headers′ = data["headers"]
            @test header(headers′, "Accept") == "application/tar"
        end

        @testset "override default header with empty value" begin
            headers = ["Accept" => ""]
            data = download_json(url, headers)
            @test "headers" in keys(data)
            headers′ = data["headers"]
            @test header(headers′, "Accept") == ""
        end

        @testset "delete default header" begin
            headers = ["Accept" => nothing]
            data = download_json(url, headers)
            @test "headers" in keys(data)
            headers′ = data["headers"]
            @test !("Accept" in keys(headers′))
        end
    end

    @testset "errors" begin
        @test_throws ArgumentError Downloads.download("ba\0d")
        @test_throws ArgumentError Downloads.download("good", "ba\0d")

        err = @exception Downloads.download("xyz://invalid")
        @test err isa ErrorException
        @test startswith(err.msg, "Protocol \"xyz\" not supported")

        err = @exception Downloads.download("https://invalid")
        @test err isa ErrorException
        @test startswith(err.msg, "Could not resolve host")

        err = @exception Downloads.download("$server/status/404")
        @test err isa ErrorException
        @test contains(err.msg, r"^HTTP/\d+(?:\.\d+)?\s+404\b")

        path = tempname()
        @test_throws ErrorException Downloads.download("$server/status/404", path)
        @test !ispath(path)
    end

    @testset "concurrent requests" begin
        delay = 2
        count = 100
        url = "$server/delay/$delay"
        t = @elapsed @sync for id = 1:count
            @async begin
                data = download_json("$url?id=$id")
                @test "args" in keys(data)
                @test get(data["args"], "id", nothing) == ["$id"]
            end
        end
        @test t < 0.9*count*delay
    end

    @testset "referer" begin
        dest = "$server/headers"
        url = "$server/redirect-to?url=$(url_escape(dest))"
        data = download_json(url)
        @test "headers" in keys(data)
        headers′ = data["headers"]
        @test header(headers′, "Referer") == url
    end

    @testset "request API" begin
        @testset "basic request usage" begin
            for status in (200, 300, 400)
                url = "$server/status/$status"
                resp = request_body(multi, url)[1]
                @test resp.url == url
                @test resp.status == status
                test_response_string(resp.response, status)
                @test all(hdr isa Pair{String,String} for hdr in resp.headers)
                headers = Dict(resp.headers)
                @test "content-length" in keys(headers)
            end
        end

        @testset "custom headers" begin
            url = "$server/response-headers?FooBar=VaLuE"
            resp, data = request_body(multi, url)
            @test resp.url == url
            @test resp.status == 200
            test_response_string(resp.response, 200)
            headers = Dict(resp.headers)
            @test "foobar" in keys(headers)
            @test headers["foobar"] == "VaLuE"
        end

        @testset "url for redirect" begin
            dest = "$server/headers"
            url = "$server/redirect-to?url=$(url_escape(dest))"
            resp, data = request_json(multi, url)
            @test resp.url == dest
            @test resp.status == 200
            test_response_string(resp.response, 200)
            @test "headers" in keys(data)
            headers′ = data["headers"]
            @test header(headers′, "Referer") == url
        end

        @testset "progress" begin
            # https://httpbingo.org/drip doesn't work
            # see https://github.com/mccutchen/go-httpbin/issues/40
            # fixed, but their deployed setup is still broken
            url = "https://httpbin.org/drip"
            progress = Downloads.Curl.Progress[]
            req = Request(devnull, url, String[])
            Downloads.request(req, multi, p -> push!(progress, p))
            unique!(progress)
            @test 11 ≤ length(progress) ≤ 12
            shift = length(progress) - 10
            @test all(p.dl_total == (i==1 ? 0 : 10) for (i, p) in enumerate(progress))
            @test all(p.dl_now   == max(0, i-shift) for (i, p) in enumerate(progress))
        end
    end
end


text:
module handle_validvalues

export parse_validvalues, get_datatype, value_is_valid

using CategoricalArrays
using DataAPI
using Dates

"Returns: The type of each valid value (after validvalues has been parsed)."
get_datatype(validvalues::DataType) = validvalues
get_datatype(validvalues::Set)      = eltype(validvalues)
get_datatype(validvalues::T) where {T <: AbstractRange} = typeof(validvalues[1])

"Returns: True if value is in validvalues"
value_is_valid(value::T, validvalues::DataType) where {T <: CategoricalValue} = typeof(DataAPI.unwrap(value)) == validvalues
value_is_valid(value::T, validvalues::DataType) where {T <: CategoricalValue{String, <:Integer}} = typeof(DataAPI.unwrap(value)) == validvalues
value_is_valid(value::T, validvalues)           where {T <: CategoricalValue} = value_is_valid(DataAPI.unwrap(value), validvalues)
value_is_valid(value::T, validvalues)           where {T <: CategoricalValue{String, <:Integer}} = value_is_valid(DataAPI.unwrap(value), validvalues)
value_is_valid(value, validvalues::DataType) = typeof(value) == validvalues
value_is_valid(value, validvalues) = typeof(value) == eltype(validvalues) && in(value, validvalues)

"""
Returns: An instance of ColumnSchema.validvalues.

validvalues could be a:
- DataType. E.g., "Int64"
- Range of a Core type. E.g., "1:10"
- Range of a non-Core type, represented with a tuple: "(val1, stepsize, val2)"
"""
function parse_validvalues(parser::Function, returntype::DataType, validvalues::String)
    validvalues[1] == '(' && return parse_noncore_range(parser, returntype, validvalues)  # Range of a non-Core type

    # DataType
    try
        eval(Meta.parse(validvalues)) == returntype && return returntype  # "Int" becomes Int64
    catch e
    end

    # Range of a Core type
    vals = strip.(String.(split(validvalues, ':')))
    val1 = parser(vals[1])
    val2 = parser(vals[2])
    length(vals) == 2 && return val1:val2       # Step-size not supplied
    val3 = parser(vals[3])
    length(vals) == 3 && return val1:val2:val3  # Step-size supplied

    # Else error
    error("Cannot parse validvalues: $(validvalues)")
end

function parse_validvalues(parser::Function, returntype::DataType, validvalues::T) where {T <: Vector}
    eltype(validvalues) == returntype ? validvalues : parser.(validvalues)
end

"""
Parse String to a range of values with non-Core type.

The format is: "(start, stepsize, stop)"

Example: (2017-10-01 09:00+10:00, Day(1), 2017-12-08 23:00+10:00), where the entries have type TimeZones.ZonedDateTime.
"""
function parse_noncore_range(parser::Function, returntype::DataType, validvalues::String)
    @assert length(validvalues) >= 5  # "(a,b)" contains 5 characters
    @assert validvalues[1] == '('
    @assert validvalues[end] == ')'
    validvalues = strip.(String.(split(validvalues[2:(end-1)], ",")))
    length(validvalues) != 3 && error("Range of non-Base type requires 3 elements. It has $(length(validvalues)).")
    val1 = parser(validvalues[1])
    val2 = returntype <: Dates.TimeType ? eval(Meta.parse(validvalues[2])) : parser(validvalues[2])
    val3 = parser(validvalues[3])
    val1:val2:val3
end

end


text:
using Documenter, ColorSchemes, Luxor

makedocs(
    modules = [ColorSchemes],
    sitename = "ColorSchemes",
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        assets = ["assets/colorschemes-docs.css"]),
    pages = Any[
        "Introduction"      => "index.md",
        "Catalogue"         => "catalogue.md",
        "Basic usage"       => "basics.md",
        "Good practice"     => "goodpractice.md",
        "Finding colors"    => "finding.md",
        "Plotting"          => "plotting.md",
        "Images"            => "images.md",
        "Index"             => "functionindex.md"
    ]
)

deploydocs(
    push_preview = true,
    repo = "github.com/JuliaGraphics/ColorSchemes.jl.git",
    target = "build"
)

max_stars_repo_name: cormullion/ColorSchemes

max_issues_repo_name: cormullion/ColorSchemes

max_forks_repo_name: cormullion/ColorSchemes


text:
# Julia LIIC & other useful scripts
# Written on 26 November 2019
# Update on 28 March 2022
# Execute by running `julia liic.jl --flags path_to_geom1 path_to_geom2`

using ArgParse

include("utils.jl")
import .Utils

# include("interpol.jl")
# include("kabsch.jl")
# include("converter.jl")
# include("structs.jl")
# using .Interpol, .Kabsch, .Converter, ..Structs

parse_settings = ArgParseSettings()

@add_arg_table parse_settings begin
    "--cartesian", "-c"
        help = "Specifies cartesian coordinate interpolation"
        action = :store_true
    "--distance", "-d"
        help = "Specifies distance coordinate interpolation"
        action = :store_true
    "--internal", "-i"
        help = "Specifies interatomic distance interpolation"
        action = :store_true
    "--steps", "-s"
        help = "Number of steps between structures in interpolation"
        arg_type = Int
        default = 20
    "--output", "-o"
        help = "Path to output .xyz file of the LIIC"
        arg_type = String
        default = "liic_output.xyz"
    "--kabsch", "-k"
        help = "Whether kabsch should run before interpolation"
        action = :store_true
    "geom1"
        help = "path to .xyz file containing first geometry"
        required = true
    "geom2"
        help = "path to .xyz file containing first geometry"
        required = true
end

function main()
    parsed_args = parse_args(ARGS, parse_settings)
    stp = parsed_args["steps"]
    file = parsed_args["output"]
    in_file_1 = parsed_args["geom1"]
    in_file_2 = parsed_args["geom2"]
    MoleculeA = Utils.import_molecule(in_file_1)
    MoleculeB = Utils.import_molecule(in_file_2)
    @assert MoleculeA.Atoms == MoleculeB.Atoms
    if parsed_args["cartesian"]
        println("Performing cartesian interpolation...")
        # in_mat_1, atomNames1 = Interpol.xyz2matrix(in_file_1)
        # in_mat_2, atomNames2 = Interpol.xyz2matrix(in_file_2)
        # Perform kabsch algorithm on geom2 if kabsch==True
        if parsed_args["kabsch"]
            println("Running Kabsch algorithm on second geometry prior to LIIC")
            MoleculeA,MoleculeB = Utils.kabsch_rotate(MoleculeA,MoleculeB)
        end
        outputArray = Utils.cartesian(MoleculeA,MoleculeB,stp)
    elseif parsed_args["distance"]
        println("Performing interpolation in internal distance matrix...")
    elseif parsed_args["internal"]
        println("Performing interpolation in internal coordinates...")
        ZMatA = Utils.xyz_to_zmat(MoleculeA)
        ZMatB = Utils.xyz_to_zmat(MoleculeB)
        @assert ZMatA.VarNames == ZMatB.VarNames
        outputArray = Utils.internal(ZMatA,ZMatB,stp)
		# internalCoords1,header = Utils.xyz2internal(in_file_1)
		# internalCoords2,header = Utils.xyz2internal(in_file_2)
		# @assert internalCoords1[:,1] == internalCoords2[:,1]
		# outputArray, atomNames1 = Interpol.internal_babel(internalCoords1,internalCoords2,stp,header)
    else
        println("ERROR")
    end

    Utils.write_liic(outputArray,file)

end

# Off we ... off we ... off we fucking GO!
main()

max_stars_repo_name: dasha-shchep/LIIC

max_issues_repo_name: dasha-shchep/LIIC

max_forks_repo_name: dasha-shchep/LIIC


text:
using Documenter, AeroAcoustics

makedocs(
    modules = [AeroAcoustics],
    checkdocs = :none,
    sitename = "AeroAcoustics.jl",
    pages = Any["index.md"],
    repo = "https://gitlab.windenergy.dtu.dk/ollyl/AeroAcoustics.jl/blob/{commit}{path}#{line}"
)


text:
"""
"""
struct DiscreteModelPortion{Dc,Dp} <: DiscreteModel{Dc,Dc}
  model::DiscreteModel{Dc,Dp}
  parent_model::DiscreteModel{Dc,Dp}
  d_to_dface_to_parent_dface::Vector{Vector{Int}}
end

get_grid(model::DiscreteModelPortion) = get_grid(model.model)

get_grid_topology(model::DiscreteModelPortion) = get_grid_topology(model.model)

get_face_labeling(model::DiscreteModelPortion) = get_face_labeling(model.model)

get_face_to_parent_face(model::DiscreteModelPortion,d::Integer) = model.d_to_dface_to_parent_dface[d+1]

get_cell_to_parent_cell(model::DiscreteModelPortion) = get_face_to_parent_face(model,num_cell_dims(model))

get_parent_model(model::DiscreteModelPortion) = model.parent_model

"""
"""
function DiscreteModelPortion(model::DiscreteModel, cell_to_parent_cell::AbstractVector{<:Integer})
  topo = get_grid_topology(model)
  labels = get_face_labeling(model)
  grid_p =  GridPortion(get_grid(model),cell_to_parent_cell)
  topo_p, d_to_dface_to_parent_dface = _grid_topology_portion(topo,cell_to_parent_cell)
  labels_p = _setup_labels_p(labels,d_to_dface_to_parent_dface)
  model_p = DiscreteModel(grid_p,topo_p,labels_p)
  DiscreteModelPortion(model_p,model,d_to_dface_to_parent_dface)
end

function DiscreteModelPortion(model::DiscreteModel, cell_to_is_in::AbstractArray{Bool})
  cell_to_parent_cell = findall(collect1d(cell_to_is_in))
  DiscreteModelPortion(model,cell_to_parent_cell)
end

function DiscreteModelPortion(model::DiscreteModel, cell_to_is_in::AbstractVector{Bool})
  cell_to_parent_cell = findall(cell_to_is_in)
  DiscreteModelPortion(model,cell_to_parent_cell)
end

function _grid_topology_portion(topo,cell_to_parent_cell)
  D = num_cell_dims(topo)

  n_m_to_parent_nface_to_parent_mface = [
    Table(get_faces(topo,n,m)) for n in 0:D, m in 0:D ]

  d_to_dface_to_parent_dface = [
    _setup_dface_to_parent_dface(
      n_m_to_parent_nface_to_parent_mface[D+1,d+1],
      num_faces(topo,d),
      cell_to_parent_cell) for d in 0:D]

  d_to_dface_to_vertices  = [
    _setup_connectivities_d(
    n_m_to_parent_nface_to_parent_mface[d+1,0+1],
    d_to_dface_to_parent_dface[d+1],
    d_to_dface_to_parent_dface[0+1],
    num_faces(topo,0)) for d in 0:D]

  D = num_cell_dims(topo)
  vertex_coordinates = get_vertex_coordinates(topo)[d_to_dface_to_parent_dface[0+1]]
  cell_type = get_cell_type(topo)[d_to_dface_to_parent_dface[D+1]]
  polytopes = get_polytopes(topo)
  orientation_style = OrientationStyle(topo)

  topo_p = UnstructuredGridTopology(
    vertex_coordinates,
    d_to_dface_to_vertices,
    cell_type,
    polytopes,
    orientation_style)

  topo_p, d_to_dface_to_parent_dface
end

function _setup_dface_to_parent_dface(
  parent_cell_to_parent_dface::Table,
  num_parent_dfaces,
  cell_to_parent_cell)

  parent_dface_touched = fill(false,num_parent_dfaces)
  dface = 0
  for parent_cell in cell_to_parent_cell
    pini = parent_cell_to_parent_dface.ptrs[parent_cell]
    pend = parent_cell_to_parent_dface.ptrs[parent_cell+1]-1
    for p in pini:pend
      parent_dface = parent_cell_to_parent_dface.data[p]
      parent_dface_touched[parent_dface] = true
    end
  end
  dface_to_parent_dface = findall(parent_dface_touched)
  dface_to_parent_dface
end

function _setup_connectivities_d(
  parent_nface_parent_mface::Table,
  nface_to_parent_nface,
  mface_to_parent_mface,
  num_parent_mfaces)

  parent_mface_to_mface = fill(-1,num_parent_mfaces)
  parent_mface_to_mface[mface_to_parent_mface] .= 1:length(mface_to_parent_mface)

  nface_to_mface = parent_nface_parent_mface[nface_to_parent_nface]
  for p in 1:length(nface_to_mface.data)
    parent_mface = nface_to_mface.data[p]
    mface = parent_mface_to_mface[parent_mface]
    @check mface > 0
    nface_to_mface.data[p] = mface
  end
  nface_to_mface
end

function _setup_labels_p(labels,d_to_dface_to_parent_dface)

  D = length(d_to_dface_to_parent_dface)-1

  d_to_dface_to_entity = [
    zeros(
      eltype(labels.d_to_dface_to_entity[d+1]),
      length(d_to_dface_to_parent_dface[d+1]))
    for  d in 0:D]

  for d in 0:D
    face_to_label = d_to_dface_to_entity[d+1]
    oface_to_label = labels.d_to_dface_to_entity[d+1]
    face_to_oface = d_to_dface_to_parent_dface[d+1]
    _update_labels_dim!(face_to_label,oface_to_label,face_to_oface)
  end

  FaceLabeling(d_to_dface_to_entity,labels.tag_to_entities,labels.tag_to_name)
end

function _update_labels_dim!(face_to_label,oface_to_label,face_to_oface)
  for face in 1:length(face_to_label)
    oface = face_to_oface[face]
    label = oface_to_label[oface]
    face_to_label[face] = label
  end
end


text:
# GOKU-net model
#
# Based on
# https://arxiv.org/abs/2003.10775

struct GOKU <: LatentDE end

@doc raw"""
    apply_feature_extractor(encoder::Encoder{GOKU}, x)
    apply_feature_extractor(encoder::Encoder{LatentODE}, x)

Applies the feature extractor layer contained in the `encoder` to the batch of input data x, usually reducing its dimensionality (i.e. extracting features).

# Arguments
`encoder`: Encoder structure containing all the encoder layers.\
`x`: Input data. When using the default architecture, `x` correspond to an array of size `pixels` x `batch size` x `time`.
"""
apply_feature_extractor(encoder::Encoder{GOKU}, x) = encoder.feature_extractor(x)

@doc raw"""
    apply_pattern_extractor(encoder::Encoder{GOKU}, fe_out)

Passes `fe_out` through the pattern_extractor layer contained in the `encoder`, returning a tuple associated to the initial latent states and parameters, respectively.

# Arguments
`encoder`: Encoder structure containing all the encoder layers.\
`fe_out`: Output of feature extractor layer.
"""
function apply_pattern_extractor(encoder::Encoder{GOKU}, fe_out)
    pe_z₀, pe_θ_forward, pe_θ_backward = encoder.pattern_extractor
    fe_out = Flux.unstack(fe_out, 3)

    # reverse sequence
    fe_out_rev = reverse(fe_out)

    # pass it through the recurrent layers
    pe_z₀_out = [pe_z₀(x) for x in fe_out_rev][end]
    pe_θ_out_f = [pe_θ_forward(x) for x in fe_out][end]
    pe_θ_out_b = [pe_θ_backward(x) for x in fe_out_rev][end]
    pe_θ_out = vcat(pe_θ_out_f, pe_θ_out_b)

    # reset hidden states
    Flux.reset!(pe_z₀)
    Flux.reset!(pe_θ_forward)
    Flux.reset!(pe_θ_backward)

    return pe_z₀_out, pe_θ_out
end

@doc raw"""
    apply_latent_in(encoder::Encoder{GOKU}, pe_out)
    apply_latent_in(encoder::Encoder{LatentODE}, pe_out)

Applies the `encoder`'s `latent_in` layer to `pe_out`, returning the mean and log-variance of the latent variables to use for sampling.

# Arguments
`encoder`: Encoder structure containing all the encoder layers.\
`pe_out`: Output of pattern extractor layer.
"""
function apply_latent_in(encoder::Encoder{GOKU}, pe_out)
    pe_z₀_out, pe_θ_out = pe_out
    li_μ_z₀, li_logσ²_z₀, li_μ_θ, li_logσ²_θ = encoder.latent_in

    z₀_μ = li_μ_z₀(pe_z₀_out)
    z₀_logσ² = li_logσ²_z₀(pe_z₀_out)

    θ_μ = li_μ_θ(pe_θ_out)
    θ_logσ² = li_logσ²_θ(pe_θ_out)

    return (z₀_μ, θ_μ), (z₀_logσ², θ_logσ²)
end

@doc raw"""
    apply_latent_out(decoder::Decoder{GOKU}, l̃)

Applies the `decoder`'s `latent_out` layer to `l̃`, returning a tuple with initial conditions and parameters for use in the differential equation layer.

# Arguments
`decoder`: Decoder structure containing all the decoder layers.\
`l̃`: Tuple containing sampled abstract representations of the initial conditions and parameters, respectively.
"""
function apply_latent_out(decoder::Decoder{GOKU}, l̃)
    z̃₀, θ̃ = l̃
    lo_z₀, lo_θ = decoder.latent_out

    ẑ₀ = lo_z₀(z̃₀)
    θ̂ = lo_θ(θ̃)

    return ẑ₀, θ̂
end

@doc raw"""
    diffeq_layer(decoder::Decoder{GOKU}, l̂, t)

Solves the differential equations contained in the `diffeq` layer of the `decoder` using the initial conditions and parameters contained in `l̂`, and saving at times `t`.
"""
function diffeq_layer(decoder::Decoder{GOKU}, l̂, t)
    ẑ₀_, θ̂_ = l̂

    # make sure the diff eq  solving is done on cpu
    ẑ₀ = cpu(ẑ₀_)
    θ̂ = cpu(θ̂_)

    prob = decoder.diffeq.prob
    solver = decoder.diffeq.solver
    sensealg = decoder.diffeq.sensealg
    kwargs = decoder.diffeq.kwargs

    # Function definition for ensemble problem
    prob_func(prob,i,repeat) = remake(prob, u0=ẑ₀[:,i], p = θ̂[:,i])

    # Check if solve was successful, if not, return NaNs to avoid problems with dimensions matches
    output_func(sol, i) = sol.retcode == :Success ? (Array(sol), false) : (fill(NaN32,(size(ẑ₀, 1), length(t))), false)

    ## Adapt problem to given time span and create ensemble problem definition
    prob = remake(prob; tspan = (t[1],t[end]))
    ens_prob = EnsembleProblem(prob, prob_func = prob_func, output_func = output_func)

    ## Solve
    ẑ = solve(ens_prob, solver, EnsembleThreads(); sensealg = sensealg, trajectories = size(θ̂, 2), saveat = t, kwargs...)

    # Optionally transform the latent state variables
    ẑ = transform_after_diffeq(Array(ẑ), decoder.diffeq)
    ẑ = permutedims(ẑ, [1,3,2])

    # go back to gpu if it corresponds
    ẑ = device(ẑ₀_, ẑ)
    return ẑ
end

device(x::Flux.CUDA.CuArray, y) = gpu(y)
device(x::Array, y) = y

# Identity by default
transform_after_diffeq(x, diffeq) = x

@doc raw"""
    apply_reconstructor(decoder::Decoder{GOKU}, ẑ)
    apply_reconstructor(decoder::Decoder{LatentODE}, ẑ)

Passes latent trajectories `ẑ` through the reconstructor layer contained in the `decoder`.

# Arguments
`decoder`: Decoder structure containing all the decoder layers.\
`ẑ`: Latent trajectories, consists of matrices corresponding to different time frames and having size `latent data dimension` x `batch size`.
"""
apply_reconstructor(decoder::Decoder{GOKU}, ẑ) = decoder.reconstructor(ẑ)

@doc raw"""
    sample(μ::T, logσ²::T, model::LatentDiffEqModel{LatentODE}) where T <: Array

Samples latent variables from the normal distribution with mean μ and variance exp(logσ²).
"""
function sample(μ::T, logσ²::T, model::LatentDiffEqModel{GOKU}) where T <: Tuple{Array, Array}
    z₀_μ, θ_μ = μ
    z₀_logσ², θ_logσ² = logσ²

    ẑ₀ = z₀_μ + randn(Float32, size(z₀_logσ²)) .* exp.(z₀_logσ²/2f0)
    θ̂ =  θ_μ + randn(Float32, size( θ_logσ²)) .* exp.(θ_logσ²/2f0)

    return ẑ₀, θ̂
end

function sample(μ::T, logσ²::T, model::LatentDiffEqModel{GOKU}) where T <: Tuple{Flux.CUDA.CuArray, Flux.CUDA.CuArray}
    z₀_μ, θ_μ = μ
    z₀_logσ², θ_logσ² = logσ²

    ẑ₀ = z₀_μ + gpu(randn(Float32, size(z₀_logσ²))) .* exp.(z₀_logσ²/2f0)
    θ̂ =  θ_μ + gpu(randn(Float32, size( θ_logσ²))) .* exp.(θ_logσ²/2f0)

    return ẑ₀, θ̂
end

@doc raw"""
    default_layers(model_type, input_dim::Int, diffeq;
        device = cpu,
        hidden_dim_resnet = 200, rnn_input_dim = 32,
        rnn_output_dim = 16, latent_dim = 16,
        latent_to_diffeq_dim = 200, θ_activation = softplus,
        output_activation = σ, init = Flux.kaiming_uniform(gain = 1/sqrt(3)))

Generates default encoder and decoder layers that are to be fed into the LatentDiffEqModel.

# Arguments
`model_type`: GOKU() or LatentODE()\
`input_dim`: Dimension of input\
`diffeq`: Differential equations structure, containing fields `prob`, `solver` and `sensealg`, which correspond to DifferentialEquations.jl's problem, solver and sensitivity algorithm, respectively.

# Keyword Arguments
`device`: Flux.jl's `cpu` or `gpu`.
`hidden_dim_resnet`: Hidden dimension of the feature_extractor's resnet .
`rnn_input_dim`: Input dimension of the pattern_extractor layers.
`latent_dim`: Diemension of the latent variables.
`latent_to_diffeq_dim`: Hidden dimension of the dense layers from latent_out.
`θ_activation`: Activation function used in the last dense layer from latent_out corresponding to the parameters of the differential equation `θ`, which can be useful for imposing contrains.
`output_activation`: Activation function used in the last layer of the reconstructor.
"""
function default_layers(model_type::GOKU, input_dim, diffeq; device=cpu,
                            hidden_dim_resnet = 200, rnn_input_dim = 32,
                            rnn_output_dim = 16, latent_dim = 16,
                            latent_to_diffeq_dim = 200, θ_activation = softplus,
                            output_activation = σ, init = Flux.kaiming_uniform(gain = 1/sqrt(3)))

    z_dim = length(diffeq.prob.u0)
    θ_dim = length(diffeq.prob.p)

    ######################
    ### Encoder layers ###
    ######################
    # Resnet
    l1 = Dense(input_dim, hidden_dim_resnet, relu, init = init)
    l2 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)
    l3 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)
    l4 = Dense(hidden_dim_resnet, rnn_input_dim, relu, init = init)
    feature_extractor = Chain(l1,
                                SkipConnection(l2, +),
                                SkipConnection(l3, +),
                                l4) |> device

    # RNN
    pe_z₀ = Chain(RNN(rnn_input_dim, rnn_output_dim, relu, init = init),
                       RNN(rnn_output_dim, rnn_output_dim, relu, init = init)) |> device

    # Bidirectional LSTM
    pe_θ_forward = Chain(LSTM(rnn_input_dim, rnn_output_dim, init = init),
                       LSTM(rnn_output_dim, rnn_output_dim, init = init)) |> device

    pe_θ_backward = Chain(LSTM(rnn_input_dim, rnn_output_dim, init = init),
                        LSTM(rnn_output_dim, rnn_output_dim, init = init)) |> device

    pattern_extractor = (pe_z₀, pe_θ_forward, pe_θ_backward)

    # final fully connected layers before sampling
    li_μ_z₀ = Dense(rnn_output_dim, latent_dim, init = init) |> device
    li_logσ²_z₀ = Dense(rnn_output_dim, latent_dim, init = init) |> device

    li_μ_θ = Dense(rnn_output_dim*2, latent_dim, init = init) |> device
    li_logσ²_θ = Dense(rnn_output_dim*2, latent_dim, init = init) |> device

    latent_in = (li_μ_z₀, li_logσ²_z₀, li_μ_θ, li_logσ²_θ)

    encoder_layers = (feature_extractor, pattern_extractor, latent_in)

    ######################
    ### Decoder layers ###
    ######################

    # after sampling in the latent space but before the differential equation layer
    lo_z₀ = Chain(Dense(latent_dim, latent_to_diffeq_dim, relu, init = init),
                        Dense(latent_to_diffeq_dim, z_dim, init = init)) |> device

    lo_θ = Chain(Dense(latent_dim, latent_to_diffeq_dim, relu, init = init),
                        Dense(latent_to_diffeq_dim, θ_dim, θ_activation, init = init)) |> device

    latent_out = (lo_z₀, lo_θ)

    # going back to the input space
    # Resnet
    l1 = Dense(z_dim, hidden_dim_resnet, relu, init = init)
    l2 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)
    l3 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)
    l4 = Dense(hidden_dim_resnet, input_dim, output_activation, init = init)
    reconstructor = Chain(l1,
                            SkipConnection(l2, +),
                            SkipConnection(l3, +),
                            l4)  |> device

    decoder_layers = (latent_out, diffeq, reconstructor)

    return encoder_layers, decoder_layers
end

max_issues_repo_name: gabrevaya/LatentDiffEq

max_forks_repo_name: gabrevaya/LatentDiffEq


text:
basepath = abspath("./results/test1_830-900am")
f = open("$basepath/BNlog.txt", "w")
write(f, "Hello world")

max_stars_repo_name: zhangzhangzf/IntentionPrediction

max_issues_repo_name: zhangzhangzf/IntentionPrediction

max_forks_repo_name: zhangzhangzf/IntentionPrediction


text:
using Distances

using mjWrap
using MuJoCo

################################################## Experiment Centric Functions
## accepts model and trajectory data
## returns sum of trajectory rewards for each traj; saves per timestep reward
function rewardfunction(x::mjWrap.mjSet,
                        s0::AbstractVector{Float64}, # pass in states
                        s1::AbstractVector{Float64},
                        o0::AbstractVector{Float64}, # observations
                        o1::AbstractVector{Float64},
                        ctrl::AbstractVector{Float64})::Float64
   tid = Threads.threadid()
   d = x.datas[tid]

   height = s1[3]
   t_height = 1.25
   t_speed = 0.8

   reward = 0.0

   #reward -= 1.0*abs(s1[1]) #(s1[1]^2)         #(10*s1[1])^2
   #reward -= 1.0*abs(s1[2]) #(s1[2]^2)         #(10*s1[2])^2
   if height < t_height
      reward -= 2.0*abs(t_height - s1[3])
   #else
      #reward -= abs( s1[x.nq+1] )
      #reward -= abs( s1[x.nq+2] )
      #reward -= 1.0*abs(s1[1]) #(s1[1]^2)         #(10*s1[1])^2
      #reward -= 1.0*abs(s1[2]) #(s1[2]^2)         #(10*s1[2])^2
   #else
   #   reward += 1.0 - abs(t_speed-s1[x.nq+1])
   end

   return reward
end

function obsfunction(x::mjWrap.mjSet,
                     s::AbstractVector{Float64},
                     o::AbstractVector{Float64})
end

## accepts model,
## default state,
## preallocated matrix of states (nqnv x K)
function initfunction!(x::mjWrap.mjSet,
                       s::mjWrap.TrajSamples)
   mag = 0.2
   init_state = view( s.state, :, 1, : )
   for t=1:s.K
      #@. init_state[:,t] = s.s0 #+ rand() * mag - mag/2.0
      #init_state[3,t] -= 0.1
      #init_state .= 0.0
      #init_state[1:x.nq,t] = LAYING_QPOS
      init_state[:,t] = s.s0 #+ rand() * mag - mag/2.0
   end
end

function ctrlfunction!(x::mjWrap.mjSet, ctrl::Array{Float64})
   randn!(ctrl)
end

# modelfunc should takes in time index (iteration index) to know change curric
function modelfunction!(x::mjWrap.mjSet, iter::Int=0)
end

## evaluation function: reward function shapes exploration, evaluation
## function defines success or failure
function evalfunction(x::mjWrap.mjSet, d::mjWrap.TrajSamples)
   #return mean(d.state[1,end,:]) # x position at time T, for all K
   return mean(@view( d.state[3,end,:] )) # z position at time T, for all K
end

using ExpFunctions
myfuncs = FunctionSet(modelfunction!,
                      initfunction!,
                      ctrlfunction!,
                      obsfunction,
                      rewardfunction,
                      evalfunction)


text:
using Convex, SCS
set_default_solver(SCSSolver(verbose=0, max_iters=1000000))

println("small test to compile and see if we throw warnings when warmstarting what can't be warmstarted.")
n = 10
y = rand(n)
x = Variable(n)

lambda = [100]
problem = minimize(sumsquares(y - x) + lambda[1] * sumsquares(x - 10))
@time solve!(problem, warmstart = true)

println("\nnow try first solving and then warmstarting")
n = 1000
y = rand(n)
x = Variable(n)

lambda = [100]
problem = minimize(sumsquares(y - x) + lambda[1] * sumsquares(x - 10))
@time solve!(problem)
@show problem.optval
lambda[1] = 105
println("this run should be faster than the last and have slightly higher optimal value")
@time solve!(problem, warmstart=true)
@show problem.optval

println("\nnow try with scalar variables")
lambda = 100
problem = minimize(sumsquares(y - x) + lambda * sumsquares(x - 10))
@time solve!(problem)
@show problem.optval
lambda = 105
println("this run should be faster than the last and have slightly higher optimal value")
@time solve!(problem, warmstart=true)
@show problem.optval


text:
using PyCall
using DelimitedFiles

function loadcsv(inputcsv)
    contents = readdlm(inputcsv, ',', Float64, '\n'; header = true)
    return contents
end

function coldata(inputcsv, theset)
    contents = loadcsv(inputcsv)
    return contents[1][:, theset]
end

# Example
inputcsv = "LE1NAFEMS_MSH8_convergence.CSV"

@pyimport matplotlib.pyplot as plt
plt.style[:use]("seaborn-whitegrid")
fig = plt.figure()
ax = plt.axes()
x = coldata(inputcsv, 1)
y = coldata(inputcsv, 2)
ax[:loglog](x, abs.(y), color = "blue", marker = "x")
y = coldata(inputcsv, 3)
ax[:loglog](x, abs.(y), color = "red", marker = "o")
ax[:set_xlabel]("Number of elements")
ax[:set_ylabel]("Approximate error")
plt.show()


text:
for (fname, elty) in (
    (Symbol("cublasDgemm_v2"), :Float64),
    (Symbol("cublasSgemm_v2"), :Float32),
    (Symbol("cublasZgemm_v2"), :ComplexF64),
    (Symbol("cublasCgemm_v2"), :ComplexF32),
)
    @eval begin
        function _gemm!(
            transA::Char,
            transB::Char,
            alpha::Number,
            A::StridedCuMatrix{$elty},
            B::StridedCuMatrix{$elty},
            beta::Number,
            C::StridedCuMatrix{$elty},
        )
            m = size(A, transA == 'N' ? 1 : 2)
            k = size(A, transA == 'N' ? 2 : 1)
            n = size(B, transB == 'N' ? 2 : 1)
            if m != size(C, 1) || n != size(C, 2) || k != size(B, transB == 'N' ? 1 : 2)
                throw(DimensionMismatch(""))
            end
            if stride(A, 1) != 1 || stride(B, 1) != 1 || stride(C, 1) != 1
                throw(DimensionMismatch("cuBLAS requires matrices to be in column-major order"))
            end
            lda = max(1, stride(A, 2))
            ldb = max(1, stride(B, 2))
            ldc = max(1, stride(C, 2))
            CUDA.CUBLAS.$fname(
                CUDA.CUBLAS.handle(),
                transA,
                transB,
                m,
                n,
                k,
                alpha,
                A,
                lda,
                B,
                ldb,
                beta,
                C,
                ldc,
            )
            C
        end

    end
end
for elty in (:Float32, :Float64, :ComplexF32, :ComplexF64)
    @eval begin
        function _mul_with_strides!(
            C::StridedCuMatrix{$elty},
            A::StridedCuMatrix{$elty},
            B::StridedCuMatrix{$elty},
            a::$elty,
            b::$elty,
        )
            if C isa CuArray
                mul!(C, A, B, a, b)
            else
                _gemm!('N', 'N', a, A, B, b, C)
            end
        end
        _mul_with_strides!(C::StridedMatrix{$elty}, A, B, a::$elty, b::$elty) =
            mul!(C, A, B, a, b)
    end
end

@doc raw"""
    ThreeBlockMatrix

!!! warning
    This is an internal data structure

A dense matrix with top-left and bottom-right blocks assumed to be zero:
![three-block-structure](three-blocks.jpg)

`ThreeBlockMatrix` stores blocks 1, 2, and 3 as dense matrices. There is also a special
overload of [`LinearAlgebra.mul!`](@ref) function for faster matrix-matrix multiplication.
"""
struct ThreeBlockMatrix{T, M <: AbstractMatrix{T}}
    block₁::M
    block₂::M
    block₃::M
end

Base.ndims(::ThreeBlockMatrix) = 2

Base.eltype(::Type{ThreeBlockMatrix{T, M}}) where {T, M} = T
# Base.eltype(::Type{ThreeBlockMatrix{T}}) where {T} = T
Base.eltype(::Type{ThreeBlockMatrix}) = Any

Base.size(x::ThreeBlockMatrix) = (size(x, 1), size(x, 2))
function Base.size(x::ThreeBlockMatrix, dim::Integer)
    dim > 0 || throw(DimensionMismatch("dimension out of range: $dim"))
    dim <= 2 ? size(x.block₁, 1) + size(x.block₂, 1) : 1
end

Base.real(x::ThreeBlockMatrix) =
    ThreeBlockMatrix(real(x.block₁), real(x.block₂), real(x.block₃))
Base.imag(x::ThreeBlockMatrix) =
    ThreeBlockMatrix(imag(x.block₁), imag(x.block₂), imag(x.block₃))

Base.show(io::IO, x::ThreeBlockMatrix) = print(io, "ThreeBlockMatrix")

function Base.convert(::Type{M₁}, x::ThreeBlockMatrix{T, M₂}) where {M₁, T, M₂ <: M₁}
    N = size(x.block₁, 1) + size(x.block₂, 1)
    n₁ = size(x.block₂, 1)
    n₂ = size(x.block₃, 2)
    out = convert(M₁, similar(x.block₁, N, N))
    fill!(view(out, 1:n₁, 1:n₁), zero(eltype(out)))
    fill!(view(out, (N - n₂ + 1):N, (N - n₂ + 1):N), zero(eltype(out)))
    copy!(view(out, (n₁ + 1):N, 1:(N - n₂)), x.block₁)
    copy!(view(out, 1:n₁, (n₁ + 1):(N - n₂)), x.block₂)
    copy!(view(out, 1:(N - n₂), (N - n₂ + 1):N), x.block₃)
    return out
end

Adapt.adapt_structure(to, x::ThreeBlockMatrix) =
    ThreeBlockMatrix(adapt(to, x.block₁), adapt(to, x.block₂), adapt(to, x.block₃))

@doc raw"""
    sparsity(x::ThreeBlockMatrix)

Get the fraction of zeros which are optimized away during GEMM.
"""
function sparsity(x::ThreeBlockMatrix)
    N = size(x.block₁, 1) + size(x.block₂, 1)
    n₁ = size(x.block₂, 1)
    n₂ = size(x.block₃, 2)
    (n₁^2 + n₂^2) / N^2
end

@doc raw"""
    ThreeBlockMatrix(G, n₁, n₂)

!!! warning
    This is an internal function

Given a square dense matrix `G` and sizes `n₁` and `n₂` of top-left and bottom-right blocks
respectively, construct a `ThreeBlockMatrix`. All three blocks are copied (this ensure
contiguous storage for faster matrix-matrix products).
"""
function ThreeBlockMatrix(G::AbstractMatrix, n₁::Integer, n₂::Integer)
    if size(G, 1) != size(G, 2)
        throw(DimensionMismatch("'G' must be a square matrix, but got a matrix of shape $(size(G))"))
    end
    n₁ >= 0 || throw(ArgumentError("Invalid 'n₁': $n₁"))
    n₂ >= 0 || throw(ArgumentError("Invalid 'n₂': $n₂"))
    N = size(G, 1)
    if N == 0
        throw(DimensionMismatch("Empty matrices are not (yet) supported: 'G' has shape $(size(G))"))
    end
    if (n₁ + n₂) >= N
        throw(ErrorException("Overlapping zero regions are not (yet) supported: $n₁ + $n₂ >= $N"))
    end
    ThreeBlockMatrix(
        G[(n₁ + 1):N, 1:(N - n₂)],
        G[1:n₁, (n₁ + 1):(N - n₂)],
        G[1:(N - n₂), (N - n₂ + 1):N],
    )
end

@doc raw"""
    ThreeBlockMatrix(G)

!!! warning
    This is an internal function

Analyze matrix `G` and construct three-block version of it for faster matrix-matrix
multiplication. This function computes `n₁` and `n₂` and constructs `ThreeBlockMatrix`
using [`ThreeBlockMatrix(G, n₁, n₂)`](@ref).
"""
function ThreeBlockMatrix(G::AbstractMatrix)
    if size(G, 1) != size(G, 2)
        throw(DimensionMismatch("Expected a square matrix, but G has size $(size(G))"))
    end
    N = size(G, 1)
    N > 1 || return ThreeBlockMatrix(G, 0, 0)
    # All elements smaller than ε·‖G‖ are assumed to be negligible and are set to zero.
    cutoff = mapreduce(abs, max, G) * eps(real(eltype(G)))
    n₁ = _analyze_top_left(G, cutoff)
    n₁ < N - 1 || (n₁ = N - 2)
    n₂ = _analyze_bottom_right(G, cutoff)
    n₂ < N - 1 || (n₂ = N - 2)
    if n₁ + n₂ >= N
        @warn "Temperature is so low that G contains overlapping zero regions. " *
              "Reducing n₁ or n₂. This will hurt performance..."
        # Overlapping zero regions are not yet supported... We reduce one of the blocks to
        # ensure that n₁ + n₂ == size(G, 1) - 1
        if n₁ >= n₂
            n₂ = N - n₁ - 1
        else
            n₁ = N - n₂ - 1
        end
    end
    ThreeBlockMatrix(G, n₁, n₂)
end
# function ThreeBlockMatrix(G::CuArray)
#     @warning "There are no CUDA kernels for _analyze_top_left and _analyze_bottom_right. " *
#              "Matrix 'G' will be temporarily copied to CPU..."
# end

@doc raw"""
    _analyze_top_left(G::AbstractMatrix{ℂ}, cutoff::ℝ) -> Int

!!! warning
    This is an internal function!

We assume that `G` is a square matrix with the following block structure:
![top-left-block](top-left-block.jpg)
I.e. the top-left corner of size `n₁ x n₁` consists of zeros. This function determines `n₁`
by treating all matrix elements smaller than `cutoff` as zeros.
"""
function _analyze_top_left(
    G::AbstractMatrix{ℂ},
    cutoff::ℝ,
)::Int where {ℝ <: Real, ℂ <: Union{ℝ, Complex{ℝ}}}
    # Find an upper bound on n. We follow the first column of G and find the first element
    # which exceeds cutoff.
    n = size(G, 1)
    @inbounds for i in 1:size(G, 1)
        if abs(G[i, 1]) >= cutoff
            n = i - 1
            break
        end
    end
    # Fine-tune it by considering other columns.
    j = 2
    @inbounds while j <= n
        i = 1
        while i <= n
            if abs(G[i, j]) >= cutoff
                n = max(i, j) - 1
                break
            end
            i += 1
        end
        j += 1
    end
    return n
end

@doc raw"""
    _analyze_bottom_right(G::AbstractMatrix{ℂ}, cutoff::ℝ) -> Int

!!! warning
    This is an internal function!

Very similar to [`_analyze_top_left`](@ref) except that now the bottom right corner of G is
assumed to contain zeros. This function determines the size of this block.
"""
function _analyze_bottom_right(
    G::AbstractMatrix{ℂ},
    cutoff::ℝ,
)::Int where {ℝ <: Real, ℂ <: Union{ℝ, Complex{ℝ}}}
    # Find a lower bound on n
    n = 1
    @inbounds for i in size(G, 1):-1:1
        if abs(G[i, size(G, 2)]) >= cutoff
            n = i + 1
            break
        end
    end
    # Fine-tune it
    j = size(G, 2) - 1
    @inbounds while j >= n
        i = size(G, 1)
        while i >= n
            if abs(G[i, j]) >= cutoff
                n = min(i, j) + 1
                break
            end
            i -= 1
        end
        j -= 1
    end
    return size(G, 1) + 1 - n
end

"""
    mul!(C, A, B::_ThreeBlockMatrix, α, β) -> C

Overload of matrix-matrix multiplication for ThreeBlockMatrix. We can spare some flops
because of the zero blocks.

We replace one bigg GEMM by three smaller:

  1) ![first-gemm](gemm-1.jpg)
  2) ![second-gemm](gemm-2.jpg)
  3) ![third-gemm](gemm-3.jpg)
"""
function LinearAlgebra.mul!(C::AbstractMatrix, A::AbstractMatrix, B::ThreeBlockMatrix, α, β)
    N = size(B, 1)
    if size(A, 2) != N
        throw(DimensionMismatch("Dimensions of 'A' and 'B' are incompatible: $(size(A)) vs $(size(B))"))
    end
    if size(C) != size(A)
        throw(DimensionMismatch("'C' has wrong dimension: $(size(C)), expected $(size(A))"))
    end
    n₁ = size(B.block₂, 1)
    n₂ = size(B.block₃, 2)
    _mul_with_strides!(
        view(C, :, 1:size(B.block₁, 2)),
        view(A, :, (n₁ + 1):N),
        B.block₁,
        α,
        β,
    )
    if n₁ > 0
        # NOTE: If n₁ is zero, then A B α + C β amount to C β, and since β is 1 we can drop
        # it altogerher.
        #! format: off
        _mul_with_strides!(view(C, :, (n₁ + 1):size(B.block₁, 2)), view(A, :, 1:n₁), B.block₂, α, one(eltype(C)))
        #! format: on
    end
    if n₂ > 0
        #! format: off
        _mul_with_strides!(view(C, :, (size(B.block₁, 2) + 1):size(C, 2)), view(A, :, 1:size(B.block₁, 2)), B.block₃, α, β)
        #! format: on
    end
    return C
end

dot_batched(A, B) = dot_batched!(similar(A, size(A, 1)), A, B)
function dot_batched!(
    out::StridedVector{T},
    A::StridedMatrix{T},
    B::StridedMatrix{T},
) where {T}
    if size(A) != size(B)
        throw(DimensionMismatch("Dimensions of 'A' and 'B' do not match: $(size(A)) != $(size(B))"))
    end
    if length(out) != size(A, 1)
        throw(DimensionMismatch("'out' has wrong length: $(length(out)); expected $(size(A, 1))"))
    end
    (n, m) = size(A)
    fill!(out, zero(T))
    @inbounds for j in 1:m
        @simd for i in 1:n
            out[i] += conj(A[i, j]) * B[i, j]
        end
    end
    out
end

function dot_kernel_batched!(
    batch_size,
    n,
    out::CuDeviceMatrix{T},
    _A::CuDeviceMatrix{T},
    _B::CuDeviceMatrix{T},
) where {T}
    A = CUDA.Const(_A)
    B = CUDA.Const(_B)
    block_acc = @cuDynamicSharedMem(T, (blockDim().x, blockDim().y))
    # Initialize block_acc for each block within the "first grid"
    if blockIdx().x <= gridDim().x && blockIdx().y <= gridDim().y
        @inbounds block_acc[threadIdx().x, threadIdx().y] = zero(T)
    end

    # Perform thread-local accumulation in each block
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    j = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    if i <= batch_size
        thread_acc = zero(T)
        @inbounds while j <= n
            thread_acc += conj(A[i, j]) * B[i, j]
            j += blockDim().y * gridDim().y
        end
        @inbounds block_acc[threadIdx().x, threadIdx().y] += thread_acc
    end
    sync_threads()

    # Reduction within each block
    stride = blockDim().y
    while stride > 1
        stride >>= 1
        if threadIdx().y <= stride
            @inbounds block_acc[threadIdx().x, threadIdx().y] +=
                block_acc[threadIdx().x, threadIdx().y + stride]
        end
        sync_threads()
    end
    if i <= batch_size && threadIdx().y == 1
        @inbounds out[i, blockIdx().y] = block_acc[threadIdx().x, 1]
    end
    nothing
end
function dot_batched!(
    out::StridedCuVector{T},
    A::StridedCuMatrix{T},
    B::StridedCuMatrix{T},
) where {T}
    if size(A) != size(B)
        throw(DimensionMismatch("Dimensions of 'A' and 'B' do not match: $(size(A)) != $(size(B))"))
    end
    if length(out) != size(A, 1)
        throw(DimensionMismatch("'out' has wrong length: $(length(out)); expected $(size(A, 1))"))
    end
    (n, m) = size(A)
    function num_threads(threads)
        threads_y = 32
        threads_x = div(threads, threads_y)
        threads_x, threads_y
    end
    num_blocks(threads) = cld.(size(A), threads)
    amount_shmem(threads) = prod(threads) * sizeof(T)
    @assert stride(A, 1) == 1 && stride(B, 1) == 1
    _A = unsafe_wrap(CuArray{T, 2}, Base.unsafe_convert(CuPtr{T}, A), (stride(A, 2), size(A, 2)))
    _B = unsafe_wrap(CuArray{T, 2}, Base.unsafe_convert(CuPtr{T}, B), (stride(B, 2), size(B, 2)))
    kernel = @cuda launch = false dot_kernel_batched!(n, m, _A, _A, _B)
    config = launch_configuration(kernel.fun, shmem = amount_shmem)
    threads = num_threads(config.threads)
    blocks = num_blocks(threads)
    shmem = amount_shmem(threads)
    temp = similar(A, n, blocks[2])
    GC.@preserve A B kernel(
        n,
        m,
        temp,
        _A,
        _B;
        threads = threads,
        blocks = blocks,
        shmem = shmem,
    )
    sum!(out, temp)
end

function _eigen!(A::Hermitian{T, CuArray{T, 2, B}}) where {T, B}
    eigenvalues, eigenvectors = T <: Complex ? CUSOLVER.heevd!('V', A.uplo, A.data) :
        CUSOLVER.syevd!('V', A.uplo, A.data)
    return eigenvalues, eigenvectors
end
function _eigen!(A::Hermitian{T, Array{T, 2}}) where {T}
    factorization = eigen!(A)
    return factorization.values, factorization.vectors
end


text:
@testset "CSV loading" begin
    files = [refDataPath * "/refBatchDfCodes.csv", refDataPath * "/refParallelDfCodes.csv"]

    rows, cols = getCSVSize(files[1], header = true)
    @test rows == 10
    @test cols == 10

    data1 = loadCSV(files[1], header = true)
    data2 = loadCSV(files[2], header = true)
    @test typeof(data1) == Matrix{Float64}
    @test size(data1) == (10, 10)

    sizes = loadCSVSizes(files, header = true)
    @test sizes == [10, 10]

    W = addprocs(3)
    @everywhere using GigaSOM
    di = loadCSVSet(:csvTest, files, W, header = true)
    @test gather_array(di) == vcat(data1, data2)
    sizes = dmapreduce(di, size, vcat)
    dims = map(last, sizes)
    counts = map(first, sizes)
    @test all(dims .== 10)
    @test minimum(counts) + 1 >= maximum(counts)
    rmprocs(W)
end


text:
# setting working directory
# cd("/Volumes/SSD Hans/Github/MacroFall2020/AdvMacro")
# mkpath("Assignment2")
# cd("/Volumes/SSD Hans/Github/MacroFall2020/AdvMacro/Assignment2")
# mkpath("graphs")

# Pkg.add("Parameters")
using Plots
using Parameters
# Pkg.add("Roots")

#Pkg.add("IntervalArithmetic")
#Pkg.add("IntervalRootFinding")
using Roots
using IntervalArithmetic, IntervalRootFinding

# defining parameters

# Paramters
    # Generate structure for parameters using Parameters module
    # We can set default values for our parameters
    @with_kw struct Par
        # Model Parameters
        z::Float64 = 1    ; # Productivity
        α::Float64 = 1/3  ; # Production function
        β::Float64 = 0.98 ; # Discount factor
        σ::Float64 = 2 ;
        η::Float64 = 1 ;
        δ::Float64 = 1 ;
        l::Float64 = 0.4 ;
        # VFI Paramters
        max_iter::Int64   = 2000  ; # Maximum number of iterations
        dist_tol::Float64 = 1E-9  ; # Tolerance for distance
        # Howard's Policy Iterations
        H_tol::Float64    = 1E-9  ; # Tolerance for policy function iteration
    end

p = Par()

@unpack z, α, β, σ, η, l, δ = p
global M = ((z*α*β)/(1-β+β*δ))^(1/(1-α))
global N = z*M^α-δ*M

# Getting the right χ
function get_me_chi(p::Par)
    @unpack z, α, β, σ, η, l = p
    #f(x) = x*(N^σ)*(l^(σ+η))-(z*(1-α)*M^α)
    chi = (z*(1-α)*M^α)/(l^(σ+η)*N^σ)
    return chi
end

global χ = get_me_chi(p)

# function get_me_l(k,kp,p::Par)
#     @unpack z, α, β, M, N, σ, η, χ = p
#     f(l) = z*(1-α)*(k^α)/((z*(k^α)*(l^(1-α))-kp)^σ)-χ*(l^(η+α))
#     A = range(0,1;length=200)
#     B = f.(A)
#     index = findmin(B[B.>=0])[2]
#     l = A[index]
#     return l
# end
#
# get_me_l(k_grid_20[1],k_grid_20[G_kp_20][1],p)

# Function to find optimal labor
function get_l(k,kp,p::Par)
    A = range(0,1;length=100)
    B = [utility(k,kp,A[i],p) for i in 1:100]
    index = findmax(B)[2]
    l = A[index]
    return l
end

# Steady state values (funciton)
function SS_values(p::Par)
    # This function takes in parameters and provides steady state values
    # Parameters: productivity (z), returns to scale (a) and discount factor (b)
    # Output: values for capital, production, consumption, rental rate, wage
    @unpack z, α, β, l = p
    l_ss = ((z*(1-α)*M^α)/(χ*N^σ))^(1/(σ+η))
    k_ss = M*l_ss
    c_ss = N*l_ss
    y_ss = z*l_ss*M^α
    r_ss = α*z*M^(α-1)
    w_ss = (1-α)*z*M^α
    return l_ss,k_ss,y_ss,c_ss,r_ss,w_ss
end

# Test steady state function
l_ss,k_ss,y_ss,c_ss,r_ss,w_ss = SS_values(p)
println(" ")
println("------------------------")
println(" Steady State variables")
println("   Quantities: l = $l_ss; k = $k_ss; y = $y_ss; c = $c_ss;")
println("   Prices:     r = $r_ss; w = $w_ss;")
println("------------------------")
println(" ")

# Steady state values (funciton)
function path(k,kp,l,p::Par)
    # This function takes in parameters and provides steady state values
    # Parameters: productivity (z), returns to scale (a) and discount factor (b)
    # Output: values for capital, production, consumption, rental rate, wage
    @unpack z, α, β, M, N, l = p
    w = (a,b) -> (1-α)*(a^α)*(b^(-α))
    r = (a,b) -> α*(a^(a-1))*(b^(1-α))
    y = (a,b) -> (a^α)*b^(1-α)
    c = (a,b,d) -> y(a,b)-d
    l = get_l.(k,k_p,fill(p,length(k)))
    wage = w.(k,l)
    int = r.(k,l)
    cons = c.(k,kp,l)
    output = y.(k,l)
    labor = l
    return wage, int, cons, labor, output
end

# Utility function
function utility(k,kp,l,p::Par)
    @unpack z, α, η, σ = p
    y = (a,b) -> z*(a^α)*b^(1-α)
    c = (a,b,d) -> y(a,b)-d
    lb = (l) -> (χ*l.^(η+1)/(η+1))
    utils = c(k,l,kp).^(1-σ)/(1-σ)-lb(l)
    if c(k,l,kp)>0
    return utils
    else
    return -Inf
    end
end

# VFI with Grid
    # Grid
    function Make_K_Grid(n_k,p::Par)
        # Get SS
        l_ss,k_ss,y_ss,c_ss,r_ss,w_ss = SS_values(p)
        # Get k_grid
        k_grid = range(1E-5,2*k_ss;length=n_k) ; # Equally spaced grid between 0 and 2*k_ss
        # Return
        return k_grid
    end

    # Solve VFI with grid search and loops
    function VFI_grid(T::Function,k_grid,p::Par)
        # VFI paramters
        @unpack max_iter, dist_tol = p
        # Initialize variables for loop
        n_k    = length(k_grid) ; # Number of grid nodes
        V_old  = zeros(n_k)     ; # Initial value, a vector of zeros
        V_dist = 1              ; # Initialize distance
        println(" ")
        println("------------------------")
        println("VFI - Grid Search - n_k=$n_k")
        for iter=1:max_iter
            # Update value function
            V_new, G_kp, G_l = T(V_old)
            # Update distance and iterations
            V_dist = maximum(abs.(V_new./V_old.-1))
            iter  += 1
            # Update old function
            V_old  = V_new
            # Report progress
            if mod(iter,100)==0
                println("   VFI Loop: iter=$iter, dist=",100*V_dist,"%")
            end
            # Check convergence and return results
            if V_dist<=dist_tol
                println("VFI - Grid Search - n_k=$n_k")
                println("Iterations = $iter and Distance = ",100*V_dist,"%")
                println("------------------------")
                println(" ")
                return V_new, G_kp, G_l
            end
        end
        # If loop ends there was no convergence -> Error!
        error("Error in VFI - Grid Search - Solution not found")
    end

#-----------------------------------------------------------
#-----------------------------------------------------------
# VFI - Grid Search - Loops

# Define function for Value update and policy functions
    function T_grid_loop(V_old,k_grid,p::Par)
        @unpack z, α, β = p
        n_k  = length(k_grid)
        V    = zeros(n_k)
        G_kp = fill(0,n_k)
        G_l = fill(0.0,n_k)
        for i = 1:n_k
            V_aux = zeros(n_k) ;
            l_aux = zeros(n_k) # Empty vector for auxiliary value of V(i,j)
            for j = 1:n_k
                # Evaluate potential value function for combinations of
                # current capital k_i and future capital k_j
                l_aux[j] = get_l(k_grid[i],k_grid[j],p)
                V_aux[j] = utility(k_grid[i],k_grid[j],l_aux[j],p) + β*V_old[j]
                #println(V_aux[j]," ",k_grid[i]," ",k_grid[j]," ",utility(k_grid[i],k_grid[j],z,a,b) + b*V_old[j])
            end
            # Choose maximum value given current capital k_i
            V[i], G_kp[i] = findmax(V_aux)
            G_l[i] = l_aux[G_kp[i]]
        end
        return V, G_kp, G_l
    end

# Solve VFI with grid search and loops
    function Solve_VFI_loop(n_k,p::Par)
        println("----------------------------------------------------")
        println("Value Function Iteration - Grid Search with Loops")
        # Get Grid
        k_grid = Make_K_Grid(n_k,p)
        # Solve VFI
        V, G_kp, G_l = VFI_grid(x->T_grid_loop(x,k_grid,p),k_grid,p)
        # Return Solution
        return V,G_kp, G_l, k_grid
    end

    # Euler Equation
    function Euler_Error(k,kp,kpp,l,lp,p::Par)
        # Return percentage error in Euler equation
        @unpack z, α, β, σ = p
        LHS = 1/(z*k^α*l^(1-α)-kp)^σ
        RHS = β*α*z*kp^(α-1)*lp^(1-α)/(z*kp^α*lp^(1-α)-kpp)^σ
        return (RHS/LHS-1)*100
    end

    # Euler error by number of points in grid
    function Euler_grid(k_grid::AbstractArray,G_kp::AbstractArray)
        # Euler error of numerical policy function on grid
        n_k = length(k_grid)
        Euler = zeros(n_k)
        for i=1:n_k
            k   = k_grid[i]
            kp  = k_grid[G_kp[i]]
            kpp = k_grid[G_kp[G_kp[i]]]
            l = get_l(k,kp,p)
            lp = get_l(kp,kpp,p)
            Euler[i] = Euler_Error(k,kp,kpp,l,lp,p)
        end
        # Return
        return Euler
    end

#-----------------------------------------------------------
    # VFI - Grid Search - Matrices - Howard's Policy Iteration

    # Define function for Value update and policy functions
        function HT_grid_mat(V_old,U_mat,k_grid,p::Par,n_H)
            @unpack z, α, β, H_tol = p
            # Get Policy Function
            n_k    = length(V_old)
            V,G_kp = findmax( U_mat .+ β*repeat(V_old',n_k,1) , dims=2 )
            V_old  = V
            # "Optimal" U for Howard's iteration
                U_vec = U_mat[G_kp]
            # Howard's policy iteration
            # G_kp is a Cartesian Index
            for i=1:n_H
                V = U_vec .+ β*repeat(V_old',n_k,1)[G_kp]
                if maximum(abs.(V./V_old.-1))<=H_tol
                    break
                end
                V_old = V
            end
            # Recover Policy Functions
            G_kp   = [G_kp[i][2] for i in 1:n_k] # G_kp is a Cartesian index
            G_l = get_l.(k_grid,k_grid[G_kp],fill(p,n_k))

            # Return output
            return V, G_kp, G_l
        end

    # Solve VFI with Howard's policy iteration
        function Solve_VFI_HPI(n_H,n_k,p::Par)
            println("----------------------------------------------------")
            println("Value Function Iteration - Howard's Policy Iteration")
            # Get Grid
            k_grid = Make_K_Grid(n_k,p)
            # Utility matrix
            U_mat = [utility(k_grid[i],k_grid[j],get_l(k_grid[i],k_grid[j],p),p) for i in 1:n_k, j in 1:n_k]
            # Solve VFI
            V, G_kp, G_l = VFI_grid(x->HT_grid_mat(x,U_mat,k_grid,p,n_H),k_grid,p)
            # Return Solution
            return V,G_kp, G_l, k_grid
        end

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

# Fixed point with MPB
    function VFI_grid_MPB(T::Function,k_grid,p::Par)
        @unpack β, max_iter, dist_tol = p
        # Initialize variables for loop
        n_k    = length(k_grid) ; # Number of grid nodes
        V_old  = zeros(n_k)     ; # Initial value, a vector of zeros
        iter   = 0              ; # Iteration index
        V_dist = 1              ; # Initialize distance
        println(" ")
        println("------------------------")
        println("VFI - Grid Search - MPB - n_k=$n_k")
        for iter=1:max_iter
            # Update value function
            V_new, G_kp, G_l = T(V_old)
            # MPB and Distance
            MPB_l  = β/(1-β)*minimum(V_new-V_old)
            MPB_h  = β/(1-β)*maximum(V_new-V_old)
            V_dist = MPB_h - MPB_l
            # Update old function
            V_old  = V_new
            # Report progress
            if mod(iter,100)==0
                println("   VFI Loop: iter=$iter, dist=",100*V_dist,"%")
            end
            # Check Convergence
            if (V_dist<=dist_tol)
                # Recover value and policy functions
                V = V_old .+ (MPB_l+MPB_h)/2
                # Return
                println("VFI - Grid Search - MPB - n_k=$n_k")
                println("Iterations = $iter and Distance = ",100*V_dist)
                println("------------------------")
                println(" ")
                return V, G_kp, G_l
            end
        end
        # Report error for non-convergence
        error("Error in VFI - Grid Search - MPB - Solution not found")
    end

    # Define function for Value update and policy functions
        function T_grid_mat(V_old,U_mat,k_grid,p::Par)
            @unpack z, α, β = p
            n_k    = length(V_old)
            V,G_kp = findmax( U_mat .+ β*repeat(V_old',n_k,1) , dims=2 )
            G_kp   = [G_kp[i][2] for i in 1:n_k] # G_kp is a Cartesian index
            G_l = get_l.(k_grid,k_grid[G_kp],fill(p,n_k))
            return V, G_kp, G_l
        end

# Solve VFI with MPB
    function Solve_VFI_MPB(n_k,p::Par)
        println("----------------------------------------------------")
        println("Value Function Iteration - MacQueen-Porteus Bounds")
        # Get Grid
        k_grid = Make_K_Grid(n_k,p)
        # Utility matrix
        U_mat = [utility(k_grid[i],k_grid[j],get_l(k_grid[i],k_grid[j],p),p) for i in 1:n_k, j in 1:n_k]
        # Solve VFI
        V, G_kp, G_l = VFI_grid_MPB(x->T_grid_mat(x,U_mat,k_grid,p),k_grid,p)
        # Return Solution
        return V,G_kp, G_l, k_grid
    end

max_stars_repo_name: py-r-hans/AdvMacro

max_issues_repo_name: py-r-hans/AdvMacro

max_forks_repo_name: py-r-hans/AdvMacro