text:
# This file contains functions for the Revenue Efficiency Reverse DDF DEA model
"""
    RevenueReverseDDFDEAModel
An data structure representing a cost RDDF DEA model.
"""
struct RevenueReverseDDFDEAModel <: AbstractRevenueDEAModel
    n::Int64
    m::Int64
    s::Int64
    measure::Symbol
    rts::Symbol
    monetary::Bool
    dmunames::Union{Vector{String},Nothing}
    eff::Vector
    lambda::SparseMatrixCSC{Float64, Int64}
    techeff::Vector
    alloceff::Vector
    normalization::Vector
    Xtarget::Matrix
    Ytarget::Matrix
    Gxrddf::Union{Vector,Matrix}
    Gyrddf::Union{Vector,Matrix}
end

"""
    dearevenuerddf(X, Y, P, measure)
Compute profit efficiency using data envelopment analysis Reverse DDF model for
inputs `X`, outputs `Y`, price of outputs `P`, and efficiency `measure`.

# Measure specification:
- `:ERG`: Enhanced Russell Graph Slack Based Measure.

# Optional Arguments
- `rts=:VRS`: choose between constant returns to scale `:CRS` or variable returns to scale `:VRS`.
- `atol=1e-6`: tolerance for DMU to be considered efficient.
- `names`: a vector of strings with the names of the decision making units.

# Examples
```jldoctest
julia> X = [1; 1; 1; 1; 1; 1; 1; 1];

julia> Y = [7 7; 4 8; 8 4; 3 5; 3 3; 8 2; 6 4; 1.5 5];

julia> P = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1];

julia> dearevenuerddf(X, Y, P, :ERG)
Revenue Reverse DDF DEA Model
DMUs = 8; Inputs = 1; Outputs = 2
Returns to Scale: VRS
Associated efficiency measure = ERG
────────────────────────────────────
    Revenue  Technical    Allocative
────────────────────────────────────
1  0.0        0.0        0.0
2  0.25       0.0        0.25
3  0.25       0.0        0.25
4  0.464286   0.464286   0.0
5  0.571429   0.571429   0.0
6  0.666667   0.333333   0.333333
7  0.314286   0.314286  -5.55112e-17
8  0.818182   0.672727   0.145455
────────────────────────────────────
```
"""
function dearevenuerddf(X::Union{Matrix,Vector}, Y::Union{Matrix,Vector},
    P::Union{Matrix,Vector},
    measure::Symbol;
    rts::Symbol = :VRS, monetary::Bool = false, atol::Float64 = 1e-6,
    names::Union{Vector{String},Nothing} = nothing,
    optimizer::Union{DEAOptimizer,Nothing} = nothing)::RevenueReverseDDFDEAModel

    # Check parameters
    nx, m = size(X, 1), size(X, 2)
    ny, s = size(Y, 1), size(Y, 2)

    np, sp = size(P, 1), size(P, 2)

    if nx != ny
        throw(DimensionMismatch("number of rows in X and Y ($nx, $ny) are not equal"));
    end
    if np != ny
        throw(DimensionMismatch("number of rows in P and Y ($np, $ny) are not equal"));
    end
    if sp != s
        throw(DimensionMismatch("number of columns in P and Y ($sp, $s) are not equal"));
    end

    n = nx

    # Get directions based on efficiency measure
    Gxrddf, Gyrddf = DataEnvelopmentAnalysis.dearddfdirections(X, Y, measure, orient = :Output, rts = rts, optimizer = optimizer)

    # Get normalization of efficiency measure model
    if measure == :ERG
        measurerevenuemodel = dearevenuerussell(X, Y, P, rts = rts, optimizer = optimizer)
    end

    normmeasurerev = normfactor(measurerevenuemodel)

    # Cost model using calculated directions
    rddfrevenuemodel = dearevenueddf(X, Y, P, Gy = Gyrddf, rts = rts, optimizer = optimizer)

    refficiency = efficiency(rddfrevenuemodel)
    rlambdaeff = peersmatrix(rddfrevenuemodel)
    techefficiency = efficiency(rddfrevenuemodel, :Technical)
    allocefficiency = efficiency(rddfrevenuemodel, :Allocative)
    normalization = normfactor(rddfrevenuemodel)
    Xtarget = targets(rddfrevenuemodel, :X)
    Ytarget = targets(rddfrevenuemodel, :Y)

    # Change normalization factor of efficient DMUs
    for i=1:n
        if abs(techefficiency[i] ) <= atol
            refficiency[i] = refficiency[i] ./ normmeasurerev[i] .* normalization[i]
            normalization[i] = normmeasurerev[i]
        end
    end

    if monetary
        refficiency = refficiency .* normalization
        techefficiency = techefficiency .* normalization
    end

    allocefficiency = refficiency - techefficiency

    return RevenueReverseDDFDEAModel(n, m, s, measure, rts, monetary, names, refficiency, rlambdaeff, techefficiency, allocefficiency, normalization, Xtarget, Ytarget, Gxrddf, Gyrddf)

end

function Base.show(io::IO, x::RevenueReverseDDFDEAModel)
    compact = get(io, :compact, false)

    n = nobs(x)
    m = ninputs(x)
    s = noutputs(x)
    dmunames = names(x)

    eff = efficiency(x)
    techeff = efficiency(x, :Technical)
    alloceff = efficiency(x, :Allocative)

    if !compact
        print(io, "Revenue Reverse DDF DEA Model \n")
        print(io, "DMUs = ", n)
        print(io, "; Inputs = ", m)
        print(io, "; Outputs = ", s)
        print(io, "\n")
        print(io, "Returns to Scale: ", string(x.rts))
        print(io, "\n")
        print(io, "Associated efficiency measure = ", string(x.measure))
        print(io, "\n")
        show(io, CoefTable(hcat(eff, techeff, alloceff), ["Revenue", "Technical", "Allocative"], dmunames))
    end

end


text:
export readMesh, save, load, read_data, write_data, convert_mat, read_strain_stress
function readMesh(gmshFile::String)
    fp = open(gmshFile);
    boundaries = Dict{String, Array}()
    physicalnames = Dict{Int64, String}()
    nodes = nothing
    elements = nothing

    line = readline(fp)
    if (line == "\$MeshFormat")
        format = readline(fp)
        @assert(readline(fp) == "\$EndMeshFormat")
    end

    line = readline(fp)
    if (line == "\$PhysicalNames")
        nphys = parse(Int64,readline(fp))
        for i = 1:nphys
            line = readline(fp)
            physicalid = parse(Int64, split(line)[2])
            physicalnames[physicalid] = split(line)[3]
            boundaries[split(line)[3]] = []
        end
        @assert(readline(fp) == "\$EndPhysicalNames")
    end

    line = readline(fp)
    if (line == "\$Nodes")
        nnodes = parse(Int64, readline(fp))
        nodes = zeros(nnodes,2)
        for i = 1:nnodes
            l = readline(fp)|>split
            nodes[i,1] = parse(Float64, l[2])
            nodes[i,2] = parse(Float64, l[3])
        end
        @assert(readline(fp) == "\$EndNodes")
    end
    # println(physicalnames)
    # error()
    line = readline(fp)
    if (line == "\$Elements")
        nelems = readline(fp)
        nelems = parse(Int64,nelems)
        elements = []
        for i = 1:nelems
            l = readline(fp)|>split
            l4 = parse(Int64, l[4])
            physicalname = physicalnames[l4]
            if startswith(physicalname,"\"Dirichlet")
                k1,k2 = parse(Int64, l[6]),parse(Int64, l[7])
                push!(boundaries[physicalname],[k1;k2])
            elseif startswith(physicalname,"\"Neumann")
                k1,k2 = parse(Int64, l[6]),parse(Int64, l[7])
                push!(boundaries[physicalname],[k1;k2])
            else
                # println(physicalnames[l[4]], startswith(physicalnames[l[4]],"Dirichlet"))
                # println(l)
                k = [parse(Int64, l[i]) for i = 6:9]
                push!(elements, k)
            end
        end
        # elem = zeros(length(elements),4)
        # for i = 1:length(elements)
        #     elem[i,:] = elements[i]
        # end
        @assert(readline(fp) == "\$EndElements")
    end
    @warn("Jobs are not finished!!! Users need to compute EBC, g, NBC, f.")
    return elements, nodes, boundaries
end

function save(file::String, domain::Domain, globaldata::GlobalData)
    @save file domain globaldata
end

function load(file::String)
    @load file domain globaldata
    return domain, globaldata
end

# domain.state, domain.fint, fext
function write_data(file::String, domain::Domain)
    file = matopen(file, "w")
    write(file, "state", domain.history["state"])
    write(file, "fext", domain.history["fext"])
    write(file, "fint", domain.history["fint"])
    write(file, "strain", domain.history["strain"])
    write(file, "stress", domain.history["stress"])
    close(file)
end

# state, fext
function read_data(file::String)
    vars = matread(file)
    #use fint, strain, stress, for debugging purpose
    vars["state"], vars["fext"]
end

# state, fext
function read_strain_stress(file::String)
    vars = matread(file)
    #use fint, strain, stress, for debugging purpose
    vars["strain"], vars["stress"]
end

# type == nn2array: read mat convert to 1D array
# type == array2nn: read 1D array convert to mat
function convert_mat(type::String, nnlayers::Array{Int64}, file1::String)
    if type == "nn2array"
       nn = matread(file1)
       nnlayerNum = length(nnlayers) - 1
       theta_indptr = ones(Int64, 2*nnlayerNum+1)
       for i = 1:nnlayerNum
       	   theta_indptr[2*i]   = theta_indptr[2*i-1] + nnlayers[i]*nnlayers[i+1]
	       theta_indptr[2*i+1] = theta_indptr[2*i] + nnlayers[i+1]
       end

       theta = zeros(Float64, theta_indptr[end]-1)
       for i = 1:nnlayerNum
           for name in nn
               if occursin("connected_" * string(i) * "backslashweights", name.first)
                   theta[theta_indptr[2*(i+1)-1]:theta_indptr[2*(i+1)]-1] = transpose(name.second)[:]
               end
  	       if occursin("connected_" * string(i) * "backslashbiases", name.first)
                   theta[theta_indptr[2*(i+1)]:theta_indptr[2*(i+1)+1]-1] = transpose(name.second)[:]
               end
  	       if occursin("connectedbackslashweights", name.first)
                   theta[theta_indptr[1]:theta_indptr[2]-1] = transpose(name.second)[:]
               end
  	       if occursin("connectedbackslashbiases", name.first)
                   theta[theta_indptr[2]:theta_indptr[3]-1] = transpose(name.second)[:]
               end
           end
       end

       return theta
   end
end

ext: jl


text:
using FileIO, GeometryTypes
using Test
const tf = joinpath(dirname(@__FILE__), "testfiles")

function test_face_indices(mesh)
    for face in faces(mesh)
        for index in face
            @test firstindex(vertices(mesh)) <= index <= lastindex(vertices(mesh))
        end
    end
end

@testset "MeshIO" begin
    dirlen = 1f0
    baselen = 0.02f0
    mesh = [
        AABB{Float32}(Vec3f0(baselen), Vec3f0(dirlen, baselen, baselen)),
        AABB{Float32}(Vec3f0(baselen), Vec3f0(baselen, dirlen, baselen)),
        AABB{Float32}(Vec3f0(baselen), Vec3f0(baselen, baselen, dirlen))
    ]
    mesh = merge(map(GLPlainMesh, mesh))
    mktempdir() do tmpdir
        for ext in ["2dm", "off"]
            @testset "load save $ext" begin
                save(joinpath(tmpdir, "test.$ext"), mesh)
                mesh_loaded = load(joinpath(tmpdir, "test.$ext"), GLPlainMesh)
                @test mesh_loaded == mesh
            end
        end
        @testset "PLY ascii and binary" begin
            f = File(format"PLY_ASCII", joinpath(tmpdir, "test.ply"))
            save(f, mesh)
            mesh_loaded = load(joinpath(tmpdir, "test.ply"), GLPlainMesh)
            @test mesh_loaded == mesh
            save(File(format"PLY_BINARY", joinpath(tmpdir, "test.ply")), mesh)
        end
        @testset "STL ascii and binary" begin
            save(File(format"STL_ASCII", joinpath(tmpdir, "test.stl")), mesh)
            mesh_loaded = load(joinpath(tmpdir, "test.stl"), GLPlainMesh)
            #@test mesh_loaded == mesh
            save(File(format"STL_BINARY", joinpath(tmpdir, "test.stl")), mesh)
            mesh_loaded = load(joinpath(tmpdir, "test.stl"), GLNormalMesh)
            #@test mesh_loaded == mesh
            mesh_loaded = load(joinpath(tmpdir, "test.stl"), GLPlainMesh)
        end
    end
    @testset "Real world files" begin

        @testset "STL" begin
            msh = load(joinpath(tf, "ascii.stl"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 12
            @test length(vertices(msh)) == 36
            @test length(normals(msh)) == 36
            test_face_indices(msh)

            msh = load(joinpath(tf, "binary.stl"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 828
            @test length(vertices(msh)) == 2484
            @test length(normals(msh)) == 2484
            test_face_indices(msh)

            mktempdir() do tmpdir
                save(File(format"STL_BINARY", joinpath(tmpdir, "test.stl")), msh)
                msh1 = load(joinpath(tmpdir, "test.stl"))
                @test typeof(msh1) == GLNormalMesh
                @test faces(msh) == faces(msh1)
                @test vertices(msh) == vertices(msh1)
                @test normals(msh) == normals(msh1)
            end

            msh = load(joinpath(tf, "binary_stl_from_solidworks.STL"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 12
            @test length(vertices(msh)) == 36
            test_face_indices(msh)

            # STL Import
            msh = load(joinpath(tf, "cube_binary.stl"))
            @test length(vertices(msh)) == 36
            @test length(faces(msh)) == 12
            test_face_indices(msh)

            msh = load(joinpath(tf, "cube.stl"))
            @test length(vertices(msh)) == 36
            @test length(faces(msh)) == 12
            test_face_indices(msh)

        end
        @testset "PLY" begin
            msh = load(joinpath(tf, "ascii.ply"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 36
            test_face_indices(msh)
            @test length(vertices(msh)) == 72
            @test length(normals(msh)) == 72
            #msh = load(joinpath(tf, "binary.ply")) # still missing
            #@test typeof(msh) == GLNormalMesh
            #println(msh)
            msh = load(joinpath(tf, "cube.ply")) # quads
            @test length(vertices(msh)) == 24
            @test length(faces(msh)) == 12
            test_face_indices(msh)
        end
        @testset "OFF" begin
            msh = load(joinpath(tf, "test.off"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 28
            @test length(vertices(msh)) == 20
            @test length(normals(msh)) == 20
            test_face_indices(msh)

            msh = load(joinpath(tf, "test2.off"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 810
            @test length(vertices(msh)) == 405
            @test length(normals(msh)) == 405
            test_face_indices(msh)

            msh = load(joinpath(tf, "cube.off"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 12
            @test length(vertices(msh)) == 8
            test_face_indices(msh)

        end
        @testset "OBJ" begin
            msh = load(joinpath(tf, "test.obj"))
            @test typeof(msh) == GLNormalMesh
            @test length(faces(msh)) == 3954
            @test length(vertices(msh)) == 2248
            @test length(normals(msh)) == 2248
            test_face_indices(msh)

            msh = load(joinpath(tf, "cube.obj")) # quads
            @test length(faces(msh)) == 12
            @test length(vertices(msh)) == 8
            test_face_indices(msh)

            msh = load(joinpath(tf, "polygonal_face.obj"))
            @test length(faces(msh)) == 4
            @test length(vertices(msh)) == 6
            test_face_indices(msh)

            msh = load(joinpath(tf, "test_face_normal.obj"))
            @test length(faces(msh)) == 1
            @test length(vertices(msh)) == 3
            test_face_indices(msh)

        end
        @testset "2DM" begin
            msh = load(joinpath(tf, "test.2dm"))
            @test typeof(msh) == GLNormalMesh
            test_face_indices(msh)
            #@test length(faces(msh)) == 3954
            #@test length(vertices(msh)) == 2248
            #@test length(normals(msh)) == 2248
        end
    end
end

if false
    #using GLVisualize, GLAbstraction, FileIO
    w = glscreen()
    meshes = [visualize(load(joinpath(tf, name))) for name in readdir(tf)];
    meshes = convert(Matrix{Context}, reshape(meshes, (7, 2)));
    view(visualize(meshes))
    renderloop(w)
end
#=

amf1 = mesh(data_path*"pyramid.amf")
@test length(amf1[1].vertices) == 5
@test length(amf1[1].faces) == 4
@test length(amf1[2].vertices) == 5
@test length(amf1[2].faces) == 4

amf1 = mesh(data_path*"pyramid_zip.amf")
@test length(amf1[1].vertices) == 5
@test length(amf1[1].faces) == 4
@test length(amf1[2].vertices) == 5
@test length(amf1[2].faces) == 4
=#


text:
#!/usr/bin/env julia

##################################################
# Imports
##################################################

using Printf

# import PyPlot in version 2.8.0+
import Pkg
try
    using PyPlot
catch
    Pkg.add("PyPlot")
    using PyPlot
end
if PyPlot.version < v"2.8.0"
    error("Please upgrade PyPlot to version 2.8.0+")
end
Circle = matplotlib.patches.Circle

try
    using TimerOutputs
catch
    Pkg.add("TimerOutputs")
    using TimerOutputs
end
# Create the timer object
to = TimerOutput()

# import wrapper for the GDIP solution
include("GDIP.jl")

################################################
# Scenarios
################################################
#define planning problems:
    #  map file
    #  minimum turning radius
    #  sensing radius
    #  solver type
scenarios = [
    ("./problems/gdip-n10.txt", 1.0, 1.0, "DTSPN-GDIP"),
    #("./problems/gdip-n10.txt", 1.0, 0.5, "DTSPN-GDIP"),
    #("./problems/gdip-n10.txt", 0.5, 1.0, "DTSPN-GDIP"),
]

################################################
# Settings
################################################

visualize = true
show = true

# Save figures into a directory "images"
save_figures = false
if length(ARGS) > 0
    if "-save" == ARGS[1]
        visualize = true
        save_figures = true
    end
end
if save_figures
    mkpath("images")
end

##################################################
# Functions
##################################################

function load_map(filename)
    """Read config with goal positions"""
    goals = []
    lines = readlines(open(filename))
    for line in lines
        label, x, y = split(line)
        push!(goals, [parse(Float64,x), parse(Float64,y)])
    end
    return goals
end

function plot_points(points; specs = "b")
    x_val = [x[1] for x in points]
    y_val = [x[2] for x in points]
    plt.plot(x_val, y_val, specs)
end

function plot_circle(xy, radius)
    ax = plt.gca()
    circle = Circle(xy, radius, facecolor="yellow",edgecolor="orange", linewidth=1, alpha=0.2)
    ax.add_patch(circle)
end

function dist_euclidean(coord1, coord2)
    (x1, y1) = coord1
    (x2, y2) = coord2
    (dx, dy) = (x2 - x1, y2 - y1)
    return sqrt(dx * dx + dy * dy)
end

# cache results from computing length of the GDIP
lowerPathGDIPLenCache = Dict()

function lowerPathGDIP(s1, s2, turning_radius; cached = false)
    """Compute lower-bound path using GDIP between two configurations

    Arguments:
        s1 - start; s2 - end; turning_radius
        cached - compute only distance and cache the results
    Returns:
        Dubins maneuver (DubinsWrapper)
    """
    #if cached:
    #    key = (s1, s2)
    #    if key in lowerPathGDIPLenCache:
    #        length = lowerPathGDIPLenCache[key]
    #        return (None, length)

    interval1 = [s1.alpha1, s1.alpha2 - s1.alpha1]
    interval2 = [s2.alpha1, s2.alpha2 - s2.alpha1]
    #@timeit to "gdip"
    return gdip_init_gdip(s1.center, interval1, s1.radius, s2.center, interval2, s2.radius, turning_radius)

    #if cached:
    #    lowerPathGDIPLenCache[key] = length
    #    return (None, length)
    #else:
    #    return (path, length)
    return length
end

function upperPathGDIP(s1, s2, turning_radius; cached = false)
    """Compute feasible Dubins path two configurations

    Arguments:
        s1 - start; s2 - end; turning_radius
    Returns:
        (Dubins maneuver 'DubinsWrapper', length)
    """
    q1 = getFeasibleState(s1)
    q2 = getFeasibleState(s2)
    return gdip_init_dubins_maneuver(q1, q2, turning_radius)
end

function compute_distances(samples, dst_fce; turning_radius = 0)
    n = length(samples)
    distances = []

    for i in 1:n
        ss1 = samples[i]
        ss2 = samples[(i % n) + 1]

        n1 = length(ss1)
        n2 = length(ss2)

        sh = fill(Inf, (n1, n2))
        for i1 in 1:n1
            for i2 in 1:n2
                sh[i1,i2] = dst_fce(ss1[i1], ss2[i2], turning_radius, cached=true)
            end
        end
        push!(distances, sh)
    end

    return distances
end

function update_distances(samples, distances, selected_samples, dst_fce; turning_radius = 0)
    n = length(samples)
    for i in 1:n
        ss1 = samples[i]
        ss2 = samples[(i % n) + 1]

        n1 = length(ss1)
        n2 = length(ss2)

        sh_old = distances[i]
        siz_old = size(sh_old)
        sh = fill(NaN, (n1, n2))
        sh[1:siz_old[1],1:siz_old[2]] = sh_old

        #@show size(sh), selected_samples[i], sh[selected_samples[i], :]
        sh[selected_samples[i], :] *= NaN
        sh[:,selected_samples[(i%n)+1]] *= NaN

        # Alternative without going throw the whole distance matrix
        #=
        for j in 1:n1
            dist = dst_fce(ss1[j], ss2[n2], turning_radius, cached=true)[1]
            sh[j,n2] = dist
        end

        for j in 1:n2
            dist = dst_fce(ss1[n1], ss2[j], turning_radius, cached=true)[1]
            sh[n1,j] = dist
        end

        i1 = selected_samples[i]
        for i2 in 1:n2
            dist = dst_fce(ss1[i1], ss2[i2], turning_radius, cached=true)[1]
            sh[i1, i2] = dist
        end

        i2 = selected_samples[(i%n)+1]
        for i1 in 1:n1
            dist = dst_fce(ss1[i1], ss2[i2], turning_radius, cached=true)[1]
            sh[i1, i2] = dist
        end
        =#

        for i1 in 1:n1
            for i2 in 1:n2
                dist = sh[i1,i2]
                if isnan(dist)
                    @timeit to "dist_fce" begin
                        sh[i1,i2] = dst_fce(ss1[i1], ss2[i2], turning_radius, cached=true)
                    end
                end
            end
        end

        distances[i] = sh
    end

    return distances
end

function find_shortest_tour(distances)
    n = size(distances)[1]
    best_len = Inf
    best_tour = []

    # TODO - remove
    #global dist = distances

    # maximal number of samples
    k_max = maximum([size(x)[1] for x in distances])

    no_start = size(distances[1])[1]
    #@show no_start, k_max, n
    for start in 1:no_start
        # shortest sh[region_idx][sample_idx]
        sh = fill(Inf, (n+1, k_max))
        # used edge
        prev = fill(-1, (n+1, k_max))

        sh[1,start] = 0

        for region_idx in 1:n
            (n1, n2) = size(distances[region_idx])
            #@show n1, n2, distances[region_idx]
            for idx2 in 1:n2
                #@show distances[region_idx]
                #@show sh[region_idx,1:n1], distances[region_idx][:,idx2]
                dst = sh[region_idx,1:n1] + distances[region_idx][:,idx2]
                sh[region_idx+1,idx2] = minimum(dst)
                #@show dst
                #@show minimum(dst)
                prev[region_idx+1,idx2]= argmin(dst)
            end
        end

        #@show size(sh), n, start
        act_sol = sh[n+1,start]
        if act_sol < best_len
            best_len = act_sol
            tour = []
            act = start
            for i in 1:n
                act = prev[n-i+2,act]
                push!(tour, act)
            end
            best_tour = reverse(tour)
            #print(best_tour)
        end
    end

    return best_tour
end

function retrieve_path(samples, dst_fce, turning_radius, selected_samples)
    n = length(samples)
    path = []
    for a in 1:n
        g1 = samples[a][selected_samples[a]]
        g2 = samples[(a+1) % n][selected_samples[(a+1) % n]]
        path.append(dst_fce(g1, g2, turning_radius))
    end
    return path
end

function path_len(path)
    return sum_list([dub.get_length() for dub in path])
end

function plot_path(path, turning_radius, settings)
    step_size = 0.01 * turning_radius
    for dub in path
        configurations, _ = dub.sample_many(step_size)
        plot_points(configurations, settings)
    end
end

##################################################
# Target region given by the location and its sensing radius
##################################################
mutable struct TargetRegion
    center
    radius
end

function get_position_at_boundary(region::TargetRegion, beta)
    return region.center + region.radius * [cos(beta), sin(beta)]
end

##################################################
# Sample on the given target region
##################################################
mutable struct Sample
    target

    center
    radius

    alpha1
    alpha2
    alphaResolution

    beta1
    beta2
    betaResolution
end

Base.copy(s::Sample) = Sample(s.target, s.center, s.radius, s.alpha1, s.alpha2, s.alphaResolution, s.beta1, s.beta2, s.betaResolution)

function Sample_init(targetRegion)
    alpha1 = 0
    alpha2 = 2 * pi
    alphaResolution = 1

    beta1 = 0
    beta2 = 2 * pi
    betaResolution = 1

    return Sample(targetRegion, targetRegion.center, targetRegion.radius, alpha1, alpha2, alphaResolution, beta1, beta2, betaResolution)
end

# TODO missing functions

function Sample_plot(sample)
    ax = plt.gca()
    circle = Circle(sample.center, sample.radius, facecolor=nothing ,edgecolor="green", linewidth=1, alpha=0.2)
    ax.add_patch(circle)
end

function Sample_split(sample, resolution)
    """Split the actual sample into two new ones.
    The first is stored directly in the actual sample, and the second is returned.
    If the required resolution is already met, then nothing is done and None is returned.

    Parameters:
        resolution: the requred resolution
    Returns:
        Sample - the second new sample
    """
    # prefer splitting according position resolution
    if sample.betaResolution < resolution
        sam1 = copy(sample)
        sam2 = copy(sample)
        sam1.betaResolution = sam2.betaResolution = 2 * sample.betaResolution
        sam1.beta2 = sam2.beta1 = (sample.beta1 + sample.beta2) / 2
        update_center_radius(sam1)
        update_center_radius(sam2)
        return [sam1, sam2]
    end
    if sample.alphaResolution < resolution
        sam1 = copy(sample)
        sam2 = copy(sample)
        sam1.alphaResolution = sam2.alphaResolution = 2 * sample.alphaResolution
        sam1.alpha2 = sam2.alpha1 = (sample.alpha1 + sample.alpha2) / 2
        return [sam1, sam2]
    end
    return nothing
end

function update_center_radius(self)
    p1 = get_position_at_boundary(self.target, self.beta1)
    p2 = get_position_at_boundary(self.target, self.beta2)
    self.center = (p1 + p2) / 2
    self.radius = dist_euclidean(p1, p2) / 2
end

function getFeasibleState(self)
    pos = get_position_at_boundary(self.target, self.beta1)
    q = zeros(3)
    q[1:2] = pos
    q[3] = self.alpha1
    return q
end

##################################################
# Sampling structure which holds all the used samples
##################################################
mutable struct Sampling
    targets
    samples
end

function Sampling_init(centers, sensingRadius)
    targets = [TargetRegion(c, sensingRadius) for c in centers]
    samples = [[Sample_init(t)] for t in targets]
    return Sampling(targets, samples)
end

function refine_samples(self, selected, resolution)
    """Refine the seleted samples if the required resolution is not met.

    Parameters:
        slected: indexes of the selected samples (vector 1 x n)
        resolution: the requred resolution
    Returns:
        boolean - true if any sample is refined
    """
    n = length(self.samples)
    refined = false
    for i in 1:n
        to_split = selected[i]
        samp = self.samples[i][to_split]
        res = Sample_split(samp, resolution)
        if res != nothing
            self.samples[i][to_split] = res[1]
            push!(self.samples[i], res[2])
            refined = true
        end
    end
    return refined
end

##################################################
# The main solver class
##################################################
mutable struct GDIPSolver
    turning_radius::Float64
    sensing_radius::Float64

    sampling::Sampling

    goals

    lower_path::Array{Integer,1}
    upper_path::Array{Integer,1}

    lower_bound::Float64
    upper_bound::Float64
end

function GDIPSolver_init(turning_radius, goals, sensing_radius)
    return GDIPSolver(turning_radius, sensing_radius, Sampling_init(goals, sensing_radius), goals, [], [], 0, Inf)
end

function GDIPSolver_plot_map(solver)
    plt.clf()
    plt.axis("equal")
    plot_points(solver.goals, specs = "ro")
    if solver.sensing_radius != 0
        for goal in solver.goals
            plot_circle(goal, solver.sensing_radius)
        end
    end
end

function GDIPSolver_plot_tour_and_return_length(solver, selected_samples, maneuver_function, color)
    sampling = solver.sampling
    n = length(solver.sampling.samples)
    step_size = 0.01 * solver.turning_radius
    len = 0
    for a in 1:n
        b = (a%n)+1
        g1 = sampling.samples[a][selected_samples[a]]
        g2 = sampling.samples[b][selected_samples[b]]

        len += maneuver_function(g1, g2, solver.turning_radius)
        configurations = gdip_sample_many(step_size)
        if visualize
            plot_points(configurations, specs = color)
        end
    end
    return len
end

function GDIPSolver_plot_actual_and_return_bounds(solver)
    """Plot the actual sampling, lower and upper bound path

    Returns:
        (double, double) - lower bound, upper bound
    """
    if visualize
        GDIPSolver_plot_map(solver)
    end

    for s in solver.sampling.samples
        for ss in s
            Sample_plot(ss)
        end
    end

    lower_selected_samples = GDIPSolver_find_shortest_tour(solver, lowerPathGDIP)[1]
    upper_selected_samples = GDIPSolver_find_shortest_tour(solver, upperPathGDIP)[1]

    lower_bound = GDIPSolver_plot_tour_and_return_length(solver, lower_selected_samples, lowerPathGDIP, "r-")
    upper_bound = GDIPSolver_plot_tour_and_return_length(solver, upper_selected_samples, upperPathGDIP, "b-")
    return (lower_bound, upper_bound)
end

function GDIPSolver_find_shortest_tour(solver, fce; dist = nothing, selected = nothing)
    """Select the samples which represent the shortest lower bound tour

    Returns:
        indexes of the samples (vector 1 x n)
    """
    if dist == nothing || selected == nothing
        distances = compute_distances(solver.sampling.samples, fce, turning_radius = solver.turning_radius)
    else
        @timeit to "Update" distances = update_distances(solver.sampling.samples, dist, selected, fce, turning_radius = solver.turning_radius)
    end
    @timeit to "Find" selected_samples = find_shortest_tour(distances)
    return selected_samples, distances
end

##################################################
# Main loop over selected scenarios
##################################################
for scenario in scenarios
    # Load the problem and scenario settings
    filename = scenario[1]
    goals = load_map(filename)
    turning_radius = scenario[2]
    sensing_radius = scenario[3]
    solver_type = scenario[4]

    #tour planning part
    solver = GDIPSolver_init(turning_radius, goals, sensing_radius)
    GDIPSolver_plot_actual_and_return_bounds(solver)

    @printf("\n--- Problem: %s  Turning radius: %6.2f  Sensing radius: %6.2f  ---\n",
        filename, turning_radius, sensing_radius)

    if show
        plt.pause(0.1)
    end

    max_resolution = 64
    act_res = 4
    while act_res <= max_resolution
        refined = true
        distances = nothing
        selected_samples = nothing
        @timeit to "Iteration" begin
            while refined
                @timeit to "Tour" selected_samples, distances = GDIPSolver_find_shortest_tour(
                    solver, lowerPathGDIP, dist = distances, selected = selected_samples)
                @timeit to "Refine" refined = refine_samples(solver.sampling, selected_samples, act_res)
            end
        end
        (lower_bound, upper_bound) = GDIPSolver_plot_actual_and_return_bounds(solver)
        gap = (upper_bound - lower_bound) / upper_bound * 100.0
        @printf("Res: %4d  Lower: %6.2f  Upper: %6.2f  Gap(%%): %6.2f\n",
            act_res, lower_bound, upper_bound, gap)

        if visualize
            plt.title(@sprintf("Maximum resolution: %4d", act_res))
        end

        if show
            plt.pause(0.1)
        end

        if save_figures
            plt.savefig(@sprintf("images/dtrp-res-%04d.png", act_res))
        end

        act_res *= 2
    end

    @show to
    #show(to; allocations = false, sortby=:name)
    #show(to)
    print("\n")

    if show
        plt.pause(0.5)
    end
end

max_stars_repo_name: petvana/gdip

max_issues_repo_name: petvana/gdip

max_forks_repo_name: petvana/gdip


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

"""
    JuliaFEM.jl - an open source solver for both industrial and academia usage

The JuliaFEM software library is a framework that allows for the distributed
processing of large Finite Element Models across clusters of computers using
simple programming models. It is designed to scale up from single servers to
thousands of machines, each offering local computation and storage. The basic
design principle is: everything is nonlinear. All physics models are nonlinear
from which the linearization are made as a special cases.

# Examples

Typical workflow to use JuliaFEM to solve a partial differential equations,
is 1) read mesh, 2) create elements, 3) update element properties, 4) create
problems, 5) create analysis, and 6) run analysis. A simple linear static
analysis of elastic block is clarifying these steps.

```julia
using JuliaFEM
```

1. Usually the first thing to do is to create a geometry of domain. Typically
   this is done by reading a mesh file from disk. Currently JuliaFEM supports
   reading a mesh from Code Aster file format (using `aster_read_mesh`) and
   from ABAQUS file format (using `abaqus_read_mesh`).

```julia
mesh = aster_read_mesh("mesh.med")
```

2. Next step is to create one or several sets of elements from a mesh. This is
   done using a function `create_elements(mesh, set_name)`.

```julia
body_elements = create_elements(mesh, "body")
traction_elements = create_elements(mesh, "traction")
bc_elements = create_elements(mesh, "bc")
```

3. In JuliaFEM, all properties of elements are given using so called fields.
   Fields can depend from time or spatial coordinates elements. In special cases
   field is constant in time, spatial or both directions. Updating fields to
   element is done using `update!`-function.

```julia
update!(body_elements, "youngs modulus", 210.0e3)
update!(body_elements, "poissons ratio", 0.3)
update!(traction_elements, "surface pressure", 100.0)
update!(bc_elements, "displacement 1", 0.0)
update!(bc_elements, "displacement 2", 0.0)
update!(bc_elements, "displacement 3", 0.0)
```

4. The physics considered to be solve is given using `Problem`. Problem type can
   be e.g. `Elasticity` for hyperelasticity, `Heat` for solving the heat equation
   and so on. Problems are defined by giving the problem type as first argument,
   problem name in second argument and the last argument is giving the dimension
   of problem, meaning degrees of freedom connected to each node. After problems
   are created, elements are added to them by using function `add_elements!`.

```julia
body = Problem(Elasticity, "body", 3)
traction = Problem(Elasticity, "traction", 3)
bc = Problem(Dirichlet, "bc", 3, "displacement")
add_elements!(body, body_elements)
add_elements!(traction, traction_elements)
add_elements!(bc, bc_elements)
```

5. After geometry and physics is defined, we next define what kind of analysis
   are we going to perform. Analysis can be, for example, quasistatic analysis,
   analysis of dynamics of system, natural frequency analysis and so on. Some
   other special analysis types also exists, like performing model dimension
   reduction by creating super-elements or running optimization loop, given
   geometry, another analysis and initial conditions. For simplicity, we now
   create a linear quasistatic analysis of given problems. Problems are added
   to analysis using `add_problems!`.

```julia
analysis = Analysis(Linear)
add_problems!(analysis, body, traction, bc)
```

6. The last thing to do is to request the results of analysis to be written to
   disk for later use and actually perform the analysis. Currently, Xdmf output
   is supported, which can then be read using ParaView.

```julia
xdmf_output = Xdmf("analysis_results")
add_results_writer!(analysis, xdmf_output)
run!(analysis)
close(xdmf)
```

After analysis is ready, types and variables can be accessed using REPL or
Jupyter notebook for further postprocessing. Simulation can also be written
into a function to be a part of a larger analysis process. For more information
about JuliaFEM, please visit our website at

    www.juliafem.org

"""
module JuliaFEM

using SparseArrays, LinearAlgebra, Statistics
using Reexport, ForwardDiff, LightXML, HDF5, Parameters

import FEMSparse

@reexport using FEMBase
import FEMBase: get_unknown_field_name, get_unknown_field_dimension,
                assemble!, update!, initialize!
using FEMBase: get_problems

using TimerOutputs
export @timeit, print_timer

import Base: getindex, setindex!, convert, length, size, isapprox,
             similar, first, last, vec,
             ==, +, -, *, /, haskey, copy, push!, isempty, empty!,
             append!, read, copy

using AbaqusReader
using AsterReader

@reexport using HeatTransfer
include("problems_elasticity.jl")
export Elasticity
include("materials_plasticity.jl")
export plastic_von_mises
include("problems_dirichlet.jl")
export Dirichlet

export assemble!, postprocess!

# Structural elements: beams
@reexport using FEMBeam

### Mortar methods ###

@reexport using MortarContact2D
@reexport using MortarContact2DAD

include("problems_mortar.jl")
include("problems_mortar_3d.jl")
export calculate_normals, calculate_normals!, project_from_slave_to_master,
       project_from_master_to_slave, Mortar, get_slave_elements,
       get_polygon_clip
include("io.jl")
export Xdmf, h5file, xmffile, xdmf_filter, new_dataitem, update_xdmf!, save!
include("solvers.jl")
export AbstractSolver, Solver, Nonlinear, NonlinearSolver, Linear, LinearSolver,
       get_unknown_field_name, get_formulation_type, get_problems,
       get_field_problems, get_boundary_problems,
       get_field_assembly, get_boundary_assembly,
       initialize!, create_projection, eliminate_interior_dofs,
       is_field_problem, is_boundary_problem
include("solvers_modal.jl")
export Modal
include("problems_contact.jl")
include("problems_contact_3d.jl")
#include("problems_contact_3d_autodiff.jl")
export Contact

module Preprocess
end

using FEMBase, SparseArrays, LinearAlgebra
include("preprocess.jl")
export create_elements, Mesh, add_node!, add_nodes!,
       add_element_to_element_set!, add_node_to_node_set!,
       find_nearest_nodes, find_nearest_node, reorder_element_connectivity!,
       create_node_set_from_element_set!, filter_by_element_set
include("preprocess_abaqus_reader.jl")
export abaqus_read_mesh, create_surface_elements, create_nodal_elements
include("preprocess_aster_reader.jl")
export aster_read_mesh

# Postprocess module

module Postprocess
end

include("postprocess_utils.jl")
export calc_nodal_values!, get_nodal_vector, get_nodal_dict, copy_field!,
       calculate_area, calculate_center_of_mass, calculate_second_moment_of_mass,
       extract

include("deprecations.jl")

export SparseMatrixCOO, SparseVectorCOO, optimize!, resize_sparse
export DCTI, DVTI, DCTV, DVTV, CCTI, CVTI, CCTV, CVTV, Increment
export FieldProblem, BoundaryProblem, Problem, Node, Element, Assembly
export Poi1, Seg2, Seg3, Tri3, Tri6, Tri7, Quad4, Quad8, Quad9,
       Tet4, Tet10, Pyr5, Wedge6, Wedge15, Hex8, Hex20, Hex27
export update!, add_elements!, get_unknown_field_name, add!,
       is_field_problem, is_boundary_problem, get_gdofs,
       initialize!, get_integration_points, group_by_element_type,
       get_unknown_field_dimension, get_connectivity
export get_nonzero_rows, get_local_coordinates, inside, IP, get_element_type,
       get_elements, AbstractProblem, IntegrationPoint, filter_by_element_type,
       get_element_id, get_nonzero_columns, resize_sparse, resize_sparsevec

end


text:
module PortfolioOpt

import Base: size, length
using Distributions
import Distributions: rand
using JuMP
using LinearAlgebra
using LinearAlgebra: dot
using MathOptInterface
using OptimalBids
import MathOptInterface: LessThan, EqualTo, GreaterThan, constant, MAX_SENSE, MIN_SENSE
import OptimalBids: Market, change_bids!, clear_market!, calculate_profit
using PDMats
using Random
import Reexport
using Statistics
import Statistics: mean, cov
import Base: eachindex, eltype, size, keys
using UUIDs
using JuMP.Containers: DenseAxisArray

include("AmbiguitySet.jl")
include("formulations.jl")
include("VolumeMarket.jl")
include("DeterministicSamples.jl")
include("estimated_mean_variance.jl")
include("conditional_mean.jl")
include("robust_mean.jl")
include("dro_mean.jl")
include("backtest.jl")
include("simple_decision_rules.jl")
include("forecasts.jl")
include("test_utils/testutils.jl")

export AmbiguitySet,
    CenteredAmbiguitySet,
    PieceWiseUtility,
    Robustness,
    ExpectedReturn,
    Variance,
    SqrtVariance,
    ConditionalExpectedReturn,
    ExpectedUtility,
    RiskConstraint,
    ConeRegularizer,
    ObjectiveTerm,
    PortfolioFormulation,
    portfolio_model!,
    DeterministicSamples,
    MomentUncertainty,
    BudgetSet,
    EllipticalSet,
    EstimatedCase,
    WorstCase,
    # VolumeMarket
    VolumeMarket,
    change_bids!,
    clear_market!,
    market_budget,
    market_volume_fee,
    market_model,
    MarketHistory,
    VolumeMarketHistory,
    market_template,
    current_prices,
    risk_free_rate,
    # backtest
    sequential_backtest_market,
    get_records,
    # end-to-end
    max_sharpe,
    equal_weights,
    # forecasts
    mixed_signals_predict_return

Reexport.@reexport using JuMP

Reexport.@reexport using MathOptInterface: LessThan, EqualTo, GreaterThan

Reexport.@reexport using OptimalBids
end

max_stars_repo_name: andrewrosemberg/PortfolioOpt

max_issues_repo_name: andrewrosemberg/PortfolioOpt

max_forks_repo_name: andrewrosemberg/PortfolioOpt


text:
#!/usr/bin/env julia

import Parsers

const DATA_DIR   = "../14_data"
const INPUT_FILE = "input.txt"

# ----------------------------------------
macro assert(bool_expr)
    message = string("Assertion: ", bool_expr, " failed")
    :($(esc(bool_expr)) || error($message))
end
# ----------------------------------------

function banner(txt::String)
    println('='^10, ' ', txt, ' ', '='^10)
end

function get_adjvals(floating_bits)
    result = Int[]  # Array{Int,1}()
    num_floating_bits = length(floating_bits)
    for comb_id in 0:2^num_floating_bits-1
        comb_bitvec = get_bitvals(comb_id)
        adj_vals = [ 2^floating_bits[k] for (k,bitval) in enumerate(comb_bitvec)
                     if bitval == 1
                   ]
        adj_val = length(adj_vals) > 0 ? sum(adj_vals) : 0
        push!(result, adj_val)
    end
    return result
end

function get_bitvals(val)
    result = []
    for k in 0:100
        val_mod_2 = val % 2
        push!(result, val_mod_2 == 0 ? 0 : 1)
        val >>= 1
        if val == 0
            return result
        end
    end
end

function get_bitvec(val)
    result = BitVector()
    for k in 0:100
        val_mod_2 = val % 2
        push!(result, val_mod_2 == 0 ? 0 : 1)
        val >>= 1
        if val == 0
            return result
        end
    end
end

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

function part1(lines)
    banner("PART 1")
    mask0 = nothing
    mask1 = nothing
    memory = Dict{Int, Int}()
    for line in lines
        # println("line=$(line)")
        if line[1:4] == "mask"
            m = match(r"^[^=]+=\s([01X]+)", line)
            mask_str = m.captures[1]
            mask0 = 0
            mask1 = 0
            for k in 0:length(mask_str)-1
                c = mask_str[length(mask_str) - k]
                mask0 += (c == '0') ? 2^k : 0
                mask1 += (c == '1') ? 2^k : 0
            end
            # println("mask0 = $(mask0)")
            # println("mask1 = $(mask1)")
        elseif line[1:3] == "mem"
            m = match(r"^mem\[(\d+)\]\s=\s(\d+)", line)
            loc = parse(Int, m.captures[1])
            val = parse(Int, m.captures[2])
            val = (val | mask1) & (~mask0)
            # println("\tmem[$(loc)] = $(val)")
            memory[loc] = val
        end
    end
    mem_total = sum(values(memory))
    println("mem_total=$(mem_total)")
end

function part2(lines)
    banner("PART 2")
    mem0 = nothing
    mem1 = nothing
    mem_x = nothing
    floating_bits = nothing
    memory = Dict{Int, Int}()

    for line in lines
        # println("====================")
        # println("line=$(line)")
        if startswith(line, "mask")
            m = match(r"^[^=]+=\s([01X]+)", line)
            mask_str = m.captures[1]
            mem0 = 0
            mem1 = 0
            mem_x = 0
            mask_x_count = 0
            floating_bits = []

            for k in 0:length(mask_str)-1
                c = mask_str[length(mask_str) - k]
                mem1  += (c == '1') ? 2^k : 0
                mem_x += (c == 'X') ? 2^k : 0
                if c == 'X'
                    append!(floating_bits, k)  # length(mask_str)
                end
            end
            # println("mem0 = $(mem0)")
            # println("mem1 = $(mem1)")
        elseif startswith(line, "mem")
            m = match(r"^mem\[(\d+)\]\s=\s(\d+)", line)
            mem_raw = parse(Int, m.captures[1])
            val = parse(Int, m.captures[2])

            mem_base = (mem_raw | mem1) & ~mem_x
            for adjval in get_adjvals(floating_bits)
                mem = mem_base + adjval
                memory[mem] = val
            end
        end
    end
    mem_total = sum(values(memory))
    println("mem_total=$(mem_total)")
end

function main(filename::String)
    lines::Array{String} = readlines(filename)
    part1(lines)
    println()
    part2(lines)
end

main("$DATA_DIR/$INPUT_FILE")

max_stars_repo_name: jaycoskey/AdventOfCode

max_issues_repo_name: jaycoskey/AdventOfCode

max_forks_repo_name: jaycoskey/AdventOfCode


text:
using FileIO: Base
using JLD2
using Plots
using WindMixing
using OceanParameterizations

PATH = pwd()

EXTRACTED_FILE = "NDE_training_modified_pacanowski_philander_1sim_-1e-3_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_5e-3_rate_1e-4_extracted.jld2"
EXTRACTED_FILE_PATH = joinpath(PATH, "extracted_training_output", EXTRACTED_FILE)

OUTPUT_DIR = "D:\\University Matters\\MIT\\CLiMA Project\\OceanParameterizations.jl"
test_files = ["-1e-3"]
solve_oceananigans_modified_pacanowski_philander_nn(test_files, EXTRACTED_FILE_PATH, OUTPUT_DIR, timestep=60)

file = jldopen(EXTRACTED_FILE_PATH)
train_files = file["training_info/train_files"]
gradient_scaling = file["training_info/parameters"]["gradient_scaling"]

OUTPUT_SOL_DIR = joinpath(OUTPUT_DIR, "train_$(test_files[1])")
plot_data = NDE_profile_oceananigans(OUTPUT_SOL_DIR, train_files, test_files,
                                  ν₀=1f-1, ν₋=1f-4, ΔRi=1f-1, Riᶜ=0.25f0, Pr=1, gradient_scaling=gradient_scaling,
                                  OUTPUT_PATH=joinpath(OUTPUT_SOL_DIR, "profiles_fluxes.jld2"))

animation_type = "Training"
n_trainings = length(train_files)
training_types = "Wind Mixing, Free Convection"
VIDEO_NAME = "test2"
animate_profiles_fluxes_comparison(plot_data, joinpath(OUTPUT_SOL_DIR, VIDEO_NAME), fps=30,
                                                animation_type=animation_type, n_trainings=n_trainings, training_types=training_types)
close(file)

FILE_NAME = "NDE_training_modified_pacanowski_philander_1sim_-1e-3_diffusivity_1e-1_Ri_1e-1_weights_divide1f5_gradient_smallNN_scale_5e-3_rate_1e-4"
test_files = ["-1e-3", "-8e-4"]
timestep = 600
animate_training_results_oceananigans(test_files, timestep, FILE_NAME, OUTPUT_DIR)


text:
### Uniform(-1, 1)

_rand_nograd(rng, d, dims...) = _rand(rng, d, dims...)

struct UniformNoise{
    T<:AbstractFloat,
    D<:Union{Val{:cpu},Val{:gpu}},
    S<:Tuple{Vararg{Int}}
} <: ContinuousBatchDistribution
    size::S
    function UniformNoise(T::Type{<:AbstractFloat}, device::Symbol, size::S) where {S<:Tuple{Vararg{Int}}}
        return new{T,Val{device},S}(size)
    end
end

Base.size(d::UniformNoise) = d.size

mean_cpu(d::UniformNoise{T}) where {T} = zeros(T, d.size)
mean(d::UniformNoise{T,Val{:cpu}}) where {T} = mean_cpu(d)
var_cpu(d::UniformNoise{T}) where {T} = ones(T, d.size) / 3
var(d::UniformNoise{T,Val{:cpu}}) where {T} = var_cpu(d)

UniformNoise(T, dim::Int) = UniformNoise(T, :cpu, (dim,))
UniformNoise(size::Int...) = UniformNoise(Float64, :cpu, size)

_rand(rng, d::UniformNoise{T}, dims::Int...) where {T} = 2 * rand(rng, T, d.size..., dims...) .- 1

function rand(rng::AbstractRNG, d::UniformNoise{T,Val{:cpu}}, dims::Int...) where {T}
    return _rand_nograd(rng, d, dims...)
end

logpdf(d::UniformNoise{T}, x::AbstractArray{T}) where {T} = fill(-log(2one(T)), size(x))

###

mean(d::UniformNoise{T,Val{:gpu}}) where {T} = gpu(mean_cpu(d))
var(d::UniformNoise{T,Val{:gpu}}) where {T} = gpu(var_cpu(d))

function rand(rng::AbstractRNG, d::UniformNoise{T,Val{:gpu}}, dims::Int...) where {T}
    return gpu(_rand_nograd(rng, d, dims...))
end

### Standard Normal

struct GaussianNoise{
    T<:AbstractFloat,
    D<:Union{Val{:cpu},Val{:gpu}},
    S<:Tuple{Vararg{Int}}
} <: ContinuousBatchDistribution
    size::S
    function GaussianNoise(T::Type{<:AbstractFloat}, device::Symbol, size::S) where {S<:Tuple{Vararg{Int}}}
        return new{T,Val{device},S}(size)
    end
end

Base.size(d::GaussianNoise) = d.size

mean_cpu(d::GaussianNoise{T}) where {T} = zeros(T, d.size)
mean(d::GaussianNoise{T,Val{:cpu}}) where {T} = mean_cpu(d)
var_cpu(d::GaussianNoise{T}) where {T} = ones(T, d.size)
var(d::GaussianNoise{T,Val{:cpu}}) where {T} = var_cpu(d)

GaussianNoise(T, dim::Int) = GaussianNoise(T, :cpu, (dim,))
GaussianNoise(size::Int...) = GaussianNoise(Float64, :cpu, size)

_rand(rng, d::GaussianNoise{T}, dims::Int...) where {T} = randn(rng, T, d.size..., dims...)

function rand(rng::AbstractRNG, d::GaussianNoise{T,Val{:cpu}}, dims::Int...) where {T}
    return _rand_nograd(rng, d, dims...)
end

logpdf(d::GaussianNoise{T}, x::AbstractArray{T}) where {T} = -(log(2T(π)) .+ map(abs2, x)) / 2

###

mean(d::GaussianNoise{T,Val{:gpu}}) where {T} = gpu(mean_cpu(d))
var(d::GaussianNoise{T,Val{:gpu}}) where {T} = gpu(var_cpu(d))

function rand(rng::AbstractRNG, d::GaussianNoise{T,Val{:gpu}}, dims::Int...) where {T}
    return gpu(_rand_nograd(rng, d, dims...))
end


text:
"""
    PatchEmbedding(patch_size)
    PatchEmbedding(patch_height, patch_width)

Patch embedding layer used by many vision transformer-like models to split the input image into patches.
"""
struct PatchEmbedding
  patch_height::Int
  patch_width::Int
end

PatchEmbedding(patch_size) = PatchEmbedding(patch_size, patch_size)

function (p::PatchEmbedding)(x)
  h, w, c, n = size(x)
  hp, wp = h ÷ p.patch_height, w ÷ p.patch_width
  xpatch = reshape(x, hp, p.patch_height, wp, p.patch_width, c, n)

  return reshape(permutedims(xpatch, (1, 3, 5, 2, 4, 6)), p.patch_height * p.patch_width * c,
                 hp * wp, n)
end

@functor PatchEmbedding

"""
    ViPosEmbedding(embedsize, npatches; init = (dims) -> rand(Float32, dims))

Positional embedding layer used by many vision transformer-like models.
"""
struct ViPosEmbedding{T}
  vectors::T
end

ViPosEmbedding(embedsize, npatches; init = (dims::NTuple{2, Int}) -> rand(Float32, dims)) =
  ViPosEmbedding(init((embedsize, npatches)))

(p::ViPosEmbedding)(x) = x .+ p.vectors

@functor ViPosEmbedding

"""
    ClassTokens(dim; init = Flux.zeros32)

Appends class tokens to an input with embedding dimension `dim` for use in many vision transformer models.
"""
struct ClassTokens{T}
  token::T
end

ClassTokens(dim::Integer; init = Flux.zeros32) = ClassTokens(init(dim, 1, 1))

function (m::ClassTokens)(x)
  tokens = repeat(m.token, 1, 1, size(x, 3))
  return hcat(tokens, x)
end

@functor ClassTokens


text:
mutable struct ConnectionException <: Exception
    var::String
end
mutable struct QueryException <: Exception
    var::String
end
mutable struct NotImplementedException <: Exception
    var::String
end
mutable struct InvalidInputException <: Exception
    var::String
end

Base.showerror(io::IO, e::ConnectionException) = print(io, e.var)
Base.showerror(io::IO, e::QueryException) = print(io, e.var)
Base.showerror(io::IO, e::NotImplementedException) = print(io, e.var)
Base.showerror(io::IO, e::InvalidInputException) = print(io, e.var)

max_stars_repo_name: lokax/duckdb

max_issues_repo_name: lokax/duckdb

max_forks_repo_name: lokax/duckdb


text:
# Replication of Conesa et al "Taxing Capital"
using Parameters, QuantEcon, BasisMatrices,LinearAlgebra,Optim,DataFrames,Gadfly,SparseArrays,Arpack,Roots,Optim, BasisMatrices

####################################################################
#### Create a struct for the model parameters and setup utility ####
####################################################################

@with_kw mutable struct TaxMod
    # Demographics
    jr::Int64 = 46 # Retirement age
    J::Int64 = 81 # Maximum age
    nn::Float64 = 0.011 # Population growth
    ψ::Vector{Float64} = zeros(0) # survival probabilities
    ϵ_hansen::Vector{Float64} = zeros(0) # Age efficiency from Hansen paper
    pop::Vector{Float64} = zeros(0) # Used to calculate survival
    Nu::Vector{Float64} = zeros(0) # Number of agents in population
    μ::Vector{Float64} = zeros(0) # Fraction of agents in population
    measty::Vector{Float64} = zeros(0) # probability of each type
    topop::Float64 = 1.0 # Total population

    # Preferences
    β::Float64 = 1.00093 # Discount Factor
    σ::Float64 = 4.0 # Risk aversion
    γ::Float64 = 0.377 # Consumption share

    # Labor Productivity Process
    var_α::Float64 = 0.14 # Variance types
    ρ::Float64 = 0.98 # persistence
    var_η::Float64 = 0.0289 # variance shock
    Nη::Int64 = 7 # Number of states
    η::Vector{Float64} = zeros(0) # productivity shock states
    Π::Matrix{Float64} = zeros(0,0) # Transition matrix
    ϵ::Matrix{Float64} = zeros(0,0) # ability x age efficiency matrix

    # Technology
    α::Float64 = 0.36 # capital share
    δ::Float64 = 0.0833 # depreciation
    TFP::Float64 = 1.0 # scale parameter Z
    w::Float64 = 1.0 # wages
    r::Float64 = 0.05 # interest rate

    # Government Policy
    τc::Float64 = 0.05 # consumption tax
    τk::Float64 = 0.36 # capital tax
    κ0::Float64 = 0.258 # marginal tax
    κ1::Float64 = 0.768 # tax progressivity
    κ2::Float64 = 0.0 # deduction
    τp::Float64 = 0.124 # payroll tax
    b::Float64 = 0.5 # Social security replacement rate
    τ̄p::Float64 = 87000.0/37748.0 # Max social security rate
    ȳ::Float64  = 100000 # Social security threshold

    # Grid sizes
    ns::Int64 = 7 # number of states
    na::Int64 = 10 # asset grid
    nl::Int64 = 4  # leisure grid
    nty::Int64 = 2 # number of ability levels
    maxit::Int64 = 10000

    # Other
    a̲::Float64 = 0.0 # borrowing constraint
    umin::Float64 = -1.0E+2 # minimum utility
    penscale::Float64 = 10000000 # penalty on utility for bad c or l
    l̅::Float64 = 0.99 # maximum value for labor supply

    # Spending stuff
    G::Float64 = 6.670
    Tr::Float64 = 2.0
    SS::Float64 = 2.0

    #Solution
    agrid::Vector{Float64} = zeros(0)
    lgrid::Vector{Float64} = zeros(0)
    k::Int = 2 #type of interpolation
    Vf::Array{Interpoland} = Interpoland[] # value function given a
    Vlf::Array{Interpoland} = Interpoland[] # Value function given a,l
    cf::Array{Interpoland} = Interpoland[]
    lf::Array{Interpoland} = Interpoland[]
    a′f::Array{Interpoland} = Interpoland[]

end;

"""
    U(TM::TaxMod, c, l)

Calculates household utility given consumption, leisure, and preferences
"""
function U(TM::TaxMod, c, l)
    @unpack σ,γ,penscale,umin = TM

    if c <= 0.0 # Setting very low utility for negative consumption
		U = umin - penscale * c^2
	elseif l < 0.0 # Setting very low utility for negative labor supply
		U = umin - penscale * l^2
	elseif l >= 1.0 # Setting very low utility for labor supply over 1
		U = umin - penscale*(l-1.0)^2
	else # Real utility function
		U = (1/(1-σ)) * (((c^γ)*((1-l)^(1-γ)))^(1-σ))
	end

    return U
end

"""
    marginal_utility(TM::TaxMod, c, l)

Calculates marginal utility
"""
function marginal_utility(TM::TaxMod, c, l)
    @unpack σ,γ = TM

    if c > 0.0
        mu = (γ*(c^γ*(1.0-l)^(1.0-γ))^(1.0-σ))/c
	else # Very large MU if c is zero or negative
		mu = 1000000.0+ abs(c)^2.0
	end

end

"""
    tax_gs(y, k0, k1, k2)

Gouveia and Strauss tax function
"""
function tax_gs(y, k0, k1, k2)

    tax_gs = 0.0

    if k1 >= 0.001 && y > 0.0
        tax_gs = k0*(y - (y^(-k1)+k2)^(-1/k1))
    elseif k1 < 0.001 && y > 0.0 #
        tax_gs = k0*y + k2
    else # No taxes if income is negative/zero
        tax_gs = 0.0
    end

    return tax_gs
end

"""
    Tax3()

This tax function is separable in capital and labor income
"""

function Tax3(l0,l1,l2,ca0,lear,caear)

    Tax3=0.0

    # Labor Income Tax
    if l1 >= 0.000001 && lear >= 0.00000001
        Tax3 = l0*( lear - ( lear^(-l1)+l2 )^(-1.0/l1) )
    elseif l1 < 0.000001 && lear >= 0.00000001
        Tax3 = l0*lear+l2
    else
        Tax3 = 0.0
    end

    # Capital Income Tax
    Tax3 = Tax3 + ca0*caear

end

"""
    marginal_tax_gs(k0,k1,k2,cap_earn)
"""
function marginal_tax_gs(k0,k1,k2,cap_earn)

    mt = 0.0
    if k1 >= 0.001 && cap_earn > 0.0
        mt = k0*(1 - (1 + k2*cap_earn^k1)^(-1/k1 - 1))
    elseif k1 < 0.001 && cap_earn > 0.0
        mt = k0
    end

    return mt
end

"""
    marginal_tax3(ca0)
"""
function marginal_tax3(ca0)
    return ca0
end

####################################################################
################ Interpolating the Value function  #################
####################################################################

"""
    optimalpolicy(para::NCParameters,Vprime,k)

Computes optimal policy for aprime when retired given continuation value and
current period assets. Same for all types and states.
"""
function optimalPolicyRetired(TM::TaxMod,Vf′,a,j)
    @unpack τk,τc,SS,Tr,β,ψ,a̲,r = TM

    # Budget Constraint when retired
    c_ret(a′) = (SS+(1.0+r*(1.0-τk))*(a + Tr) - a′)/(1.0+τc)

    # Value function when retired, plugging in budget constraint for c
    Vfun(a′) = U(TM,c_ret(a′),0.) + β*ψ[j]*Vf′[1,1,j+1](a′)

    # Optimizing over a′∈[a̲,a(c=0)]
    ā = (SS+(1+r*(1-τk))*(a + Tr)) - 0.001
    result = optimize(Vfun,a̲,ā)
    aprime = result.minimizer
    V = -result.minimum

    return (aprime = aprime,V=V)
end;

function optimalPolicyWork(TM::TaxMod,Vf′,a,s,ty,j,l)
    @unpack τk,τc,τp,Tr,β,ψ,a̲,r,w,ϵ,η,κ0,κ1,κ2,ȳ,Π = TM

    earn = w*ϵ[ty,j]*η[s]*l
    #earncap = r*(a + Tr)
    lab_tax = tax_gs(κ0,κ1,κ2,earn)
    #cap_tax = τk*earncap # DIFF THAN REAL CODE
    ss_tax =  τp*min(earn,ȳ)
    #mtaxes = marginal_tax_gs(κ0,κ1,κ2,earn) + τk*earncap

    # Budget Constraint when working
    c_unret(a′) = (earn - ss_tax + (1+r*(1-τk))*(a + Tr) - lab_tax - a′)/(1+τc)

    # Objective function
    Vfun(a′) = U(TM,c_unret(a′),l) + β*ψ[j]*sum(map(eta -> Vf′[eta,ty,j+1](a′) * Π[s,eta],1:tm.ns))

    # Maximize choosing a′
    ā = (earn - ss_tax + (1+r*(1-τk))*(a + Tr) - lab_tax)*(1+τc) - 0.001
    result = optimize(Vfun,a̲,ā)
    aprime = result.minimizer
    V = -result.minimum

    return (aprime = aprime,V=V)
end;

"""
    iterateBellman(TM::TaxMod,Vf′)

Apply the bellman map given continuation value function Vprime
"""
function iterateBellman(TM::TaxMod,Vf′)
    @unpack ns,nty,J,SS,r,τk,Tr,τc,nl,lgrid,jr = TM

    # Setting up
    Vf = Array{Interpoland}(undef,ns,nty,J)
    anodes = nodes(Vf′[1,1,1].basis)
    len_node = length(anodes[1])
    V = zeros(len_node,ns,nty,J)
    #lf = zeros(len_node,ns,nty,J)
    Vr = zeros(len_node)
    aprime = zeros(len_node,ns,nty,J)
    Vl = zeros(len_node,nl)

    # Finding value function for given aprime and l
    for j in 1:J
        if j == J # Last period consume everything
            aprime[:,:,:,J] .= 0.
            for a in 1:len_node
                V[a,:,:,J] .= U(TM,(SS+(1.0+r*(1.0-τk))*(anodes[1][a] + Tr))/(1.0+τc),0.)
            end
            for s in 1:ns, ty in 1:nty
                basis = Vf′[s,ty,j].basis
                Vf[s,ty,j]= Interpoland(basis,V[:,s,ty,J])
            end
            #println("J done")
        elseif j >= jr # Retired folk
            basis = Vf′[1,1,j].basis
            Vr = [optimalPolicyRetired(TM,Vf′,a,j).V for a in nodes(basis)[1]]
            for s in 1:ns, ty in 1:nty
                V[:,s,ty,J] = Vr
                Vf[s,ty,j]= Interpoland(basis,V[:,s,ty,J])
            end
            #println("Retired $j done")
        else # Working folk
            Threads.@threads for s in 1:ns
                for  ty in 1:nty
                    basis = Vf′[s,ty,j].basis
                    for l in 1:nl # find optimal policy given a and l
                        Vl[:,l] = [optimalPolicyWork(TM,Vf′,a,s,ty,j,lgrid[l]).V for a in nodes(basis)[1]]
                    end
                    # for each a, choose l to maximize V
                    for ap in 1:len_node
                        V[ap,s,ty,j] = findmax(Vl[ap,:])[1]
                    end
                    Vf[s,ty,j]= Interpoland(basis,V[:,s,ty,J])
                end
            end
            #println("Working $j done")
        end
    end
    return Vf
end;

"""
    solvebellman!(TM::TaxMod)

Solves the bellman equation by iteration
"""
function solvebellman!(TM::TaxMod)
    @unpack ns,nty,J = TM

    diff = 1.
    counter = 0
    Vf′ = copy(TM.Vf)
    while diff > 1e-8
        TM.Vf = iterateBellman(TM,Vf′);
        diff = maximum([norm(Vf′[s,ty,j].coefs - TM.Vf[s,ty,j].coefs,Inf) for s in 1:ns, ty in 1:nty, j in 1:J]);
        counter += 1;
        println("Iteration $counter, Diff: $diff")
        Vf′ = TM.Vf;
    end

end;

"""
    setup_demographics!(TM::TaxMod)

Generates the age efficieny and population vectors
"""
function setup_demographics!(TM::TaxMod)
    @unpack J, nty, nn = TM

    # Age-Efficiency Units from Hansen (1993)
    ephansen = zeros(J)

    ephansen[1]=1.0000
    ephansen[2]=1.0719
    ephansen[3]=1.1438
    ephansen[4]=1.2158
    ephansen[5]=1.2842
    ephansen[6]=1.3527
    ephansen[7]=1.4212
    ephansen[8]=1.4897
    ephansen[9]=1.5582
    ephansen[10]=1.6267
    ephansen[11]=1.6952
    ephansen[12]=1.7217
    ephansen[13]=1.7438
    ephansen[14]=1.7748
    ephansen[15]=1.8014
    ephansen[16]=1.8279
    ephansen[17]=1.8545
    ephansen[18]=1.8810
    ephansen[19]=1.9075
    ephansen[20]=1.9341
    ephansen[21]=1.9606
    ephansen[22]=1.9623
    ephansen[23]=1.9640
    ephansen[24]=1.9658
    ephansen[25]=1.9675
    ephansen[26]=1.9692
    ephansen[27]=1.9709
    ephansen[28]=1.9726
    ephansen[29]=1.9743
    ephansen[30]=1.9760
    ephansen[31]=1.9777
    ephansen[32]=1.9700
    ephansen[33]=1.9623
    ephansen[34]=1.9546
    ephansen[35]=1.9469
    ephansen[36]=1.9392
    ephansen[37]=1.9315
    ephansen[38]=1.9238
    ephansen[39]=1.9161
    ephansen[40]=1.9084
    ephansen[41]=1.9007
    ephansen[42]=1.8354
    ephansen[43]=1.7701
    ephansen[44]=1.7048
    ephansen[45]=1.6396

    # Population Numbers from Bell and Miller (2002)
    pop = zeros(J)
    pop[1]=	197316
    pop[2]=	197141
    pop[3]=	196959
    pop[4]=	196770
    pop[5]=	196580
    pop[6]=	196392
    pop[7]=	196205
    pop[8]=	196019
    pop[9]=	195830
    pop[10]=195634
    pop[11]=195429
    pop[12]=195211
    pop[13]=194982
    pop[14]=194739
    pop[15]=194482
    pop[16]=194211
    pop[17]=193924
    pop[18]=193619
    pop[19]=193294
    pop[20]=192945
    pop[21]=192571
    pop[22]=192169
    pop[23]=191736
    pop[24]=191271
    pop[25]=190774
    pop[26]=190243
    pop[27]=189673
    pop[28]=189060
    pop[29]=188402
    pop[30]=187699
    pop[31]=186944
    pop[32]=186133
    pop[33]=185258
    pop[34]=184313
    pop[35]=183290
    pop[36]=182181
    pop[37]=180976
    pop[38]=179665
    pop[39]=178238
    pop[40]=176689
    pop[41]=175009
    pop[42]=173187
    pop[43]=171214
    pop[44]=169064
    pop[45]=166714
    pop[46]=164147
    pop[47]=161343
    pop[48]=158304
    pop[49]=155048
    pop[50]=151604
    pop[51]=147990
    pop[52]=144189
    pop[53]=140180
    pop[54]=135960
    pop[55]=131532
    pop[56]=126888
    pop[57]=122012
    pop[58]=116888
    pop[59]=111506
    pop[60]=105861
    pop[61]=99957
    pop[62]=93806
    pop[63]=87434
    pop[64]=80882
    pop[65]=74204
    pop[66]=67462
    pop[67]=60721
    pop[68]=54053
    pop[69]=47533
    pop[70]=41241
    pop[71]=35259
    pop[72]=29663
    pop[73]=24522
    pop[74]=19890
    pop[75]=15805
    pop[76]=12284
    pop[77]=9331
    pop[78]=6924
    pop[79]=5016
    pop[80]=3550
    pop[81]=2454

    # Labor efficiency
    ep = zeros(nty,J)
    for i in 1:nty
        ep[i,1:J] = ephansen[1:J]
    end

    # WHat is measty? Probability of high and low ability
    measty = zeros(nty)
    measty[1:nty] .= 1.0/nty

    # Survival probabilities: surv(i)=prob(alive in i+1|alive in i)
    surv = zeros(J)
    for i in  1:(J-1)
    surv[i] = pop[i+1]/pop[i]
    end
    surv[J] = 0.0

    # Number of Agents in population
    Nu = zeros(J)
    Nu[1] = 1.0
    for i in 2:J
        Nu[i] = surv[i-1]*Nu[i-1]/(1.0+nn)
    end

    # Fraction of agents in population
    mu = zeros(J)
    for i in 1:J
        mu[i] = Nu[i]/sum(Nu)
    end

    # Total population
    topop=sum(Nu)

    TM.ϵ_hansen = ephansen;
    TM.pop = pop;
    TM.ψ = surv;
    TM.Nu = Nu;
    TM.μ = mu;
    TM.ϵ = ep;
    TM.measty = measty
    TM.topop = topop
end

"""
    setup_grid!(TM::TaxMod, a̅ = 75, curve = 2.5)

Sets up grids, shocks, demographics
"""
function setup_grid!(TM::TaxMod, a̅ = 75, curve = 2.5)
    @unpack na,nl,a̲,l̅,J,nty,nn,ns,ρ,var_η,var_α,Vf,cf,k,β,ϵ_hansen = TM

    # Asset grid: with curvature
    agrid = TM.agrid = (a̅ - a̲).*LinRange(0,1,na).^curve .+ a̲;
    # Labor grid: evenly spaced
    TM.lgrid = LinRange(0,l̅,nl);

    # Labor Productivity Markov Chain
    mc = rouwenhorst(ns,ρ,var_η);
    TM.Π = mc.p;
    TM.η = exp.(mc.state_values);

    # Ability fixed effects and age efficiency
    TM.ϵ = zeros(nty,J);
    TM.ϵ[1,1:J] = exp(-sqrt(var_α)) .* ϵ_hansen[1:J];
    TM.ϵ[2,1:J] = exp(sqrt(var_α)) .* ϵ_hansen[1:J];

    # Setting up basis functions and interpoland
    abasis = Basis(SplineParams(agrid,0,k));
    anodes = nodes(abasis);
    Vf = Array{Interpoland}(undef,ns,nty,J);
    cons = zeros(length(anodes[1]),ns,nty,J);
    V = zeros(length(anodes[1]),ns,nty,J);

    # Creating initial guess for value function
    for s in 1:ns
        for ty in 1:nty
            for j in 1:J
                cons[:,s,ty,j] .=  anodes[1].*(1-((J-j)/J));
                V[:,s,ty,j] = map(x->U(TM,x,0.5)./(1-β),cons[:,s,ty,j]);
                Vf[s,ty,j]= Interpoland(abasis,V[:,s,ty,j]);
            end
        end
    end

    TM.Vf  = Vf;
      #, c0 = cf)
end

tm = TaxMod();
setup_demographics!(tm);
setup_grid!(tm);
solvebellman!(tm)

plot(layer(x->tm.Vf[3,2,60].(x),0,1000,color=["s=1,ty=1,j=1 Vf"]),
    Guide.xlabel("x"),Guide.ylabel("f(x)"),Guide.colorkey(title=""))

####################################################################
############### Find the stationary distribution  ##################
####################################################################

# Initializing
Φ = zeros(undef,ns,nty,J)

# 1 - Have policy function aprime
#

# For the first generation

Φ[a,s,ty,j]

####################################################################
############## Check if markets clear for policies  ################
####################################################################

function market_residuals(TM::TaxMod)
    @unpack r,N,κ2,Tr,SS,α,TFP,r,δ,Nu = TM

    # Find capital, output, and wages
    K = N*((α*TFP)/(r+δ))^(1.0/(1.0-α))
    Y = TFP*(K^α)*(N^(1.0-α))
    w = (1.0-α)*Y/N

    # Social security
    ȳ = τ̄p * Y / sum(Nu)

    # Solve the household problem given parameters
    solvebellman!(TM)

    # Find the stationary distribution

    # Compute the residuals
    return(
        fv1 = As-K*(1.0+nn),
        fv2 = LabS-N,
        fv3 = Govcons-tauc*C-Totinctax,
        fv4 = TrB-TrBn,
        fv5 = SS-SSn
    )

    # Goods Market Clearing
    Y = TFP * (K^α)*(LabS^(1.0-α))
    exdem = (C + As - (1.0-δ)*K + G - Y)

end

function iterateBellman_Serial(TM::TaxMod,Vf′)
    @unpack ns,nty,J,SS,r,τk,Tr,τc,nl,lgrid,jr = TM

    # Setting up
    Vf = Array{Interpoland}(undef,ns,nty,J)
    anodes = nodes(Vf′[1,1,1].basis)
    len_node = length(anodes[1])
    V = zeros(len_node,ns,nty,J)
    #lf = zeros(len_node,ns,nty,J)
    Vr = zeros(len_node)
    aprime = zeros(len_node,ns,nty,J)
    Vl = zeros(len_node,nl)

    # Finding value function for given aprime and l
    for j in 1:J
        if j == J # Last period consume everything
            aprime[:,:,:,J] .= 0.
            for a in 1:len_node
                V[a,:,:,J] .= U(TM,(SS+(1.0+r*(1.0-τk))*(anodes[1][a] + Tr))/(1.0+τc),0.)
            end
            for s in 1:ns, ty in 1:nty
                basis = Vf′[s,ty,j].basis
                Vf[s,ty,j]= Interpoland(basis,V[:,s,ty,J])
            end
            #println("J done")
        elseif j >= jr # Retired folk
            basis = Vf′[1,1,j].basis
            Vr = [optimalPolicyRetired(TM,Vf′,a,j).V for a in nodes(basis)[1]]
            for s in 1:ns, ty in 1:nty
                V[:,s,ty,J] = Vr
                Vf[s,ty,j]= Interpoland(basis,V[:,s,ty,J])
            end
            #println("Retired $j done")
        else # Working folk
            for s in 1:ns
                for  ty in 1:nty
                    basis = Vf′[s,ty,j].basis
                    for l in 1:nl # find optimal policy given a and l
                        Vl[:,l] = [optimalPolicyWork(TM,Vf′,a,s,ty,j,lgrid[l]).V for a in nodes(basis)[1]]
                    end
                    # for each a, choose l to maximize V
                    for ap in 1:len_node
                        V[ap,s,ty,j] = findmax(Vl[ap,:])[1]
                    end
                    Vf[s,ty,j]= Interpoland(basis,V[:,s,ty,J])
                end
            end
            #println("Working $j done")
        end
    end
    return Vf
end;

tm_p = TaxMod();
tm_s = TaxMod();
setup_grid!(tm_p);
setup_grid!(tm_s);

@time iterateBellman2(tm_p,tm_p.Vf);
@time iterateBellman(tm_s,tm_s.Vf);


text:
using DataFrames, CSV
data1 = CSV.read("../data/purchaseData.csv")
data2 = CSV.read("../data/purchaseData.csv", copycols=true)

try data1[1, :Name] = "YARDEN" catch; @warn "Cannot: data1 is immutable" end

data2[1, :Name] = "YARDEN"
println("\n", first(data2, 3), "\n")

data1[!, :Price] ./= 1000
rename!(data1, :Price=>Symbol("Price(000's)"))
println(first(data1, 3), "\n")

replace!(data2[!,:Grade], ["E"=>"F", "D"=>"E"]...)
println(first(data2, 3),"\n")

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook


text:
module Integrals

using OffsetArrays: OffsetArray
using StaticArrays: SVector, MMatrix
using BasisFunctions: Gaussian, cartesianToSphericalMxs
using Common: Float
using Overlap: buildCartesianOverlaps

export OneElectronIntegral, OverlapIntegral, contractedIntegrals, assembleCartesianIntegrals

"""
    OneElectronIntegral

Abstract type that has all one electron integrals as its subtypes.
"""
abstract type OneElectronIntegral end

"""
    OverlapIntegral <: OneElectronIntegral

Abstract type to represent one electron overlap integrals.
"""
abstract type OverlapIntegral <: OneElectronIntegral end

"""
    contractedIntegrals(::Type{OverlapIntegral}, Rμ::SVector{3}, cμ::SVector, ζμ::SVector, Rν::SVector{3}, cν::SVector, ζν::SVector, lμ::Int, lν::Int)

Calculates the overlap integrals between contracted cartesian gaussians at centers Rμ, Rν with expansion coeffiencts cμ, cν and exponents ζμ, ζν, up to angular momentums lμ and lν.
"""
function contractedIntegrals(::Type{OverlapIntegral}, Rμ::SVector{3, Float}, cμ::SVector{Nμ, Float}, ζμ::SVector{Nμ, Float}, Rν::SVector{3, Float}, cν::SVector{Nν, Float}, ζν::SVector{Nν, Float}, lμ::Int, lν::Int) where {Nμ, Nν}
    contracted = OffsetArray(zeros(Float, lμ+1, lν+1, 3), -1, -1, 0)
    for i=1:Nμ, j=1:Nν
        contracted += cμ[i] * cν[j] * buildCartesianOverlaps(Rμ, ζμ[i], lμ, Rν, ζν[j], lν)
    end
    return contracted
end

"""
    assembleCartesianIntegrlas(::Type{<:Integral}, S::IntsInCartDirs, lμ, lν)

Assemble all (lμ+1)*(lμ+2)÷2 * (lν+1)*(lν+2)÷2 cartesian integrals from S which contains the integrals factorized into cartesian directions.
"""
function assembleCartesianIntegrals(::Type{OneElectronIntegral}, lμ::Int, lν::Int)
#  function assembleCartesianGTOIntegrals(::Type{OneElectronIntegral}, S::AbstractArray, lμ::Int, lν::Int)
    Nμ = (lμ+1)*(lμ+2)÷2
    Nν = (lν+1)*(lν+2)÷2
    cartesianGTO = MMatrix{Nμ, Nν, Float}(undef)
    nν=0
    S = contractedIntegrals(OverlapIntegral, SVector{3, Float}(randn(3)), SVector{10, Float}(randn(10)), SVector{10, Float}(abs.(randn(10))), SVector{3, Float}(randn(3)), SVector{6, Float}(randn(6)), SVector{6, Float}(abs.(randn(6))), lμ, lν)
    for kν=0:lν
        for jν=0:lν-kν
            iν=lν-kν-jν
            nν += 1
            nμ = 0
            for kμ=0:lμ
                for jμ=0:lμ-kμ
                    iμ=lμ-kμ-jμ
                    nμ += 1
                    cartesianGTO[nμ, nν] = S[iμ, iν, 1] * S[jμ, jν, 2] * S[kμ, kν, 3]
                end
            end
        end
    end
    return cartesianGTO
end

function overlap()
    l = 8
    Scart = assembleCartesianIntegrals(OneElectronIntegral, l, l)
    cartToSpher = cartesianToSphericalMxs(l)
    Sspher = transpose(cartToSpher[l]) * Scart * cartToSpher[l]
    return Sspher
end

end