text:
"""
    abstract type AbstractChargeDensity{T <: SSDFloat} end

Struct defining the charge density inside a [`Passive`](@ref).

For each charge density, there should be a method [`get_charge_density`](@ref)
which returns the charge density in SI units (C/m³) at a given point `pt`.

## Examples
* [`ConstantChargeDensity`](@ref)
* [`LinearChargeDensity`](@ref)
* [`CylindricalChargeDensity`](@ref)
"""
abstract type AbstractChargeDensity{T <: SSDFloat} end

@inline function ChargeDensity(T::DataType, dict::Union{Dict{String, Any}, Dict{Any, Any}}, input_units::NamedTuple)
    return ChargeDensity(T, Val{Symbol(dict["name"])}(), dict, input_units)
end

"""
    get_charge_density(cd::AbstractChargeDensity, pt::AbstractCoordinatePoint)

Returns the charge density at a given point, `pt`, based on the charge density model `cd`.

## Arguments
* `cd::AbstractChargeDensity`: The [`AbstractChargeDensity`](@ref) defining the charge density inside a [`Passive`](@ref).
* `pt::AbstractCoordinatePoint`: The point at which `cd` is to be evaluated.

!!! note
    The value returned by `get_charge_density` is in units of C/m³ (SI units).
"""
function get_charge_density end

include("ConstantChargeDensity.jl")
include("LinearChargeDensity.jl")
include("CylindricalChargeDensity.jl")

ext: jl


text:
#X = randn(12,500)
#X[1,:] = linspace(0,1,500)
#X[2,:] = linspace(0,2,500)
#X[4,:] = linspace(-0.1,0,500)
#X[5,:] = linspace(-0.1,0,500)
#X[6,:] = linspace(0,3,500)
#t = linspace(0,5,500)
#options = Dict()
#options[:U] = randn(3,500)
#filename = "test.mp4"
using ProgressMeter

#plot(rand(100,4), layout = 4)

function visualise_ballbot(X, t, filename::String, des_traj; options::Dict = Dict())
    #pyplot(leg=false, ticks=nothing)

    # Find plotting limits
    offset = 1.5
    xmin = minimum(X[1, :])-offset
    xmax = maximum(X[1, :])+offset
    x_lim = (xmin, xmax)
    x_mean = sum(x_lim)/2
    x_range = abs.(x_mean-xmin)*1

    ymin = minimum(X[2, :])-offset
    ymax = maximum(X[2, :])+offset
    y_lim = (ymin, ymax)
    y_mean = sum(y_lim) / 2
    y_range = abs.(y_mean-ymin)*1

    #Idea scale equally by having same range:
    range_max = max(x_range, y_range)
    dx = 0.5*(range_max / x_range)*x_range
    dy = 0.5*(range_max / y_range)*y_range

    x_lim = (x_mean - dx, x_mean + dx)
    y_lim = (y_mean - dy, y_mean + dy)

    fps = get(options, :fps, 25)
    U = get(options, :U, zeros(3, length(t)))

    # Create animation
    plotsize = get(options, :size, (1280, 720))
    anim = Animation()

    # Create timestamps of the visualisation and interpolate states
    t_plot = t
    X_plot = X

    n_time = 4
    prog = Progress(length(t_plot))
    for i = 1:n_time:length(t_plot)
        fig = plot(;size = plotsize, layout = grid(3, 1, heights=[0.8,0.1,0.1]))
        plot_ballbot_3d!(fig[1], X[:,i], x_lim, y_lim, des_traj, i, X)
        plot!(fig[2], t[1:i], U[1,1:i], lab = "u1[V]", xlim = (0, t[end]), ylim = (-11, 11))
        plot!(fig[2], t[1:i], U[2,1:i], lab = "u2[V]")
        plot!(fig[2], t[1:i], U[3,1:i], lab = "u3[V]")
        plot!(fig[3], t[1:i], X[5,1:i] .* (180/pi), lab = "α[°]", xlim = (0, t[end]), ylim = (-20, 20))
        plot!(fig[3], t[1:i], X[4,1:i] .* (180/pi), lab = "β[°]")
        plot!(fig[3], t[1:i], X[6,1:i] .* (180/pi), lab = "γ[°]")
        frame(anim)

        #fig = plot(;size = plotsize)
        #plot_ballbot_3d!(fig, X_plot[:, i], x_lim, y_lim, des_traj, i, X)
        #frame(anim)

        ProgressMeter.update!(prog,i)
    end

    run(`ffmpeg -v 0 -framerate $(fps) -loop 0 -i $(anim.dir)/%06d.png -y $(filename)`)
    return nothing
end

function plot_ballbot_3d(xt, x_lim, y_lim, size)
    fig = plot(size=size)
    plot_ballbot_3d!(fig, xt, x_lim, y_lim)
    return fig
end

function plot_ballbot_3d!(fig, xt, x_lim, y_lim, des_traj, i, X)
    #Trajectory to do:
    plot!(fig,des_traj[1:i,1],des_traj[1:i,2],des_traj[1:i,1] .*0, lab ="desired position")

    #Trajectory done:
    plot!(fig,X[1,1:i],X[2,1:i],X[1,1:i].*0, lab="actual position")

    scale_it = 0.5

    range = x_lim[2] - x_lim[1]
    z_lim = (0, max(3*scale_it,range))
    plot!(fig, xlim = x_lim, ylim = y_lim, zlim = z_lim, ratio = :equal) # TODO zlim problematesch

    dx = xt[1]
    dy = xt[2]

    ϕ_y = xt[4]#*(pi/180)
    ϕ_x = -xt[5]#*(pi/180)
    ϕ_z = xt[6]#*(pi/180)

    radi = scale_it * 0.5

    fac_x = scale_it * 0.4
    fac_y = scale_it * 0.4
    fac_z = scale_it * 1

    dz   = 2 * radi + fac_z

    #Draw the Circle:
    u = linspace(0.0,2pi,20);
    v = linspace(0.0,pi,20);

    lu = length(u);
    lv = length(v);

    x = zeros(lu,lv);
    y = zeros(lu,lv);
    z = zeros(lu,lv);

    colors = zeros(lu,lv,3) ;

    colortube = 0.1:0.1:1.0;

    for uu=1:lu
        for vv=1:lv
            x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
            y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
            z[uu,vv]= radi*cos(v[vv]) + radi
        end
    end

    plot!(x,y,z, st = [:surface]);
    #surf(x,y,z)

    #Draw the Quader on Top
    x = -1:1:1
    y = -1:1:1

    lx = length(x)
    ly = length(y)

    X = zeros(lx,ly)
    Y = zeros(lx,ly)
    Z = zeros(lx,ly)

    for i_x = 1:lx
        for i_y = 1:ly
            X[i_x,i_y] = x[i_x]
            Y[i_x,i_y] = y[i_y]
            Z[i_x,i_y] = 1
        end
    end

    arr_X = [X,Y,Z,-Y,-Z,-Y]
    arr_Y = [Y,Z,X,-Z,-X,-Z]
    arr_Z = [Z,X,Y,-X,-Y,-Z]

    for i_side = 1:6
        #x-Koordinate Calculation
        x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
        x_koo += sin(ϕ_y)* arr_Z[i_side] * fac_z
        x_koo += dx + sin(ϕ_y)*2*radi
        #y-Koordinate Calculation
        y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
        y_koo += sin(ϕ_x)*arr_Z[i_side]* fac_z
        y_koo += dy + sin(ϕ_x)*2*radi
        #z-Koordinate Calculation
        z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
        z_koo -= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
        z_koo -= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
        z_koo += dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

        surface!(fig,x_koo,y_koo,z_koo, zcolor=5)
    end

    #xlabel!("x-position")
    #ylabel!("y-position")

end

function plot_all_ballbots(all_traj,filename::String, params; options::Dict = Dict())
    # Find plotting limits
    xmin = -1
    xmax = +1
    xlim = (xmin, xmax)
    x_mean = sum(xlim)/2
    x_range = abs.(x_mean-xmin)*2

    ymin = -1
    ymax = +1
    ylim = (ymin, ymax)
    y_mean = sum(ylim) / 2
    y_range = abs.(y_mean-ymin)*2

    #Idea scale equally by having same range:
    range_max = max(x_range, y_range)
    dx = 0.5*(range_max / x_range)*x_range
    dy = 0.5*(range_max / y_range)*y_range

    xlim = (x_mean - dx, x_mean + dx)
    ylim = (y_mean - dy, y_mean + dy)

    #Preselect good trajectory
    max_len = 0
    for i_traj = 1: params.n_traj
        if  all_traj.traj[i_traj].len > max_len
            max_len = all_traj.traj[i_traj].len
            i_best =  i_traj
        end
    end

    #Take data for video
    t_plot = 0:params.dt:params.t_horizon;
    t_plot = t_plot[1:max_len]

    fps::Int = get(options, :fps, 25)

    # Create animation
    plotsize::Tuple{Int, Int} = get(options, :size, (1280, 720))
    anim = Animation()

    n_step = 10
    n_time = 1

    #Go over all timesteps
    for i_t = 1:n_time:max_len
        fig = plot(size=plotsize)

        #Take all selected trajectories
        for i_traj = 1:n_step:params.n_traj
            i_des = 1
            X = zeros(2)

            #Check if trajectory is still active
            if all_traj.traj[i_traj].len >= i_t
                xt = all_traj.traj[i_traj].s_t[i_t,:]
                plot_ballbot_3d!(fig, xt, x_lim, y_lim, des_traj, i_des, X)
                #plot_segway_3d!(fig, xt, xlim, ylim)
            end
        end
        frame(anim)
    end

    run(`ffmpeg -v 0 -framerate $(fps) -loop 0 -i $(anim.dir)/%06d.png -y $(filename)`)
    return nothing
end

function plot_ball(fig, xt, x_lim, y_lim)
    scale_it = 0.5

    range = x_lim[2] - x_lim[1]
    z_lim = (0, max(3*scale_it,range))
    plot!(fig, xlim = x_lim, ylim = y_lim, zlim = z_lim, ratio = :equal) # TODO zlim problematesch

    dx = xt[1]
    dy = xt[2]

    ϕ_y = xt[4]#*(pi/180)
    ϕ_x = -xt[5]#*(pi/180)
    ϕ_z = xt[6]#*(pi/180)

    radi = scale_it * 0.5

    fac_x = scale_it * 0.4
    fac_y = scale_it * 0.4
    fac_z = scale_it * 1

    dz   = 2 * radi + fac_z

    #Draw the Circle:
    u = linspace(0.0,2pi,20);
    v = linspace(0.0,pi,20);

    lu = length(u);
    lv = length(v);

    x = zeros(lu,lv);
    y = zeros(lu,lv);
    z = zeros(lu,lv);

    colors = zeros(lu,lv,3) ;

    colortube = 0.1:0.1:1.0;

    for uu=1:lu
        for vv=1:lv
            x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
            y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
            z[uu,vv]= radi*cos(v[vv]) + radi
        end
    end

    plot!(x,y,z, st = [:surface]);
    #surf(x,y,z)

    #Draw the Quader on Top
    x = -1:1:1
    y = -1:1:1

    lx = length(x)
    ly = length(y)

    X = zeros(lx,ly)
    Y = zeros(lx,ly)
    Z = zeros(lx,ly)

    for i_x = 1:lx
        for i_y = 1:ly
            X[i_x,i_y] = x[i_x]
            Y[i_x,i_y] = y[i_y]
            Z[i_x,i_y] = 1
        end
    end

    arr_X = [X,Y,Z,-Y,-Z,-Y]
    arr_Y = [Y,Z,X,-Z,-X,-Z]
    arr_Z = [Z,X,Y,-X,-Y,-Z]

    for i_side = 1:6
        #x-Koordinate Calculation
        x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
        x_koo += sin(ϕ_y)* arr_Z[i_side] * fac_z
        x_koo += dx + sin(ϕ_y)*2*radi
        #y-Koordinate Calculation
        y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
        y_koo += sin(ϕ_x)*arr_Z[i_side]* fac_z
        y_koo += dy + sin(ϕ_x)*2*radi
        #z-Koordinate Calculation
        z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
        z_koo -= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
        z_koo -= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
        z_koo += dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

        surface!(fig,x_koo,y_koo,z_koo, zcolor=5)
    end

    xlabel!("x")
    ylabel!("y")
    gui()
end

function visualise_ballbott_compare(X,X2,u,u2,t, filename::String, des_traj; options::Dict = Dict())
    # Find plotting limits
    xmin = minimum(X[1, :])-2
    xmax = maximum(X[1, :])+2
    xlim = (xmin, xmax)
    x_mean = sum(xlim)/2
    x_range = abs.(x_mean-xmin)*2

    ymin = minimum(X[2, :])-2
    ymax = maximum(X[2, :])+2
    ylim = (ymin, ymax)
    y_mean = sum(ylim) / 2
    y_range = abs.(y_mean-ymin)*2

    #Idea scale equally by having same range:
    range_max = max(x_range, y_range)
    dx = 0.5*(range_max / x_range)*x_range
    dy = 0.5*(range_max / y_range)*y_range

    xlim = (x_mean - dx, x_mean + dx)
    ylim = (y_mean - dy, y_mean + dy)

    fps::Int = get(options, :fps, 25)

    # Create animation
    plotsize::Tuple{Int, Int} = get(options, :size, (1280, 720))
    anim = Animation()

    # Create timestamps of the visualisation and interpolate states
    t_plot = t[1]:(1/fps):t[end]
    X_plot = Array{Float64}(size(X, 1), length(t_plot))

    prog = Progress(length(t))
    n_step = 4
    for i = 1:n_step:length(t)

        fig = plot(;size = plotsize, layout = grid(3, 1, heights=[0.8,0.1,0.1]))
        plot_ballbot_3d_compare!(fig[1], X[:,i], X2[:,i], xlim, ylim, des_traj, i, X, X2)
        plot!(fig[2], t[1:i], u[1,1:i], lab = "u1_hard", xlim = (0, t[end]), ylim = (-4, 4))
        plot!(fig[2], t[1:i], u[2,1:i], lab = "u2_hard")
        plot!(fig[2], t[1:i], u[3,1:i], lab = "u3_hard")
        plot!(fig[3], t[1:i], u2[1,1:i], lab = "u1_soft", xlim = (0, t[end]), ylim = (-4, 4))
        plot!(fig[3], t[1:i], u2[2,1:i], lab = "u2_soft")
        plot!(fig[3], t[1:i], u2[2,1:i], lab = "u3_soft")
        frame(anim)

        ProgressMeter.update!(prog,i)
    end

    run(`ffmpeg -v 0 -framerate $(fps) -loop 0 -i $(anim.dir)/%06d.png -y $(filename)`)
    return nothing
end

function plot_ballbot_3d_compare!(fig, xt, xt2, xlim, ylim, des_traj, i, X1, X2)
    #Trajectory to do:
    plot!(fig,des_traj[1:i,1],des_traj[1:i,2],des_traj[1:i,1] .*0, lab ="desired position", c=:black)

    #Trajectory done:
    plot!(fig,X1[1,1:i],X1[2,1:i],X1[1,1:i].*0, lab="position PPO", c=:blue)

    #Trajectory done:
    plot!(fig,X2[1,1:i],X2[2,1:i],X2[1,1:i].*0, lab="position PPO-soft", c=:green)

    p_close = [2.0, 0.0]

    d_sim  = (p_close[1] - xt[1] )^2 + (p_close[2] - xt[2] )^2
    d_real = (p_close[1] - xt2[1])^2 + (p_close[2] - xt2[2])^2

    if d_sim > d_real
        """
        1st Ball
        """
        plot!(fig, xlim = x_lim, ylim = y_lim, zlim = z_lim, ratio = :equal) # TODO zlim problematesch

        dx = xt[1]
        dy = xt[2]

        ϕ_y = xt[4]#*(pi/180)
        ϕ_x = -xt[5]#*(pi/180)
        ϕ_z = xt[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = linspace(0.0,2pi,20);
        v = linspace(0.0,pi,20);

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        plot!(x,y,z, st = [:surface]);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo += sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo += dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo += sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo += dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo -= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo -= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo += dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, zcolor=5)
        end

        """
        2nd ball
        """

        dx = xt2[1]
        dy = xt2[2]

        ϕ_y = xt2[4]#*(pi/180)
        ϕ_x = -xt2[5]#*(pi/180)
        ϕ_z = xt2[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = linspace(0.0,2pi,20);
        v = linspace(0.0,pi,20);

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        plot!(x,y,z, st = [:surface]);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo += sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo += dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo += sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo += dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo -= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo -= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo += dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, zcolor=2)
        end

    else

        """
        2nd ball
        """

        plot!(fig, xlim = x_lim, ylim = y_lim, zlim = z_lim, ratio = :equal) # TODO zlim problematesch

        dx = xt2[1]
        dy = xt2[2]

        ϕ_y = xt2[4]#*(pi/180)
        ϕ_x = -xt2[5]#*(pi/180)
        ϕ_z = xt2[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = linspace(0.0,2pi,20);
        v = linspace(0.0,pi,20);

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        plot!(x,y,z, st = [:surface]);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo += sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo += dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo += sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo += dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo -= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo -= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo += dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, zcolor=2)
        end

        """
        1st ball
        """

        dx = xt[1]
        dy = xt[2]

        ϕ_y = xt[4]#*(pi/180)
        ϕ_x = -xt[5]#*(pi/180)
        ϕ_z = xt[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = linspace(0.0,2pi,20);
        v = linspace(0.0,pi,20);

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        plot!(x,y,z, st = [:surface]);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo += sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo += dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo += sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo += dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo -= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo -= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo += dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, zcolor=5)
        end

    end

end

max_stars_repo_name: NikonPic/REINFORCEMENT-JULIA

max_issues_repo_name: NikonPic/REINFORCEMENT-JULIA

max_forks_repo_name: NikonPic/REINFORCEMENT-JULIA


text:
function generate_data(ξ_true, p = 100, N = 100, zero_ratio = 0.0)
    #ξ_true of dim p
    println("Create a full design matrix and simulate data. The ratio N/p is $(N/p)")
    println("The ratio N/alpha is $(N/sum(ξ_true.!= 0))")
    X = [ones(N) randn(N,p-1)]
    ξ_true
    q = exp.(X*ξ_true)
    z = q ./(1 .+ q)
    y = Float64.(rand(N) .< z)
    return(X, y)
end


text:
"LU Decomposition"
module LUDecomposition

using LinearAlgebra: dot

function get_n(A::Matrix)::Int
    (n_r, n_c) = size(A)
    @assert n_r == n_c "Input matrix is not square"
    n_r
end

"Row permutation"
struct Permutation
    indices::Vector{Int}
    n_pivots::Int
end

raw"""
LU factorization with partial pivoting.

`A` is changed to $(L - E) + U$ s.t. $P A = L U$.
"""
function lu_decompose!(A::Matrix{<:AbstractFloat})::Union{Permutation, Nothing}
    n = get_n(A)
    p = collect(1:n)
    n_pivots = n  # for computing the determinant

    @inbounds for c in 1:n
        # Using ArrayView here doesn't help.
        # See https://github.com/JuliaLang/julia/issues/19198#issuecomment-257986870.
        (max_value, c_max) = findmax(abs.(A[c:n, c]))
        if max_value < eps()
            return nothing
        end
        c_max += c - 1

        if c_max != c
            @views A[[c, c_max], :] .= A[[c_max, c], :]
            p[c], p[c_max] = p[c_max], p[c]
            n_pivots += 1
        end

        # `A[r, c]` is assigned to 0 by `A[r, c] -= A[c, c] * (A[r, c] / A[c, c])`, so the element
        # in `L` is `(A[r, c] / A[c, c])`.
        rg = (c + 1):n;
        @views A[rg, c] /= A[c, c]
        @. @views A[rg, rg] -= A[rg, c] * A[c, rg]'
    end

    Permutation(p, n_pivots)
end

raw"""
Solves $A x = b$.

**Precondition**: `LU` is obtained by appling `lu_decompose!` on `A`.
"""
function lu_solve(LU::Matrix{T}, perm::Permutation, b::Vector{T})::Vector{T} where T<:AbstractFloat
    n = get_n(LU)

    # transforms `b`, stored in `x`
    x = b[perm.indices]
    for r in 1:n
        @views x[r] -= dot(LU[r, 1:(r - 1)], x[1:(r - 1)])
    end

    # computes `x`
    for r in n:-1:1
        p = @views dot(LU[r, (r + 1):n], x[(r + 1):n])
        x[r] = (x[r] - p) / LU[r, r]
    end

    x
end

function lu_det(LU::Matrix{T}, n_pivots::Int)::T where T<:AbstractFloat
    n = get_n(LU)
    det = foldl(*, LU[i, i] for i in 1:n)
    if iseven(n - n_pivots)
        det
    else
        -det
    end
end

#===== Unit Tests =====#
using Test

using LinearAlgebra: Diagonal, I, LowerTriangular, UpperTriangular
using LinearAlgebra: det  # for reference implementation
using Random: seed!

@testset "LU decomposition" begin
    n = 32
    seed!()
    A = rand(n, n)
    b = rand(n)

    LU = copy(A)
    perm = lu_decompose!(LU)
    if perm != nothing
        L = LowerTriangular(LU) - Diagonal(LU) + I
        U = UpperTriangular(LU)
        @test L * U ≈ A[perm.indices[:], :]

        @test lu_det(LU, perm.n_pivots) ≈ det(A)

        x = lu_solve(LU, perm, b)
        @test A * x ≈ b
    end
end

end # module

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
function maximal_excursion(r::AbstractMatrix, τ::Integer, k::Integer)
    @argcheck (k >= 1) && (τ >= 0)
    @argcheck size(r, 1) >= 2
    T = size(r, 1)
    m = displacement(r, 1)
    @inbounds for δt = 1:τ
        mₜ = displacement(r, δt)
        if mₜ > m
            m = mₜ
        end
    end
    m^k
end

function mme(
    df::DataFrame,
    id::Symbol = :TrackID,
    x::Symbol = :POSITION_X,
    y::Symbol = :POSITION_Y,
)
    result = DataFrame(n = Int[], τ = Int[], m1 = Float64[])
    @inbounds for n = 1:maximum(df[!, id])
        m = extract(df, n, id, [x, y])
        @simd for τ = 1:size(m, 1)-1
            m1 = maximal_excursion(m, τ, 1)
            push!(result, [n, τ, m1])
        end
    end
    result
end

function StatsBase.moment(
    df::DataFrame,
    idlabel::Symbol = :TrackID,
    xlabel::Symbol = :POSITION_X,
    ylabel::Symbol = :POSITION_Y,
)
    result = DataFrame(
        TrackID = Int64[],
        fourth_moment = Float64[],
        second_moment = Float64[],
        delta_t = Int64[],
        n = Int64[],
    )
    @inbounds for n in sort(collect(Set(df[!, idlabel])))
        m = extract(df, n, id, [xlabel, ylabel])
        T = size(m, 1)
        for δ = 1:T-1
            c⁴ = 0.0
            c² = 0.0
            @simd for t = 1:(T-δ)
                c⁴ += abs2(abs2(displacement(m, t, δ)))
                c² += squared_displacement(m, t, δ)
            end
            c⁴ /= (T - δ)
            c² /= (T - δ)
            push!(result, [n, c⁴, c², δ, T - δ])
        end
    end
    result
end

max_stars_repo_name: SosUts/DiffusionMoments

max_issues_repo_name: SosUts/DiffusionMoments

max_forks_repo_name: SosUts/DiffusionMoments


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "ConnectFourSolver"
version = v"0.1.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/findmyway/connect4.git", "db9b13d14951d35e4f3f21e4f2ecaa93a906f277"),
    ArchiveSource("https://julialang-s3.julialang.org/bin/linux/armv7l/1.3/julia-1.3.1-linux-armv7l.tar.gz", "965c8fab2214f8ce1b3d449d088561a6de61be42543b48c3bbadaed5b02bf824"; unpack_target="julia-arm-linux-gnueabihf"),
    ArchiveSource("https://julialang-s3.julialang.org/bin/linux/x64/1.3/julia-1.3.1-linux-x86_64.tar.gz", "faa707c8343780a6fe5eaf13490355e8190acf8e2c189b9e7ecbddb0fa2643ad"; unpack_target="julia-x86_64-linux-gnu"),
    ArchiveSource("https://github.com/Gnimuc/JuliaBuilder/releases/download/v1.3.0/julia-1.3.0-x86_64-apple-darwin14.tar.gz", "f2e5359f03314656c06e2a0a28a497f62e78f027dbe7f5155a5710b4914439b1"; unpack_target="julia-x86_64-apple-darwin14"),
    ArchiveSource("https://github.com/Gnimuc/JuliaBuilder/releases/download/v1.3.0/julia-1.3.0-x86_64-w64-mingw32.tar.gz", "c7b2db68156150d0e882e98e39269301d7bf56660f4fc2e38ed2734a7a8d1551"; unpack_target="julia-x86_64-w64-mingw32")
]

# Bash recipe for building across all platforms
script = raw"""
case "$target" in
	arm-linux-gnueabihf|x86_64-linux-gnu)
        Julia_PREFIX=${WORKSPACE}/srcdir/julia-$target/julia-1.3.1
        ;;
    x86_64-apple-darwin14|x86_64-w64-mingw32)
        Julia_PREFIX=${WORKSPACE}/srcdir/julia-$target/juliabin
        ;;
esac

mkdir connect4/build
cd connect4/build
cmake -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} -DJulia_PREFIX=$Julia_PREFIX -DCMAKE_PREFIX_PATH=$prefix/destdir/lib/cmake/JlCxx -DCMAKE_INSTALL_PREFIX=$prefix -DCMAKE_FIND_ROOT_PATH=${prefix} ..
VERBOSE=ON cmake --build . --config Release --target install -- -j${nproc}
install_license ${WORKSPACE}/srcdir/connect4/LICENSE
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Linux(:x86_64, libc=:glibc),
]

platforms = expand_cxxstring_abis(platforms)

# The products that we will ensure are always built
products = [
    LibraryProduct("libconnect4jl", :libconnect4jl),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("libcxxwrap_julia_jll")
]

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


text:
using LinearAlgebra, ForwardDiff, Distributions
using Convex, SCS, ECOS
using IterativeSolvers

function κ_no(z)
    return max.(0.0, z)
end

function Jκ_no(z)
    x = zero(z)
    for i = 1:length(x)
        if z[i] >= 0.0
            x[i] = 1.0
        end
    end
    return Diagonal(x)
end

function κ_soc(z)
    z1 = z[1:end-1]
    z2 = z[end]

    z_proj = zero(z)

    if norm(z1) <= z2
        z_proj = copy(z)
    elseif norm(z1) <= -z2
        z_proj = zero(z)
    else
        a = 0.5 * (1.0 + z2 / norm(z1))
        z_proj[1:end-1] = a * z1
        z_proj[end] = a * norm(z1)
    end
    return z_proj
end

function Jκ_soc(z)
    z1 = z[1:end-1]
    z2 = z[end]
    m = length(z)

    if norm(z1) <= z2
        return Diagonal(ones(m))
    elseif norm(z1) <= -z2
        return Diagonal(zeros(m))
    else
        D = zeros(m, m)
        for i = 1:m
            if i < m
                D[i, i] = 0.5 + 0.5 * z2 / norm(z1) - 0.5 * z2 * ((z1[i])^2.0) / norm(z1)^3.0
            else
                D[i, i] = 0.5
            end
            for j = 1:m
                if j > i
                    if j < m
                        D[i, j] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                        D[j, i] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                    elseif j == m
                        D[i, j] = 0.5 * z1[i] / norm(z1)
                        D[j, i] = 0.5 * z1[i] / norm(z1)
                    end
                end
            end
        end
        return D
    end
end

# problem setup
p = 1
n = p + 1

c = (1.0 + 0.0e-8) * ones(n)
G1 = Diagonal(ones(n))
h = zeros(n)
q = zeros(n)
y = (1.0 + 0.0e-8)

A = [-G1; -q']
b = [h; y]
m = size(A, 1)

"Convex.jl"
x = Variable(n)
prob = minimize(c' * x)
prob.constraints += norm(G1 * x + h) <= q' * x + y

optimizer = SCS.Optimizer
@time solve!(prob, optimizer)

@show prob.status
@show prob.constraints
@show x.value
k = n + m + 1

Q = Array([Diagonal(zeros(n)) A' c;
           -A zeros(m, m) b;
           -c' -b'      0.0])

r = m * n + m + n

function Pc(z)
    z_proj = zero(z)
    z_proj[1:n] = z[1:n]
    z_proj[n .+ (1:p+2)] = κ_soc(z[n .+ (1:p+2)])
    # z_proj[n + p + 2 + 1] = κ_no(z[n + p + 2 + 1])
    # z_proj[n + p + 2 + 2] = κ_no(z[n + p + 2 + 2])
    z_proj[n + m + 1] = max(0.0, z[n + m + 1])

    return z_proj
end

function Jc(z)
    J_proj = zeros(k, k)
    J_proj[1:n, 1:n] = Diagonal(ones(n))
    J_proj[n .+ (1:p+2), n .+ (1:p+2)] = Jκ_soc(z[n .+ (1:p+2)])
    # z_proj[n + p + 2 + 1] = κ_no(z[n + p + 2 + 1])
    # z_proj[n + p + 2 + 2] = κ_no(z[n + p + 2 + 2])
    J_proj[n + m + 1, n + m + 1] = z[n + m + 1] >= 0.0 ? 1.0 : 0.0

    return J_proj
end
function Pc_star(z)
    z_proj = zero(z)
    # z_proj[1:n] = z[1:n]
    z_proj[n .+ (1:p+2)] = κ_soc(z[n .+ (1:p+2)])
    # z_proj[n + p + 2 + 1] = κ_no(z[n + p + 2 + 1])
    # z_proj[n + p + 2 + 2] = κ_no(z[n + p + 2 + 2])
    z_proj[n + m + 1] = max(0.0, z[n + m + 1])

    return z_proj
end

function F(z)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2 * k .+ (1:k)]

    [(I + Q) * ũ - (u + v);
     u - Pc(ũ - v);
     ũ - u]
end

function Fθ(z, θ)
    A = reshape(θ[1:m * n], m, n)
    b = θ[m * n .+ (1:m)]
    c = θ[m * n + m .+ (1:n)]

    Q = Array([zeros(eltype(θ), n, n) A' c;
               -A zeros(eltype(θ), m, m) b;
               -c' -b'      0.0])

    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2 * k .+ (1:k)]

    [(I + Q) * ũ - (u + v);
     u - Pc(ũ - v);
     ũ - u]
end

function Ju(z)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2k .+ (1:k)]

    JP = zeros(m, m)
    JP[1:p+2, 1:p+2] = Jκ_soc(ũ[n .+ (1:m)][1:p+2] - v[n .+ (1:m)][1:p+2])
    # JP[p+2 .+ (1:1), p+2 .+ (1:1)] = Jκ_no(ũ[n .+ (1:m)][p+2 .+ (1:1)] - v[n .+ (1:m)][p+2 .+ (1:1)])
    # JP[p+2 + 1 .+ (1:1), p+2 + 1 .+ (1:1)] = Jκ_no(ũ[p+2 .+ (1:m)][p+2 + 1 .+ (1:1)] - v[p+2 .+ (1:m)][p+2 + 1 .+ (1:1)])

    if z[k] - z[3k] >= 0.0
        ℓ = 1.0
    else
        ℓ = 0.0
    end

    [-I zeros(n, m) zeros(n) I zeros(n, m) zeros(n) I zeros(n, m) zeros(n);
     zeros(m, n) -JP zeros(m) zeros(m, n) I zeros(m) zeros(m, n) JP zeros(m);
     [zeros(1, n) zeros(1, m) -ℓ zeros(1, n) zeros(1, m) 1.0 zeros(1, n) zeros(1, m) ℓ]]
end

function J(z)
    [(I + Q) -I -I
     Ju(z);
     I -I zeros(k, k)]
end

# ũ = zeros(k)
# u = zeros(k)
# v = zeros(k)
# ũ[end] = 1.0
# u[end] = 1.0
# v[end] = 1.0
#
# z = [ũ; u; v]
# z = rand(3k)
# F(z)
# rank(Ju(z))
# J(z)
# rank(ForwardDiff.jacobian(F, z))
# norm(J(z) - ForwardDiff.jacobian(F, z), Inf)
# gmres(J(z), -F(z))

function solve()
    ũ = 0.01 * rand(k)
    u = 0.01 * rand(k)
    v = 0.01 * rand(k)
    ũ[end] = 1.0
    u[end] = 1.0
    v[end] = 1.0

    z = [ũ; u; v]

    extra_iters = 0

    for i = 1:100
        _F = F(z)
        _J = J(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)

        iter = 0
        α = 1.0
        while norm(F(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"
                # return z
                x = z[k .+ (1:n)]
                y = z[k + n .+ (1:m)]
                τ = z[k + k]
                κ = z[3k]
                println("τ = $τ")
                println("κ = $κ")
                return x ./ τ, z_sol
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) -> norm: $(norm(F(z)))")

        z .-= α * Δ
    end

    # return z
    x = z[k .+ (1:n)]
    y = z[k + n .+ (1:m)]
    τ = z[k + k]
    κ = z[3k]
    println("τ = $τ")
    println("κ = $κ")
    return x ./ τ, z
end

x_sol, z_sol = solve()
@show norm(x_sol - x.value)
@show norm(x_sol[1:n] - x.value[1:n])
x_sol
# x.value
# F(z_sol)
# J(z_sol)

θ = [vec(A); b; c]

u_sol = z_sol[k .+ (1:k)]
v_sol = z_sol[2k .+ (1:k)]

function F_sol(z, θ)
    A = reshape(θ[1:m * n], m, n)
    b = θ[m * n .+ (1:m)]
    c = θ[m * n + m .+ (1:n)]

    Q = Array([zeros(eltype(θ), n, n) A' c;
               -A zeros(eltype(θ), m, m) b;
               -c' -b'      0.0])

   u = z[1:k]
   v = z[k .+ (1:k)]

   [Q * u - v;
    u - Pc(u - v)]
end

w_sol = [u_sol; v_sol]

F_sol(w_sol, θ)
_Fz(y) = F_sol(y, θ)
_Fθ(y) = F_sol(w_sol, y)

Jz = ForwardDiff.jacobian(_Fz, w_sol)
Jθ = ForwardDiff.jacobian(_Fθ, θ)
rank(Jz)
# (Jz \ Jθ)[1:n, m * n + (m-1) .+ (1:1+n)]

((Jz' * Jz + 1.0e-5 * I) \ (Jz' * Jθ))[1:n, m * n + (m-1) .+ (1:1+n)]

norm(z_sol[1:k] - Pc(z_sol[(1:k)]))
norm(Q * z_sol[1:k] - z_sol[2k .+ (1:k)])
norm(z_sol[2k .+ (1:k)] - Pc_star(z_sol[2k .+ (1:k)]))

Rmap(a, Q) = (Q - I) * Pc(a) + a
function _Rmap(a, θ)
    A = reshape(θ[1:m * n], m, n)
    b = θ[m * n .+ (1:m)]
    c = θ[m * n + m .+ (1:n)]

    Q = Array([zeros(eltype(θ), n, n) A' c;
               -A zeros(eltype(θ), m, m) b;
               -c' -b'      0.0])

    Rmap(a, Q)
end

function Rzmap(a, Q)
    (Q - I) * Jc(a) + I
end

function _Rzmap(a, θ)
    A = reshape(θ[1:m * n], m, n)
    b = θ[m * n .+ (1:m)]
    c = θ[m * n + m .+ (1:n)]

    Q = Array([zeros(eltype(θ), n, n) A' c;
               -A zeros(eltype(θ), m, m) b;
               -c' -b'      0.0])

    Rzmap(a, Q)
end
y_sol = z_sol[1:k] - z_sol[2k .+ (1:k)]
_Rmap(y_sol, θ)
tmp(w) = _Rmap(y_sol, w)
_Rzmap(y_sol, θ)

((_Rzmap(y_sol, θ)' * _Rzmap(y_sol, θ) + 1.0e-5 * I) \ (_Rzmap(y_sol, θ)' * ForwardDiff.jacobian(tmp, θ)))[1:n, m * n + (m-1) .+ (1:1+n)]


text:
using Test, PowerMonitor

CI = get(ENV, "CI", "false") === "true"

@show PowerMonitor.status(show_raw = true)
@testset "No batteries on CI platforms" begin
    if CI
        res = PowerMonitor.status(show_raw = true)
        @test res.pow_stat == :nobatteries
        @test res.perc == 100
    end
end

@testset "Automation: Auto-precomp" begin
    if CI
        @test_throws ErrorException PowerMonitor.autoprecomp_notbattery()
    end
    PowerMonitor.stop_automation()
end


text:
import Munkres: munkres

function neg_XᵀRY!{T}(M::Matrix{T}, X::Matrix{T}, Y::Matrix{T}, R::Matrix{T})
    n = size(X, 2)

    @inbounds for i in 1:n, j in 1:n
        s1 = X[1, i] * (R[1, 1] * Y[1, j] + R[1, 2] * Y[2, j] + R[1, 3] * Y[3, j])
        s2 = X[2, i] * (R[2, 1] * Y[1, j] + R[2, 2] * Y[2, j] + R[2, 3] * Y[3, j])
        s3 = X[3, i] * (R[3, 1] * Y[1, j] + R[3, 2] * Y[2, j] + R[3, 3] * Y[3, j])

        M[i, j] = -(s1 + s2 + s3)
    end

    return nothing
end

function wYPXᵀ!{T}(weight::T, Q::Matrix{T}, X::Matrix{T}, Y::Matrix{T}, p::Vector{Int})
    w = weight
    n = size(X, 2)

    @inbounds for i in 1:n
        Q[1, 1] += w * Y[1, p[i]] * X[1, i]
        Q[1, 2] += w * Y[1, p[i]] * X[2, i]
        Q[1, 3] += w * Y[1, p[i]] * X[3, i]

        Q[2, 1] += w * Y[2, p[i]] * X[1, i]
        Q[2, 2] += w * Y[2, p[i]] * X[2, i]
        Q[2, 3] += w * Y[2, p[i]] * X[3, i]

        Q[3, 1] += w * Y[3, p[i]] * X[1, i]
        Q[3, 2] += w * Y[3, p[i]] * X[2, i]
        Q[3, 3] += w * Y[3, p[i]] * X[3, i]
    end

    return nothing
end

function cost{T}(M::Matrix{T}, p::Vector{Int})
    n = size(M, 1)
    s = zero(T)

    @inbounds for i in 1:n
        s += M[i, p[i]]
    end

    return s
end

function cost{T}(weight::T, Q::Matrix{T}, M::Matrix{T}, X::Matrix{T}, Y::Matrix{T}, R::Matrix{T})
    neg_XᵀRY!(M, X, Y, R)
    p = munkres(M)
    wYPXᵀ!(weight, Q, X, Y, p)

    return weight * cost(M, p)
end

function partial_cost{T}(X::Matrix{T}, Y::Matrix{T})
    n = size(X, 2)
    s = zero(T)

    @inbounds for j in 1:n
        s1 = X[1, j] * X[1, j] + Y[1, j] * Y[1, j]
        s2 = X[2, j] * X[2, j] + Y[2, j] * Y[2, j]
        s3 = X[3, j] * X[3, j] + Y[3, j] * Y[3, j]

        s += s1 + s2 + s3
    end

    return s
end


text:
@testset "Sample" begin
    @testset "The Barn" begin
        for creature ∈ keys(CREATURES)
            @test getkey(CREATURES, creature, missing) == creature
        end
    end
end


text:
"""
    inputvertex(name, size)

Return an immutable input type vertex with the given `name` and `size`.

Typically used as "entry" point to a computation graph.

# Examples
```jldoctest
julia> using NaiveNASlib

julia> inputvertex("input", 5)
InputSizeVertex(InputVertex(input, outputs=[]), 5)

```
"""
inputvertex(name, size) = InputSizeVertex(name, size)

gen_vertex_doc_sig(fname) = """
    $(fname)computation, inputs::AbstractVertex...; traitdecoration=identity)
    $(fname)vname::AbstractString, computation, inputs::AbstractVertex...; traitdecoration=identity)
    $(fname)computation, vname::AbstractString, inputs::AbstractVertex...; traitdecoration=identity)
"""

"""
$(gen_vertex_doc_sig("vertex(trait::MutationTrait, "))

Return a mutable computation type vertex.

# Examples
```jldoctest
julia> using NaiveNASlib

julia> v = NaiveNASlib.vertex(NaiveNASlib.SizeInvariant(), x -> 5x, inputvertex("input", 1));

julia> v(3)
15
```
"""
vertex(trait::MutationTrait, computation, inputs::AbstractVertex...) = MutationVertex(CompVertex(computation, inputs...), trait)
function vertex(trait::MutationTrait, vname::AbstractString, computation, inputs::AbstractVertex...)
    othername = name(trait)
    if othername !== nothing
        @warn "Attach name $vname to vertex which already has name trait $(othername)!"
    end
    vertex(NamedTrait(vname, trait), computation, inputs...)
end
vertex(trait::MutationTrait, computation, vname::AbstractString, inputs::AbstractVertex...) = vertex(trait, vname, computation, inputs...)

"""
$(gen_vertex_doc_sig("immutablevertex("))

Return an immutable computation type vertex.

Use `traitdecoration` to attach other traits, such as [`named`](@ref), [`logged`](@ref) or [`validated`](@ref).

# Examples
```jldoctest
julia> using NaiveNASlib

julia> v = immutablevertex(x -> x * [1 2 3; 4 5 6], inputvertex("input", 2));

julia> v([1 2])
1×3 Matrix{Int64}:
 9  12  15

julia> v = immutablevertex("v", x -> x * [1 2 3; 4 5 6], inputvertex("input", 2));

julia> name(v)
"v"

julia> v = immutablevertex("v", inputvertex("input", 2)) do x
           x * [1 2 3; 4 5 6]
       end;

julia> name(v)
"v"
```
"""
immutablevertex(args...; traitdecoration=identity) = vertex(traitdecoration(Immutable()), args...)

"""
$(gen_vertex_doc_sig("absorbvertex("))

Return a mutable computation type vertex which absorbs size changes. Typical example of this is a neural network layer
wrapping a parameter array.

Use `traitdecoration` to attach other traits, such as [`named`](@ref), [`logged`](@ref) or [`validated`](@ref).

# Examples
```jldoctest
julia> using NaiveNASlib

julia> v = absorbvertex(x -> x * [1 2 3; 4 5 6], inputvertex("input", 2));

julia> v([1 2])
1×3 Matrix{Int64}:
 9  12  15

julia> v = absorbvertex("v", x -> x * [1 2 3; 4 5 6], inputvertex("input", 2));

julia> name(v)
"v"

julia> v = absorbvertex("v", inputvertex("input", 2)) do x
           x * [1 2 3; 4 5 6]
       end;

julia> name(v)
"v"
```
"""
absorbvertex(args...; traitdecoration=identity) = vertex(traitdecoration(SizeAbsorb()), args...)

"""
$(gen_vertex_doc_sig("invariantvertex("))

Return a mutable computation type vertex which is size invariant, i.e `nin == nout`.

Use `traitdecoration` to attach other traits, such as [`named`](@ref), [`logged`](@ref) or [`validated`](@ref).

# Examples
```jldoctest
julia> using NaiveNASlib

julia> v = invariantvertex(x -> 2 .* x, inputvertex("input", 2));

julia> nin(v)
1-element Vector{Int64}:
 2

julia> nout(v)
2

julia> v([1 2])
1×2 Matrix{Int64}:
 2  4

julia> v = invariantvertex("v", x -> 2 .* x, inputvertex("input", 2));

julia> name(v)
"v"

julia> v = invariantvertex("v", inputvertex("input", 2)) do x
           2 .* x
       end;

julia> name(v)
"v"
```
"""
invariantvertex(args...; traitdecoration=identity) = vertex(traitdecoration(SizeInvariant()), args...)

"""
    conc(v::AbstractVertex, vs::AbstractVertex...; dims, traitdecoration=identity, outwrap=identity)
    conc(vname::AbstractString, v::AbstractVertex, vs::AbstractVertex...; dims, traitdecoration=identity, outwrap=identity)

Return a mutable vertex which concatenates input along dimension `dim`.

Use `traitdecoration` to attach other traits, such as [`named`](@ref), [`logged`](@ref) or [`validated`](@ref).

Use `outwrap=f` to wrap the concatenation function in `f`.

`conc(vname::AbstractString,...;traitdecoration = f, ...)` is equivalent to `conc(...;traitdecoration=named(vname) ∘ f, ...)`.

# Examples
```jldoctest
julia> using NaiveNASlib

julia> v = conc(inputvertex.(["in1", "in2", "in3"], 1:3)..., dims=1);

julia> nin(v)
3-element Vector{Int64}:
 1
 2
 3

julia> nout(v)
6

julia> v([1], [2, 3], [4, 5, 6])
6-element Vector{Int64}:
 1
 2
 3
 4
 5
 6

julia> v = conc(inputvertex.(["in1", "in2", "in3"], 1:3)..., dims=1, outwrap = f -> (x...) -> 2f(x...));

julia> v([1], [2, 3], [4, 5, 6])
6-element Vector{Int64}:
  2
  4
  6
  8
 10
 12
```
"""
function conc(v::AbstractVertex, vs::AbstractVertex...; dims, traitdecoration=identity, outwrap=identity)
    vertex(traitdecoration(SizeStack()), outwrap((x...) -> cat(x..., dims=dims)), v, vs...)
end
function conc(vname::AbstractString, v::AbstractVertex, vs::AbstractVertex...; dims, traitdecoration=identity, outwrap=identity)
    vertex(traitdecoration(SizeStack()), vname, outwrap((x...) -> cat(x..., dims=dims)), v, vs...)
end

"""
    VertexConf
    VertexConf(; traitdecoration = identity, outwrap = identity)

Config struct to be used with element wise op syntax (`+`, `-`, `*`, `/`).

`traitdecoration` allows for decorating the vertex trait with stuff like logging, validation etc.
`outwrap` is a function which returns a function which will be applied to the computed output.
For example, the following `outwrap` scales output by a factor of 2: `outwrap = f ->  (x...) -> 2f((x...)`
"""
struct VertexConf{TD, OW}
    traitdecoration::TD
    outwrap::OW
end
traitconf(t) = VertexConf(traitdecoration=t)
outwrapconf(o) = VertexConf(outwrap=o)
VertexConf(;traitdecoration = identity, outwrap = identity)= VertexConf(traitdecoration, outwrap)

# Common wiring for all elementwise operations
function elemwise(op, conf::VertexConf, vs::AbstractVertex...)
    all(vi -> nout(vi) == nout(vs[1]), vs) || throw(DimensionMismatch("nout of all vertices input to elementwise vertex must be equal! Got $(nout.(vs))"))
    invariantvertex(conf.outwrap((x...) -> op.(x...)), vs...; conf.traitdecoration)
end

"""
    >>(conf::VertexConf, v::AbstractVertex)
    >>(name::AbstractString, v::AbstractVertex)
    >>(outwrap::Function, v::AbstractVertex)

Return inputs as a tuple. Only used to enable the `conf >> v1 op v2 ...` syntax.

"""
Base.:>>(conf::VertexConf, v::AbstractVertex) = (conf, v)
Base.:>>(name::AbstractString, v::AbstractVertex) = traitconf(named(name)) >> v
Base.:>>(outwrap::Function, v::AbstractVertex) = outwrapconf(outwrap) >> v

"""
    +(v::AbstractVertex, vs::AbstractVertex...)

Return a mutable vertex which performs (broadcasted) elementwise addition of its inputs.

An `AbstractString`, `Function` or `VertexConf` can be supplied through the `>>` operator.
An `AbstractString` will be used as the name of the vertex, a `Function f` will wrap the
output so that the vertex computation becomes `f(+.(x...))` and `VertexConf` can be used
to supply both a name and a function.

# Examples

```jldoctest
julia> using NaiveNASlib

julia> v = inputvertex("in1", 2) + inputvertex("in2", 2) + inputvertex("in3" ,2);

julia> nin(v)
3-element Vector{Int64}:
 2
 2
 2

julia> nout(v)
2

julia> v([1, 2], [3, 4], [5, 6])
2-element Vector{Int64}:
  9
 12

julia> v = "v" >> inputvertex("in1", 3) + inputvertex("in2", 3);

julia> name(v)
"v"
```
"""
Base.:+((conf, v)::Tuple{VertexConf, <:AbstractVertex}, vs::AbstractVertex...) = elemwise(+, conf, v, vs...)
Base.:+(v::AbstractVertex, vs::AbstractVertex...) = +(VertexConf() >> v, vs...)

"""
    *(v::AbstractVertex, vs::AbstractVertex...)
    *(conf::VertexConf, v::AbstractVertex, vs::AbstractVertex...)

Return a mutable vertex which performs (broadcasted) elementwise multiplication of its inputs.

An `AbstractString`, `Function` or `VertexConf` can be supplied through the `>>` operator.
An `AbstractString` will be used as the name of the vertex, a `Function f` will wrap the
output so that the vertex computation becomes `f(*.(x...))` and `VertexConf` can be used
to supply both a name and a function.

# Examples

```jldoctest
julia> using NaiveNASlib

julia> v = inputvertex("in1", 2) * inputvertex("in2", 2) * inputvertex("in3" ,2);

julia> nin(v)
3-element Vector{Int64}:
 2
 2
 2

julia> nout(v)
2

julia> v([1, 2], [3, 4], [5, 6])
2-element Vector{Int64}:
 15
 48

julia> v = "v" >> inputvertex("in1", 3) * inputvertex("in2", 3);

julia> name(v)
"v"
```
"""
Base.:*((conf, v)::Tuple{VertexConf, <:AbstractVertex}, vs::AbstractVertex...) = elemwise(*, conf, v, vs...)
Base.:*(v::AbstractVertex, vs::AbstractVertex...) = *(VertexConf() >> v, vs...)

"""
    -(v1::AbstractVertex, v2::AbstractVertex)

Return a mutable vertex which performs (broadcasted) elementwise subtraction of its inputs.

An `AbstractString`, `Function` or `VertexConf` can be supplied through the `>>` operator.
An `AbstractString` will be used as the name of the vertex, a `Function f` will wrap the
output so that the vertex computation becomes `f(-.(x...))` and `VertexConf` can be used
to supply both a name and a function.

# Examples

```jldoctest
julia> using NaiveNASlib

julia> v = inputvertex("in1", 2) - inputvertex("in2", 2);

julia> nin(v)
2-element Vector{Int64}:
 2
 2

julia> nout(v)
2

julia> v([1, 2], [3, 4])
2-element Vector{Int64}:
 -2
 -2

julia> v = "v" >> inputvertex("in1", 3) - inputvertex("in2", 3);

julia> name(v)
"v"
```
"""
Base.:-((conf, v1)::Tuple{VertexConf, <:AbstractVertex}, v2::AbstractVertex) = elemwise(-, conf, v1, v2)
Base.:-(v1::AbstractVertex, v2::AbstractVertex) = -(VertexConf() >> v1, v2)

"""
    -(v::AbstractVertex)

Return a mutable vertex which performs elementwise negation of its input.

An `AbstractString`, `Function` or `VertexConf` can be supplied through the `>>` operator.
An `AbstractString` will be used as the name of the vertex, a `Function f` will wrap the
output so that the vertex computation becomes `f(-.(x...))` and `VertexConf` can be used
to supply both a name and a function.
Due to operator precedence, this has to be done in the following order: `-(conf >> v)`

#Examples

```jldoctest
julia> using NaiveNASlib

julia> v = -inputvertex("in", 2);

julia> v([1,2])
2-element Vector{Int64}:
 -1
 -2

julia> v = -("v" >> inputvertex("in", 2));

julia> name(v)
"v"
```
"""
Base.:-((conf, v)::Tuple{VertexConf, <:AbstractVertex}) = elemwise(-, conf, v)
Base.:-(v::AbstractVertex) = -(VertexConf() >> v)

"""
    /(v1::AbstractVertex, v2::AbstractVertex)
    /(conf::VertexConf, v1::AbstractVertex, v2::AbstractVertex)

Return a mutable vertex which performs (broadcasted) elementwise division of its inputs.

An `AbstractString`, `Function` or `VertexConf` can be supplied through the `>>` operator.
An `AbstractString` will be used as the name of the vertex, a `Function f` will wrap the
output so that the vertex computation becomes `f(/.(x...))` and `VertexConf` can be used
to supply both a name and a function.

# Examples

```jldoctest
julia> using NaiveNASlib

julia> v = inputvertex("in1", 2) / inputvertex("in2", 2);

julia> nin(v)
2-element Vector{Int64}:
 2
 2

julia> nout(v)
2

julia> v([6, 8], [2, 4])
2-element Vector{Float64}:
 3.0
 2.0

julia> v = "v" >> inputvertex("in1", 3) / inputvertex("in2", 3);

julia> name(v)
"v"
```
"""
Base.:/((conf, v1)::Tuple{VertexConf, <:AbstractVertex}, v2::AbstractVertex) = elemwise(/, conf, v1, v2)
Base.:/(v1::AbstractVertex, v2::AbstractVertex) = /(VertexConf() >> v1, v2)

"""
    named(name)

Return a function `t -> NamedTrait(name, t)` intended to reduce a bit of the verbosity when using `traitdecoration`.

Intended to be composable with other similar functions through ∘.
"""
named(name) = t -> NamedTrait(name, t)

"""
    validated(args...;kwargs...)

Return a function `t -> AfterΔSizeTrait(validateafterΔsize(args...;kwargs...), t)` indended to reduce a bit of the verbosity when using `traitdecoration`.

Intended to be composable with other similar functions through ∘.
"""
validated(args...;kwargs...) = t -> AfterΔSizeTrait(validateafterΔsize(args...;kwargs...), t)

"""
    logged(args...;kwargs...)

Return a function `t -> AfterΔSizeTrait(logafterΔsize(args...;kwargs...)` indended to reduce a bit of the verbosity when using `traitdecoration`.

Intended to be composable with other similar functions through ∘.
"""
logged(args...;kwargs...) = t -> AfterΔSizeTrait(logafterΔsize(args...;kwargs...), t)


text:
#!/usr/bin/env julia

using FASTX

seqFile = "data/hla_gen.fa"

name2seq = Dict()
open(FASTA.Reader,seqFile) do reader
	for record in reader
		name2seq[FASTA.identifier(record)]=FASTA.sequence
	end
end

print(name2seq["HLA:HLA16652_A*01:01:01:09_3340_bp"])

# time ./seqbyname.jl
# sequence./seqbyname.jl  1.28s user 0.27s system 112% cpu 1.375 total

# ./seqbyname.py  0.50s user 0.09s system 99% cpu 0.598 total

max_stars_repo_name: jdurbin/sandbox

max_issues_repo_name: jdurbin/sandbox

max_forks_repo_name: jdurbin/sandbox


text:
include("utils.jl")
include("sized_abstract_array.jl")

# starting with tests of simple manifolds
include("euclidean.jl")
include("sphere.jl")
include("rotations.jl")
include("symmetric_positive_definite.jl")
include("cholesky_space.jl")
include("product_manifold.jl")
include("power_manifold.jl")
include("vector_bundle.jl")

include("metric.jl")


text:
# # Logistic Regression
# The presented example is adapted from the [Machine Learning - course by Andrew Ng](https://www.coursera.org/learn/machine-learning).
#
# In this example we use logistic regression to estimate the parameters $\theta_i$ of a logistic model in a classification problem. We will first  transform the logistic regression problem into an exponential cone optimisation problem. We will then solve the optimisation problem with COSMO  and determine the model parameters and the decision boundary.
# ## Visualizing the data
# Before we start, let's load and take a look at the example data from `examples/chip_data.txt`:

#-
using LinearAlgebra, SparseArrays, COSMO, JuMP
using Plots

## load example data
f = open(joinpath(@__DIR__, "chip_data.txt"))
lines = readlines(f)
close(f)
n_data = length(lines)
n_half = div(n_data, 2)
x1 = zeros(n_data)
x2 = zeros(n_data)
y = zeros(Float64, n_data)
for (i, l) in enumerate(lines)
    s = split(l, ",")
    x1[i] = parse(Float64, s[1])
    x2[i] = parse(Float64, s[2])
    y[i] = parse(Float64, s[3])
end

## visualize data
plot(x1[1:n_half-1], x2[1:n_half-1], color = :blue, st=:scatter, markershape = :cross, aspect_ratio=:equal, label = "Accepted", xlabel = "x1 - Microchip Test Score 1", ylabel = "x2 - Microchip Test Score 2")
plot!(x1[n_half:end], x2[n_half:end], color = :red, st=:scatter, markershape = :circle, label = "Rejected")

# The plot shows two test scores of $n$ microchip samples from a fabrication plant and whether the chip passed the quality check. Based on this data we would like to build a logistic model that takes into account the test scores and helps us predict the likelihood of a chip being accepted.
#
# ## Defining the logistic model
# The logistic regression hypothesis is given by
#
# $$
# h_\theta(x) = g(\theta^\top x)
# $$
# where $g$ is the sigmoid function:
# $$
# g(\theta^\top x) = \frac{1}{1+\exp(-\theta^\top x)}.
# $$
#
# The vector $x$ represents the independent variables and $\theta$ represents the model parameters. For our samples we set the dependent variable $y =1$ if the chip was accepted and $y = 0$ otherwise.
#
# The function $h_\theta(x)$ can be interpreted as the probability of the outcome being true rather than false. We want to find the parameters $\theta$ that maximize the log-likelihood over all (independently Bernoulli distributed) observations
#
# $$
# J(\theta) = \sum_{i, y_i = 1} \log h_\theta(x_i) + \sum_{i, y_i = 0} \log (1-h_\theta(x_i)).
# $$
# Consequently, we want to solve the following optimization problem:
# $$
# \text{minimize} \quad -J(\theta) + \mu \|\theta \|_2,
# $$
#
# where we added a regularization term with parameter $\mu$ to prevent overfitting.

# ## Feature mapping
# As our dataset only has two independent variables (the test scores) our model $y = \theta_0 + \theta_1 x_1 + \theta_2 x_2$ will have the form of a straight line. Looking at the plot one can see that a line will not perform well in separating the samples. Therefore, we will create more features based on each data point by mapping the original features ($x_1$, $x_2$) into all polynomial terms of $x_1$ and $x_2$ up to the 6th power:
#
# $$
# \text{map\_feature}(x_1,x_2) = [1, x_1, x_2, x_1^2, x_1x_2, x_2^2, x_1^3, \dots, x_1x_2^5, x_2^6 ]
# $$
#
# This will create 28 features for each sample.

#-

function map_feature(x1, x2)
  deg = 6
  x_new = ones(length(x1))
  for i = 1:deg, j = 0:i
      x_new = hcat(x_new, x1.^(i-j) .* x2.^j)
  end
  return x_new
end

X = map_feature(x1, x2);
size(X)

# ## Transformation into a conic optimisation problem
# We can rewrite above likelihood maximisation problem as a conic optimisation problem with exponential-cone-, second-order-cone-, equality-, and inequality constraints:
#
# $$
# \begin{array}{ll}
# \text{minimize}  &\sum_i^n \epsilon_i + \mu v\\
# \text{subject to}  & \|\theta \|_2 \leq v\\
# &  \log(1 + \exp(-\theta^\top x_i)) \leq \epsilon_i  \quad \text{if } y_i = 1, \\
# & \log(1 + \exp(\theta^\top x_i)) \leq \epsilon_i  \quad\text{   otherwise.}
# \end{array}
# $$
#
# Implementing the constraint $\log(1 + \exp(z)) \leq \epsilon$ for each of the $n$ samples requires two exponential cone constraints, one inequality constraint and two equality constraints. To see this, take the exponential on both sides and then divide by $\exp(\epsilon)$ to get:
#
# $$
# \exp(-\epsilon) + \exp(z - \epsilon) \leq 1.
# $$
#
# This constraint is equivalent to:
# $$
# \begin{array}{ll}
# (z - \epsilon, s_1, t_1) &\in K_{\text{exp}}, \\
# (-\epsilon, s_2, t_2) &\in K_{\text{exp}}, \\
# t_1 + t_2 &\leq 1,\\
# s_1 = s_2 &= 1,
# \end{array}
# $$
#
# where we defined the exponential cone as:
# $$
# K_{\text{exp}} = \{(r, s, t) \mid s >0, s \exp(r/s) \leq t \} \cup \{ r \leq 0, s = 0, t \geq 0 \}.
# $$
#
# Based on this transformation our optimisation problem will have $5n + n_\theta + 1$ variables, 1 SOCP constraint, $2n$ exponential cone
# constraints, $n$ inequality constraints and $2n$ equality constraints.
# Let's model the problem with JuMP and COSMO:

#-

n_theta = size(X, 2)
n = n_data
μ  = 1.

m = JuMP.Model(COSMO.Optimizer)
@variable(m, v)
@variable(m, θ[1:n_theta])
@variable(m, e[1:n])
@variable(m, t1[1:n])
@variable(m, t2[1:n])
@variable(m, s1[1:n])
@variable(m, s2[1:n])

@objective(m, Min, μ * v + sum(e))
@constraint(m, [v; θ] in MOI.SecondOrderCone(n_theta + 1))

## create the constraints for each sample points
for i = 1:n
  yi = y[i]
  x = X[i, :]
  yi == 1. ? (a = -1) : (a = 1)
  @constraint(m, [a * dot(θ, x) - e[i]; s1[i]; t1[i] ] in MOI.ExponentialCone())
  @constraint(m, [-e[i]; s2[i]; t2[i]] in MOI.ExponentialCone())
  @constraint(m, t1[i] + t2[i] <= 1)
  @constraint(m, s1[i] == 1)
  @constraint(m, s2[i] == 1)
end
JuMP.optimize!(m)
#-
theta = value.(θ)

# Once we have solved the optimisation problem and obtained the parameter vector $\theta$, we can plot the decision boundary. This can be done by evaluating our model over a grid of points $(u,v)$ and then plotting the contour line where the function $g$ returns a probability of $p=0.5$.

#-
## First we evaluate our model over a grid of points z = θ' x
u = collect(range(-1., stop = 1.5, length = 50))
v = collect(range(-1., stop = 1.5, length = 50))
z = zeros(length(u), length(v));
for i = 1:length(u), j = 1:length(v)
    z[i, j] = dot(map_feature(u[i], v[j]), theta);
end

# To add the decision boundary we have to plot the line indicating $50\%$ probability of acceptance, i.e. $g(\theta^\top x) = g(z)  = 0.5$ which we get at $z=0$.

#-
plot(x1[1:n_half-1], x2[1:n_half-1], color = :blue, st = :scatter, markershape = :cross, aspect_ratio=:equal, label = "Accepted", xlabel = "x1 - Microchip Test Score 1", ylabel = "x2 - Microchip Test Score 2")
plot!(x1[n_half:end], x2[n_half:end], color = :red, st = :scatter, markershape = :circle, label = "Rejected")
contour!(u, v, z', levels = [0.], c = :black, linewidth = 2)

# ## Solving the optimisation problem directly with COSMO
# We can solve the problem directly in COSMO by using its modeling interface. The problem will have $nn = 5 n + n_\theta + 1$ variables. Let us define the cost function  $\frac{1}{2}x^\top P x + q^\top x$:

#-
nn = 5 * n + n_theta +  1
P = spzeros(nn, nn)
q = zeros(nn)
q[1] = μ
for i = 1:n
  q[1 + n_theta + (i - 1) * 5 + 1] = 1.
end
# Next we define a function that creates the `COSMO.Constraints` for a given sample:

#-
## the order of the variables
## v, thetas, [e1 t11 t12 s11 s12] [e2 t21 t22 s21 s22] ...
## for each sample create two exponential cone constraints,
## 1 nonnegatives constraint, 2 zeroset constraints
function add_log_regression_constraints!(constraint_list, x, y, n, sample_num)
  num_thetas = length(x)
  ## 1st exponential cone constraint (zi - ei, s1, t1) in Kexp
  c_start = 1 + num_thetas + (sample_num - 1) * 5 + 1
  A = spzeros(3, n)
  A[1, c_start] = -1.
  y == 1. ? (a = -1) : (a = 1)
  for k = 1:num_thetas
    A[1, 2 + k - 1] = a * x[k]
  end
  A[2, c_start + 3] = 1.
  A[3, c_start + 1] = 1.
  b = zeros(3)
  push!(constraint_list, COSMO.Constraint(A, b, COSMO.ExponentialCone))

  ## 2nd exponential cone constraint (-e, s2, t2)
  A = spzeros(3, n)
  A[1, c_start] = -1.
  A[2, c_start + 4] = 1.
  A[3, c_start + 2] = 1.
  b = zeros(3)
  push!(constraint_list, COSMO.Constraint(A, b, COSMO.ExponentialCone))

  ## Nonnegatives constraint t1 + t2 <= 1
  A = spzeros(1, n)
  A[1, c_start + 1] = -1.
  A[1, c_start + 2] = -1.
  b = [1.]
  push!(constraint_list, COSMO.Constraint(A, b, COSMO.Nonnegatives))

  ## ZeroSet constraint s1 == 1, s2 == 1
  A = spzeros(2, n)
  A[1, c_start + 3] = 1.
  A[2, c_start + 4] = 1.
  b = -1 * ones(2)
  push!(constraint_list, COSMO.Constraint(A, b, COSMO.ZeroSet))
end
nothing

# Now we can use this function to loop over the sample points and add the constraints to our constraint list:

#-
constraint_list = Array{COSMO.Constraint{Float64}}(undef, 0)
for i = 1:n
  add_log_regression_constraints!(constraint_list, X[i, :], y[i], nn, i )
end

# It remains to add a second order cone constraint for the regularisation term:
# $\|\theta \|_2 \leq v$

#-
push!(constraint_list, COSMO.Constraint(Matrix(1.0I, n_theta + 1, n_theta + 1), zeros(n_theta + 1), COSMO.SecondOrderCone, nn, 1:n_theta+1));

# We can now create, assemble, and solve our `COSMO.Model`:

#-
model = COSMO.Model()
assemble!(model, P, q, constraint_list, settings = COSMO.Settings(verbose=true))
res = COSMO.optimize!(model);

# Let us double check that we get the same $\theta$ as in the previous section:

#-
using Test
theta_cosmo = res.x[2:2+n_theta-1]
@test norm(theta_cosmo - theta) < 1e-10


text:
using SnoopCompileCore

tinf = @snoopi_deep begin
    using MPSToolkit

    disableThreading()

    model = TransverseIsing1D(8)
    options = DMRGOptions()
    result = runDMRG(model, options)

    storeDMRGResult("dmrg.h5", result)

    model.parameters["hx"] = 1.5

    state = result.states[1]

    tebd_options = TEBDOptions(1, 2.0, 0.1, 10, 1e-12)
    result = runTEBD(state, model, tebd_options)
    storeTEBDResult("tebd1.h5", result)

    tebd_options = TEBDOptions(2, 2.0, 0.1, 10, 1e-12)
    result = runTEBD(state, model, tebd_options)
    storeTEBDResult("tebd2.h5", result)

    tebd_options = TEBDOptions(3, 2.0, 0.1, 10, 1e-12)
    result = runTEBD(state, model, tebd_options)
    storeTEBDResult("tebd3.h5", result)
end

using SnoopCompile

ttot, pcs = SnoopCompile.parcel(tinf)
SnoopCompile.write(joinpath(pwd(), "src", "Precompile"), pcs)


text:
immutable Point
x
y
end

max_stars_repo_path: Module 2/Chapter03/point.jl

max_issues_repo_path: Module 2/Chapter03/point.jl

max_forks_repo_path: Module 2/Chapter03/point.jl


text:
using PowerModels

abstract AbstractNFForm <: PowerModels.AbstractDCPForm

type StandardNFForm <: AbstractNFForm end
typealias NFPowerModel GenericPowerModel{StandardNFForm}

# default DC constructor
function NFPowerModel(data::Dict{AbstractString,Any}; kwargs...)
    return GenericPowerModel(data, StandardNFForm(); kwargs...)
end

function post_pfls{T}(pm::GenericPowerModel{T})

    PowerModels.variable_voltage(pm) # overloaded: phase angle variables for DC
    variable_generation(pm)
    variable_load_shed(pm) # TODO: add a PR to PowerModels; also overloaded
    if T != StandardNFForm
        PowerModels.variable_line_flow(pm)
    else
        variable_line_flow(pm)
    end

    objective_min_load_shed(pm) # TODO: add PR to PowerModels

    PowerModels.constraint_theta_ref(pm) # overloaded: does nothing for SOC
    PowerModels.constraint_voltage(pm) # overloaded: |W_ij|² ≦ Wᵢ⋅Wⱼ

    for (i, bus) in pm.ref[:bus] # pm.ref[:bus] replaces pm.set
        constraint_kcl_ls(pm, bus)
    end

    for (i, branch) in pm.ref[:branch]
        constraint_ohms_from(pm, branch)
        constraint_ohms_to(pm, branch)

        PowerModels.constraint_phase_angle_difference(pm, branch)
        PowerModels.constraint_thermal_limit_from(pm, branch)
        PowerModels.constraint_thermal_limit_to(pm, branch)
    end
end

function variable_generation{T}(pm::GenericPowerModel{T}; kwargs...)
    variable_active_generation(pm; kwargs...)
    variable_reactive_generation(pm; kwargs...) # does nothing for DC
end

function variable_active_generation{T}(pm::GenericPowerModel{T}; bounded = true)
    if bounded
        @variable(pm.model, 0 <= pg[i in keys(pm.ref[:gen])] <= pm.ref[:gen][i]["pmax"], start = PowerModels.getstart(pm.ref[:gen], i, "pg_start"))
    else
        @variable(pm.model, pg[i in keys(pm.ref[:gen])], start = PowerModels.getstart(pm.ref[:gen], i, "pg_start"))
    end
    return pg
end

function variable_reactive_generation{T <: PowerModels.AbstractWRForm}(pm::GenericPowerModel{T}; bounded = true)
    for i in keys(pm.ref[:gen])
        if (pm.ref[:gen][i]["qmin"] > 0 && pm.ref[:gen][i]["qmax"] > 0)
            pm.ref[:gen][i]["qmin"] = 0
        end
        if (pm.ref[:gen][i]["qmin"] < 0 && pm.ref[:gen][i]["qmax"] < 0)
            pm.ref[:gen][i]["qmax"] = 0
        end
    end
    if bounded
        @variable(pm.model, pm.ref[:gen][i]["qmin"] <= qg[i in keys(pm.ref[:gen])] <= pm.ref[:gen][i]["qmax"], start = PowerModels.getstart(pm.ref[:gen], i, "qg_start"))
    else
        @variable(pm.model, qg[i in keys(pm.ref[:gen])], start = PowerModels.getstart(pm.ref[:gen], i, "qg_start"))
    end
    return qg
end

function variable_reactive_generation{T <: PowerModels.AbstractDCPForm}(pm::GenericPowerModel{T})
    # do nothing for dc model - no reactive variables
end

function variable_reactive_generation{T <: PowerModels.AbstractACPForm}(pm::GenericPowerModel{T}; bounded = true)
    for i in keys(pm.ref[:gen])
        if (pm.ref[:gen][i]["qmin"] > 0 && pm.ref[:gen][i]["qmax"] > 0)
            pm.ref[:gen][i]["qmin"] = 0
        end
        if (pm.ref[:gen][i]["qmin"] < 0 && pm.ref[:gen][i]["qmax"] < 0)
            pm.ref[:gen][i]["qmax"] = 0
        end
    end
    if bounded
        @variable(pm.model, pm.ref[:gen][i]["qmin"] <= qg[i in keys(pm.ref[:gen])] <= pm.ref[:gen][i]["qmax"], start = PowerModels.getstart(pm.ref[:gen], i, "qg_start"))
    else
        @variable(pm.model, qg[i in keys(pm.ref[:gen])], start = PowerModels.getstart(pm.ref[:gen], i, "qg_start"))
    end
    return qg
end

function variable_load_shed{T <: PowerModels.AbstractDCPForm}(pm::GenericPowerModel{T})
    ld_min = Dict(i => 0.0 for i in keys(pm.ref[:bus]))
    ld_max = Dict(i => 0.0 for i in keys(pm.ref[:bus]))
    for i in keys(pm.ref[:bus])
        if (pm.ref[:bus][i]["pd"] > 0)
            ld_max[i] = 1.0
        end
    end
    @variable(pm.model, ld_min[i] <= ld[i in keys(pm.ref[:bus])] <= ld_max[i], start = 0.0)
    return ld
end

function variable_load_shed{T <: PowerModels.AbstractWRForm}(pm::GenericPowerModel{T})
    ld_min = Dict(i => 0.0 for i in keys(pm.ref[:bus]))
    ld_max = Dict(i => 0.0 for i in keys(pm.ref[:bus]))
    for i in keys(pm.ref[:bus])
        if (pm.ref[:bus][i]["pd"] > 0 || pm.ref[:bus][i]["qd"] != 0)
            ld_max[i] = 1.0
        end
    end
    @variable(pm.model, ld_min[i] <= ld[i in keys(pm.ref[:bus])] <= ld_max[i], start = 0.0)
    return ld
end

function variable_load_shed{T <: PowerModels.AbstractACPForm}(pm::GenericPowerModel{T})
    ld_min = Dict(i => 0.0 for i in keys(pm.ref[:bus]))
    ld_max = Dict(i => 0.0 for i in keys(pm.ref[:bus]))
    for i in keys(pm.ref[:bus])
        if (pm.ref[:bus][i]["pd"] > 0 || pm.ref[:bus][i]["qd"] != 0)
            ld_max[i] = 1.0
        end
    end
    @variable(pm.model, ld_min[i] <= ld[i in keys(pm.ref[:bus])] <= ld_max[i], start = 0.0)
    return ld
end

function variable_line_flow{T <: AbstractNFForm}(pm::GenericPowerModel{T}; bounded = true)
    if bounded
        @variable(pm.model, -pm.ref[:branch][l]["rate_a"] <= p[(l,i,j) in pm.ref[:arcs_from]] <= pm.ref[:branch][l]["rate_a"], start = PowerModels.getstart(pm.ref[:branch], l, "p_start"))
    else
        @variable(pm.model, p[(l,i,j) in pm.ref[:arcs_from]], start = PowerModels.getstart(pm.ref[:branch], l, "p_start"))
    end

    p_expr = Dict([((l,i,j), 1.0*p[(l,i,j)]) for (l,i,j) in pm.ref[:arcs_from]])
    p_expr = merge(p_expr, Dict([((l,j,i), -1.0*p[(l,i,j)]) for (l,i,j) in pm.ref[:arcs_from]]))

    pm.model.ext[:p_expr] = p_expr
end

# objective definition
function objective_min_load_shed{T}(pm::GenericPowerModel{T})
    c_pd = Dict(bp => 1 for bp in keys(pm.ref[:bus]))
    for (i, bus) in pm.ref[:bus]
        if (bus["pd"] < 0)
            c_pd[i] = 0
        end
    end
    ld = getvariable(pm.model, :ld)
    @objective(pm.model, Min, sum(pm.ref[:bus][i]["pd"] * c_pd[i] * ld[i] for i in keys(pm.ref[:bus])))
end

# constraint templates
function constraint_kcl_ls{T}(pm::GenericPowerModel{T}, bus)
    i = bus["index"]
    bus_arcs = pm.ref[:bus_arcs][i]
    bus_gens = pm.ref[:bus_gens][i]

    return constraint_kcl_ls(pm, i, bus_arcs, bus_gens, bus["pd"], bus["qd"])
end

function constraint_ohms_from{T}(pm::GenericPowerModel{T}, branch)
    i = branch["index"]
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    g, b = PowerModels.calc_branch_y(branch)
    return constraint_ohms_from(pm, f_bus, t_bus, f_idx, t_idx, g, b)
end

function constraint_ohms_to{T}(pm::GenericPowerModel{T}, branch)
    i = branch["index"]
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    g, b = PowerModels.calc_branch_y(branch)
    return constraint_ohms_to(pm, f_bus, t_bus, f_idx, t_idx, g, b)
end

# overloaded constraint functions
function constraint_kcl_ls{T <: PowerModels.AbstractDCPForm}(pm::GenericPowerModel{T}, i, bus_arcs, bus_gens, pd, qd)
    pg = getvariable(pm.model, :pg)
    p_expr = pm.model.ext[:p_expr]
    ld = getvariable(pm.model, :ld)

    c = @constraint(pm.model, sum(p_expr[a] for a in bus_arcs) == sum(pg[g] for g in bus_gens) - pd*(1-ld[i]))
    # omit reactive constraint
    return Set([c])
end

function constraint_kcl_ls{T <: PowerModels.AbstractWRForm}(pm::GenericPowerModel{T}, i, bus_arcs, bus_gens, pd, qd)
    w = getvariable(pm.model, :w)[i]
    p = getvariable(pm.model, :p)
    q = getvariable(pm.model, :q)
    pg = getvariable(pm.model, :pg)
    qg = getvariable(pm.model, :qg)
    ld = getvariable(pm.model, :ld)

    c1 = @constraint(pm.model, sum(p[a] for a in bus_arcs) == sum(pg[g] for g in bus_gens) - pd*(1-ld[i]))
    c2 = @constraint(pm.model, sum(q[a] for a in bus_arcs) == sum(qg[g] for g in bus_gens) - qd*(1-ld[i]))
    return Set([c1, c2])
end

function constraint_kcl_ls{T <: PowerModels.AbstractACPForm}(pm::GenericPowerModel{T}, i, bus_arcs, bus_gens, pd, qd)
    v = getvariable(pm.model, :v)[i]
    p = getvariable(pm.model, :p)
    q = getvariable(pm.model, :q)
    pg = getvariable(pm.model, :pg)
    qg = getvariable(pm.model, :qg)
    ld = getvariable(pm.model, :ld)

    c1 = @constraint(pm.model, sum(p[a] for a in bus_arcs) == sum(pg[g] for g in bus_gens) - pd*(1-ld[i]))
    c2 = @constraint(pm.model, sum(q[a] for a in bus_arcs) == sum(qg[g] for g in bus_gens) - qd*(1-ld[i]))
    return Set([c1, c2])
end

function constraint_ohms_from{T <: AbstractNFForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
     # Do nothing, this model is symmetric
     return Set()
end

function constraint_ohms_from{T <: PowerModels.AbstractDCPForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
    p_fr = getvariable(pm.model, :p)[f_idx]
    t_fr = getvariable(pm.model, :t)[f_bus]
    t_to = getvariable(pm.model, :t)[t_bus]

    c = @constraint(pm.model, p_fr == -b*(t_fr - t_to))
    return Set([c])
end

function constraint_ohms_to{T <: PowerModels.AbstractDCPForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
     # Do nothing, this model is symmetric
     return Set()
end

function constraint_ohms_from{T <: PowerModels.AbstractWRForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
    p_fr = getvariable(pm.model, :p)[f_idx]
    q_fr = getvariable(pm.model, :q)[f_idx]
    w_fr = getvariable(pm.model, :w)[f_bus]
    w_to = getvariable(pm.model, :w)[t_bus]
    wr = getvariable(pm.model, :wr)[(f_bus, t_bus)]
    wi = getvariable(pm.model, :wi)[(f_bus, t_bus)]

    c1 = @constraint(pm.model, p_fr == g*w_fr + (-g*wr) + (-b*wi) )
    c2 = @constraint(pm.model, q_fr == -b*w_fr - (-b*wr) + (-g*wi) )
    return Set([c1, c2])
end

function constraint_ohms_to{T <: PowerModels.AbstractWRForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
    q_to = getvariable(pm.model, :q)[t_idx]
    p_to = getvariable(pm.model, :p)[t_idx]
    w_fr = getvariable(pm.model, :w)[f_bus]
    w_to = getvariable(pm.model, :w)[t_bus]
    wr = getvariable(pm.model, :wr)[(f_bus, t_bus)]
    wi = getvariable(pm.model, :wi)[(f_bus, t_bus)]

    c1 = @constraint(pm.model, p_to == g*w_to + (-g*wr) + (b*wi) )
    c2 = @constraint(pm.model, q_to == -b*w_to - (-b*wr) + (g*wi) )
    return Set([c1, c2])
end

function constraint_ohms_from{T <: PowerModels.AbstractACPForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
    p_fr = getvariable(pm.model, :p)[f_idx]
    q_fr = getvariable(pm.model, :q)[f_idx]
    v_fr = getvariable(pm.model, :v)[f_bus]
    v_to = getvariable(pm.model, :v)[t_bus]
    t_fr = getvariable(pm.model, :t)[f_bus]
    t_to = getvariable(pm.model, :t)[t_bus]

    c1 = @NLconstraint(pm.model, p_fr == g*v_fr^2 + -g*v_fr*v_to*cos(t_fr-t_to) + -b*v_fr*v_to*sin(t_fr-t_to) )
    c2 = @NLconstraint(pm.model, q_fr == -b*v_fr^2 + b*v_fr*v_to*cos(t_fr-t_to) + -g*v_fr*v_to*sin(t_fr-t_to) )
    return Set([c1, c2])
end

function constraint_ohms_to{T <: PowerModels.AbstractACPForm}(pm::GenericPowerModel{T}, f_bus, t_bus, f_idx, t_idx, g, b)
    p_to = getvariable(pm.model, :p)[t_idx]
    q_to = getvariable(pm.model, :q)[t_idx]
    v_fr = getvariable(pm.model, :v)[f_bus]
    v_to = getvariable(pm.model, :v)[t_bus]
    t_fr = getvariable(pm.model, :t)[f_bus]
    t_to = getvariable(pm.model, :t)[t_bus]

    c1 = @NLconstraint(pm.model, p_to == g*v_to^2 + -g*v_fr*v_to*cos(t_to-t_fr) + -b*v_fr*v_to*sin(t_to-t_fr) )
    c2 = @NLconstraint(pm.model, q_to == -b*v_to^2 + b*v_fr*v_to*cos(t_to-t_fr) + -g*v_fr*v_to*sin(t_to-t_fr) )
    return Set([c1, c2])
end

max_stars_repo_name: kaarthiksundar/acopf-n-k

max_issues_repo_name: kaarthiksundar/acopf-n-k

max_forks_repo_name: kaarthiksundar/acopf-n-k


text:
module GraphicsManager

using FileManager
using FileIO.query

export IGraphicsReference
export IGraphicsData
export IGraphicsShaderProgram
export IGraphicsShader
export IGraphicsShaderProperty

abstract IGraphicsReference
abstract IGraphicsData <: IGraphicsReference
abstract IGraphicsShaderProgram <: IGraphicsReference
abstract IGraphicsShader <: IGraphicsReference
abstract IGraphicsShaderProperty <: IGraphicsReference

export AbstractGraphicsReference
export AbstractGraphicsData
export AbstractGraphicsShaderProgram
export AbstractGraphicsShader
export AbstractGraphicsShaderProperty

AbstractGraphicsReference = Union{Void,IGraphicsReference}
AbstractGraphicsData = Union{Void,IGraphicsData}
AbstractGraphicsShaderProgram = Union{Void,IGraphicsShaderProgram}
AbstractGraphicsShader = Union{Void,IGraphicsShader}
AbstractGraphicsShaderProperty = Union{Void,IGraphicsShaderProperty}

type Manager
	handle ::Any
	Manager() = new()
end

MANAGER = Manager()

setGraphicsDriver(o::Any) = (MANAGER.handle = o)
getGraphicsDriver() = MANAGER.handle

GRAPHICSDRIVER() = getGraphicsDriver()

export GRAPHICSDRIVER

end #GraphicsManager

max_stars_repo_name: Gilga/JGE

max_issues_repo_name: Gilga/JGE

max_forks_repo_name: Gilga/JGE


text:
using MultiAgentPathFinding
using Test

@testset "MultiAgentPathFinding.jl" begin
    # Write your tests here.
end


text:
"""
    skipmissing(itr::Raster)

Returns an iterable over the elements in a `Raster` object, skipping any values equal to either the `missingval` or `missing`.
"""
function Base.skipmissing(itr::Raster)
    if ismissing(missingval(itr))
        Base.SkipMissing(itr)
    else
        SkipMissingVal(itr)
    end
end

struct SkipMissingVal{T}
    x::T
end
Base.IteratorSize(::Type{<:SkipMissingVal}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{SkipMissingVal{T}}) where {T} = Base.IteratorEltype(T)
Base.eltype(::Type{SkipMissingVal{T}}) where {T} = Base.nonmissingtype(eltype(T))
missingval(itr::SkipMissingVal) = missingval(itr.x)

function Base.iterate(itr::SkipMissingVal, state...)
    y = iterate(itr.x, state...)
    y === nothing && return nothing
    item, state = y
    # We check for both `missing` and the raster `missingval`
    # Mostly the compiler should elide the `missing` check?
    while _missing(item, itr)
        y = iterate(itr.x, state)
        y === nothing && return nothing
        item, state = y
    end
    item, state
end

_missing(x, itr) = ismissing(x) || x == missingval(itr)  # mind the order, as comparison with missing returns missing
function _missing(x::AbstractFloat, itr)
    if isnan(missingval(itr))
        return ismissing(x) || isnan(x)
    else
        return ismissing(x) || x == missingval(itr)
    end
end

Base.IndexStyle(::Type{<:SkipMissingVal{T}}) where {T} = IndexStyle(T)
Base.eachindex(itr::SkipMissingVal) =
    Iterators.filter(i -> !_missing(@inbounds(itr.x[i]), itr), eachindex(itr.x))
Base.keys(itr::SkipMissingVal) =
    Iterators.filter(i -> !_missing(@inbounds(itr.x[i]), itr), keys(itr.x))
@propagate_inbounds function Base.getindex(itr::SkipMissingVal, I...)
    v = itr.x[I...]
    _missing(v, itr) && throw(MissingException("the value at index $I is the raster missingval"))
    v
end

function Base.show(io::IO, s::SkipMissingVal)
    print(io, "skipmissing(")
    show(io, s.x)
    print(io, ')')
end


text:
function markdown_str(q::RBQuote)
    split_ = split(q.txt, "\n")
    return string(">", join(split_, "\n>"))
end

function markdown_str(n::RBNote)
    return n.txt
end

function markdown_str(p::RBPair)
    return string("- ", p.key, ": ", p.val)
end

function markdown_str(t::RBTagLine)
    return string("***", t.tag, "***")
end

function markdown_str(i::RBObject)
    return string("***'", nameof(i), "' not supported", "***")
end

function create_markdown(doc::RBDoc, outfile)
    lines = String[]

    # top level
    push!(lines, string("# ", dockey(doc)))
    push!(lines, "")

    for (seckey, sec) in sections(doc)
        push!(lines, string("## ", seckey))
        push!(lines, "")

        for item in items(sec)
            push!(lines, markdown_str(item))
            push!(lines, "")
        end

    end

    # write
    open(outfile, "w") do io
        for line in lines
            println(io, line)
        end
    end
end

# TODO: Add expand links
macro create_doc_readme()
    __file__ = string(__source__.file)
    outfile = joinpath(dirname(__file__), "README.md")
    create_markdown(currdoc(), outfile)
    :(nothing)
end


text:
#
@testset "FReps" begin
    N = 10
    xy  = rand(2,N)
    xyz = rand(3,N)

    x = rand(1,N)
    y = rand(1,N)
    z = rand(1,N)
    #----------------------------------

    p  = Circle(1.0)

    @test FRep(p,xy[:,1]) |> size == ()
    @test FRep(p,xy)      |> size == (1,N)
    @test FRep(p,x,y)     |> size == (1,N)
    #----------------------------------

    p = Square([1.0,1.0],center=true)

    @test FRep(p,xy[:,1]) |> size == ()
    @test FRep(p,xy)      |> size == (1,N)
    @test FRep(p,x,y)     |> size == (1,N)
    #----------------------------------

    p = Sphere(1.0)

    @test FRep(p,xyz[:,1]) |> size == ()
    @test FRep(p,xyz)      |> size == (1,N)
    @test FRep(p,x,y,z)    |> size == (1,N)
    #----------------------------------

    p = Cuboid([1.0,1.0,1.0],center=true)

    @test FRep(p,xyz[:,1]) |> size == ()
    @test FRep(p,xyz)      |> size == (1,N)
    @test FRep(p,x,y,z)    |> size == (1,N)
    #----------------------------------

    p = Cylinder(1.0,1.0,1.0)

    @test FRep(p,xyz[:,1]) |> size == ()
    @test FRep(p,xyz)      |> size == (1,N)
    @test FRep(p,x,y,z)    |> size == (1,N)
    #----------------------------------

    p = LinearExtrude(Circle(1.0),2.0)

    @test FRep(p,xyz[:,1]) |> size == ()
    @test FRep(p,xyz)      |> size == (1,N)
    @test FRep(p,x,y,z)    |> size == (1,N)
    #----------------------------------
    c = translate([1.0,0.0])Circle(1.0)
    s = translate([0.0,1.0])Square([3.0,3.0];center=true)
    p = diff(s,c)

    @test FRep(p,xy[:,1]) |> size == ()
    @test FRep(p,xy)      |> size == (1,N)
    @test FRep(p,x,y)     |> size == (1,N)
    #----------------------------------
    p = Piping(1.0, [[0,0,0],[10,0,0],[10,10,0],[10,10,10],[5,5,5]])

    @test FRep(p,xyz[:,1]) |> size == ()
    @test FRep(p,xyz)      |> size == (1,N)
    @test FRep(p,x,y,z)    |> size == (1,N)
    #----------------------------------
end


text:
# WRITTEN IN THE REPL
# ] add IJulia InstantiateFromURL # Enter package
# mode and install link to Jupyter and manager for
# QE dependencies.

# QuantEcon packages
# using InstantiateFromURL
# activate_github("QuantEcon/QuantEconLectureAllPackages", tag = "v0.9.6", add_default_environment = true)

# Status and Gadfly plotting package
# ]
#status
# add Gadfly


text:
using Plexifications
using Base.Test: @test, @test_throws

@test begin
    iszerodivisor(0)
end

@test begin
    !iszerodivisor(3.14)
end

@test_throws ErrorException begin
    inv(Cmplex(0, 0.0))
end

@test begin
    iszerodivisor(Cmplex(0,0))
end

@test begin
    !iszerodivisor(Cmplex(0,1))
end

@test begin
    length(unreal(1)) == 0
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(Cmplex(a))
end

@test begin
    a = 1
    !isreal(Cmplex(a, a))
end

@test begin
    z = random(Cmplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(Cmplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, Cmplex(1, 2))
    l = String(take!(io))
    r = "[1: 1, H: 2]"
    l == r
end

@test begin
    x = random(Cmplex{BigFloat})
    y = random(Cmplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(Cmplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(Cmplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    z = random(Cmplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    z = random(Cmplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    barH(barH(x)) == x
end

@test begin
    x = random(Cmplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = barH(x * y)
    r = barH(y) * barH(x)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = barH(x + y)
    r = barH(x) + barH(y)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = barH(x - y)
    r = barH(x) - barH(y)
    l == r
end

@test begin
    x = random(Cmplex{Rational{BigInt}})
    y = random(Cmplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Cmplex{Rational{BigInt}})
    y = random(Cmplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Cmplex{Rational{BigInt}})
    y = random(Cmplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    z = random(Cmplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(Cmplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    z = random(Cmplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    z = random(Cmplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    z = random(Cmplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    abs2(x) > big(0)
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(Cmplex{BigInt})
    l = abs2(x * y)
    r = abs2(x) * abs2(y)
    l == r
end

@test_throws ErrorException begin
    inv(Cmplex(0, 0.0))
end

@test_throws ErrorException begin
    random(Cmplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Cmplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(Cmplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(Cmplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Cmplex{BigInt})
    z == z + zero(Cmplex{BigInt})
end

@test begin
    z = random(Cmplex{BigInt})
    z == z * one(Cmplex{BigInt})
end

@test begin
    z = random(Cmplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Cmplex{BigInt})
    z == z * one(z)
end


text:
module ThreadSafeDicts

import Base.getindex, Base.setindex!, Base.get!, Base.get, Base.empty!, Base.pop!
import Base.haskey, Base.delete!, Base.print, Base.iterate, Base.length

export ThreadSafeDict

"""
    ThreadSafeDict(pairs::Vector{Pair{K,V}})
Struct and constructor for ThreadSafeDict. There is one lock per Dict struct. All functions lock this lock, pass
arguments to the d member Dict, unlock the spinlock, and then return what is returned by the Dict.
"""
struct ThreadSafeDict{K, V} <: AbstractDict{K, V}
    dlock::Threads.SpinLock
    d::Dict
    ThreadSafeDict{K, V}() where V where K = new(Threads.SpinLock(), Dict{K, V}())
    ThreadSafeDict{K, V}(itr) where V where K = new(Threads.SpinLock(), Dict{K, V}(itr))
end
ThreadSafeDict() = ThreadSafeDict{Any,Any}()

function getindex(dic::ThreadSafeDict, k)
    lock(dic.dlock)
    v = getindex(dic.d, k)
    unlock(dic.dlock)
    return v
end

function setindex!(dic::ThreadSafeDict, k, v)
    lock(dic.dlock)
    h = setindex!(dic.d, k, v)
    unlock(dic.dlock)
    return h
end

function haskey(dic::ThreadSafeDict, k)
    lock(dic.dlock)
    b = haskey(dic.d, k)
    unlock(dic.dlock)
    return b
end

function get(dic::ThreadSafeDict, k, v)
    lock(dic.dlock)
    v = get(dic.d, k, v)
    unlock(dic.dlock)
    return v
end

function get!(dic::ThreadSafeDict, k, v)
    lock(dic.dlock)
    v = get!(dic.d, k, v)
    unlock(dic.dlock)
    return v
end

function pop!(dic::ThreadSafeDict)
    lock(dic.dlock)
    p = pop!(dic.d)
    unlock(dic.dlock)
    return p
end

function empty!(dic::ThreadSafeDict)
    lock(dic.dlock)
    d = empty!(dic.d)
    unlock(dic.dlock)
    return d
end

function delete!(dic::ThreadSafeDict, k)
    lock(dic.dlock)
    p = delete!(dic.d, k)
    unlock(dic.dlock)
    return p
end

function length(dic::ThreadSafeDict)
    lock(dic.dlock)
    len = length(dic.d)
    unlock(dic.dlock)
    return len
end

function iterate(dic::ThreadSafeDict)
    lock(dic.dlock)
    p = iterate(dic.d)
    unlock(dic.dlock)
    return p
end

function iterate(dic::ThreadSafeDict, i)
    lock(dic.dlock)
    p = iterate(dic.d, i)
    unlock(dic.dlock)
    return p
end

function print(io::IO, dic::ThreadSafeDict)
    print(io, "Dict was ", islocked(dic.dlock) ? "locked" : "unlocked", ", contents: ")
    lock(dic.dlock)
    print(io, dic.d)
    unlock(dic.dlock)
end

end # module


text:
using SequentialMonteCarlo
using RNGPool
import SMCExamples.LinearGaussian.defaultLGModel
import SMCExamples.Particles.Float64Particle
import Statistics: mean, var

setRNGs(0)

model, theta, ys, ko = defaultLGModel(100)

smcio = SMCIO{model.particle, model.pScratch}(1024*1024, model.maxn,
  Threads.nthreads(), false)
smc!(model, smcio)

predMeanSMC = SequentialMonteCarlo.eta(smcio, p -> p.x, false, smcio.n)
predVarSMC = SequentialMonteCarlo.eta(smcio, p -> p.x^2, false, smcio.n) -
  predMeanSMC^2
println("Predictive mean: $(ko.predictionMeans[smcio.n])")
println("Estimate: $predMeanSMC")
println("Predictive variance: $(ko.predictionVariances[smcio.n])")
println("Estimate: $predVarSMC")

filtMeanSMC = SequentialMonteCarlo.eta(smcio, p -> p.x, true, smcio.n)
filtVarSMC = SequentialMonteCarlo.eta(smcio, p -> p.x^2, true, smcio.n) -
  filtMeanSMC^2
println("Filtering mean: $(ko.filteringMeans[smcio.n])")
println("Estimate: $filtMeanSMC")
println("Filtering variance: $(ko.filteringVariances[smcio.n])")
println("Estimate: $filtVarSMC")

println("Running many particle filters with only 128 particles in parallel...")
m = 10000
lZs = Vector{Float64}(undef, m)
Vs = Vector{Float64}(undef, m)
lZhats = Vector{Float64}(undef, m)
Vhats = Vector{Float64}(undef, m)
f1(p::Float64Particle) = 1.0

nthreads = Threads.nthreads()
smcios = Vector{SMCIO}(undef, nthreads)
Threads.@threads for i = 1:nthreads
  smcios[i] = SMCIO{model.particle, model.pScratch}(128, model.maxn, 1, true)
end
@time Threads.@threads for i = 1:m
  smcio = smcios[Threads.threadid()]
  smc!(model, smcio)
  lZs[i] = smcio.logZhats[smcio.n-1]
  Vs[i] = SequentialMonteCarlo.V(smcio, f1, false, false, smcio.n)
  lZhats[i] = smcio.logZhats[smcio.n]
  Vhats[i] = SequentialMonteCarlo.V(smcio, f1, true, false, smcio.n)
end

print("Empirical relative variance of Z_n^N: ")
println(var(exp.(lZs.-ko.logZhats[smcio.n-1])))
print("Mean of unbiased estimator of relative variance of Z_n^N: ")
println(mean(Vs .* exp.(2*(lZs.-ko.logZhats[smcio.n-1]))))
print("Estimated standard deviation of the value above: ")
println(sqrt(var(Vs .* exp.(2*(lZs.-ko.logZhats[smcio.n-1])))/m))

print("Empirical relative variance of \\hat{Z}_n^N: ")
println(var(exp.(lZhats.-ko.logZhats[smcio.n])))
print("Mean of unbiased estimator of relative variance of \\hat{Z}_n^N: ")
println(mean(Vhats .* exp.(2*(lZhats.-ko.logZhats[smcio.n]))))
print("Estimated standard deviation of the value above: ")
println(sqrt(var(Vhats .* exp.(2*(lZhats.-ko.logZhats[smcio.n])))/m))


text:
printstyled("  env!\n", color=:light_green)

# GphysChannel, gaps
printstyled("    on SeisChannel\n", color=:light_green)
C = randSeisChannel(s=true)
C.t = [ 1 0;
        2 500;
        1959 69420;
        90250 0]
C.x = randn(C.t[end,1])
env!(C)

# GphysChannel, no gaps
C = randSeisChannel(s=true)
ungap!(C)
D = env(C)

# GphysData
printstyled("    on SeisData\n", color=:light_green)
S = randSeisData(24)
ungap!(S)
U = deepcopy(S)
env!(S, v=2)
T = env(U, v=2)
printstyled("    testing that env! == DSP.hilbert\n", color=:light_green)
for i = 1:S.n
  if S.fs[i] > 0.0 && length(S.x[i]) > 256
    ex1 = S.x[i]
    ex2 = abs.(DSP.hilbert(U.x[i]))
    @test isapprox(ex1, ex2)
    @test isapprox(S.x[i], T.x[i])
  end
end

for k = 1:4
  S = randSeisData(24, s=1.0)
  deleteat!(S, findall(S.fs.<20.0))
  for i = 1:S.n
    nx = 2^16
    t = [1 0; 512 134235131; 2^14 100000; 2^15 12345678; nx 0]
    S.t[i] = t
    S.x[i] = randn(eltype(S.x[i]), nx)
  end
  U = deepcopy(S)
  env!(S)

  for i = 1:S.n
    j = rand(1:3)
    si = S.t[i][j,1]
    ei = S.t[i][j+1,1]-1
    ex1 = S.x[i][si:ei]
    ex2 = abs.(DSP.hilbert(U.x[i][si:ei]))
    @test isapprox(ex1, ex2)
  end

end

printstyled("    testing on very short data segments\n", color=:light_green)
S = randSeisData(5)
N = 2 .^ collect(1:1:5)
for i = 1:S.n
  nx = N[i]
  S.fs[i] = 100.0
  S.t[i] = [1 0; nx 0]
  S.x[i] = rand(nx)
end
env!(S)


text:
function SolveModel!(MS::ModelSolution)
  LastPeriod!(MS)
  if MS.np.p1cons == true
    FirstPeriod!(MS)
  else
    FirstPeriod_nocons!(MS)
  end

  return MS

end

function LastPeriod!(MS::ModelSolution)
  np = MS.np
  mp = MS.mp
  ia = np.na
  y = np.ygrd[ia]
  vtmp = fill(-Inf64,np.nh)
  ctmp = fill(NaN64,np.nh)
  for (ix,xv) in enumerate(np.xgrd)
    for ih in eachindex(np.hgrd)
    b = 0.0
    s = 0.0
    for (ihn,hv) in enumerate(np.hgrd) # Find choices
      ms = SalesCost(ih,ihn,mp.ms)
      xtmp = xv-ms
      ctmp[ihn] = c = findc2(xtmp,y,b,s,mp.q,hv)
      if c > 0
        vtmp[ihn] = utilh(c,hv,mp.γ,mp.η) + KKKShifter(mp.xstar,xv,mp.ψ)
      end
    end
    imax = argmax(vtmp)
    MS.h[ix,ih,ia] = np.hgrd[imax]
    MS.c[ix,ih,ia] = ctmp[imax]
    MS.V[ix,ih,ia] = vtmp[imax]
    end
  end
end

function FirstPeriod!(MS::ModelSolution)
  np = MS.np
  mp = MS.mp

  savgrd = range(0.0,stop=np.xgrd[end],length=np.nsav)
  αgrd = range(0.0,stop=1.0,length=np.nα)
  vtmp = fill(0.0,(np.nsav,np.nα))

  ia = 1
  y = np.ygrd[ia]
  for ih in eachindex(np.hgrd)
    Vnxt_intrp = CubicSplineInterpolation((np.xgrd,),MS.V[:,ih,ia+1],extrapolation_bc=Interpolations.Line())
  for (ix,xv) in enumerate(np.xgrd)
    vtmp .= -Inf64
    for (isav,sav) in enumerate(savgrd)
      for (iα,α) in enumerate(αgrd)
        b = sav*(1.0-α)
        s = sav*α
        c = findc1(xv,y,b,s,mp.q)
        if c > 0
          vtmp[isav,iα] = util(c,mp.γ)
          for irs = 1:np.nq
            vtmp[isav,iα] += mp.β*np.πrs[irs]*Vnxt_intrp(xp(b,s,mp.r,np.rsgrd[irs]))
          end
        end
      end
    end

    imax = argmax(vtmp)
    isav = imax[1]
    iα = imax[2]
    MS.α[ix,ih] = α = αgrd[iα]
    MS.s[ix,ih] = s = savgrd[isav]*α
    MS.b[ix,ih] = b = savgrd[isav]*(1.0-α)
    if savgrd[isav] == 0.0 # If there is no saving, portfolio weight is'nt defined
      MS.α[ix,ih] = NaN64
    end
    MS.c[ix,ih,ia] = findc1(xv,y,s,b,mp.q)
    MS.V[ix,ih,ia] = vtmp[imax]
  end
  end

end

function FirstPeriod_nocons!(MS::ModelSolution)
  np = MS.np
  mp = MS.mp

  αgrd = range(0.0,stop=1.0,length=np.nα)
  vtmp = fill(0.0,(np.nα))

  ia = 1
  y = np.ygrd[ia]
  for ih in eachindex(np.hgrd)
    Vnxt_intrp = CubicSplineInterpolation((np.xgrd,),MS.V[:,ih,ia+1],extrapolation_bc=Interpolations.Line())
  for (ix,xv) in enumerate(np.xgrd)
    vtmp .= -Inf64
    saving = xv + y
    for (iα,α) in enumerate(αgrd)
      sav = saving - mp.q*part(α)
      b = sav*(1.0-α)
      s = sav*α
      if sav > 0
        vtmp[iα] = 0.0
        for irs = 1:np.nq
          xpn = xp(b,s,mp.r,np.rsgrd[irs])
          vtmp[iα] += mp.β*np.πrs[irs]*(Vnxt_intrp(xpn))
        end
      end
    end

    iα = argmax(vtmp)
    MS.α[ix,ih] = α = αgrd[iα]
    sav = saving - mp.q*part(α)
    MS.s[ix,ih] = s = sav*α
    MS.b[ix,ih] = b = sav*(1.0-α)
    if sav == 0.0 # If there is no saving, portfolio weight is'nt defined
      MS.α[ix,ih] = NaN64
    end
    MS.c[ix,ih,ia] = findc1(xv,y,s,b,mp.q)
    MS.V[ix,ih,ia] = vtmp[iα]
  end
  end

end

" Find consumption as a function of choices"
function findc2(x,y,b,s,q,h) # Find second period consumption (subtract housing)
  c = findc1(x,y,b,s,q) - h
end

function findc1(x,y,b,s,q) # Find todays consumption
  c = x + y - b - s - q*part(s)
end

"Takes portfolio choices and returns possible next period wealth"
function xp(b,s,r,rs::Float64)
  b*(1.0+r) .+ s*(1.0 + rs)
end

"Returns true if household participates"
function part(s::AbstractFloat) # Whether you are participating or not
  if s>0.0
      part = true
  else
      part = false
  end

  return part
end

"Utility shifter as in Karlman, Kinnerud, Kragh-Soerensen (2021)"
function KKKShifter(xstar::AbstractFloat, x::AbstractFloat,ψ::AbstractFloat)
  ψ*(x>xstar)
end

function SalesCost(ih::Integer,ihn::Integer,κ::AbstractFloat)
  if ihn != ih
    cost = κ
  else
    cost = 0.0
  end

  return cost
end


text:
module test_julia_multiple_dispatch

using Test

f(::Int) = 2
f(::String) = "s"

@test length(methods(f)) == 2

@test Base.return_types(f, (Int,)) == Any[Int]
@test Base.return_types(f, (String,)) == Any[String]

m = first(methods(f, [String]))
Base.delete_method(m)

@test length(methods(f)) == 1

end # module test_julia_multiple_dispatch

module test_julia_multiple_dispatch_dc

using Test
using IOCapture

# http://gall.dcinside.com/mgallery/board/view/?id=github&no=1628

abstract type 사람 end

struct 학생 <: 사람
    name::String
    age::UInt8
end

struct 손님 <: 사람
    name::String
    age::UInt8
end

struct 예비군 <: 사람
    name::String
    age::UInt8
end

function study(나::사람)
    println(나.name, "(", 나.age, ")는 ", nameof(typeof(나)), "으로서 공부를 합니다.")
end

function buy_goods(나::사람)
    println(나.name, "(", 나.age, ")는 ", nameof(typeof(나)), "으로서 물건을 삽니다.")
end

function train(나::사람)
    println(나.name, "(", 나.age, ")는 ", nameof(typeof(나)), "으로서 훈련을 합니다.")
end

(value, output) = IOCapture.capture() do
    name, age = "다메즈마", 27
    buy_goods(손님(name, age))
    study(학생(name, age))
    train(예비군(name, age))

    buy_goods(학생(name, age))
end

@test output == """
다메즈마(27)는 손님으로서 물건을 삽니다.
다메즈마(27)는 학생으로서 공부를 합니다.
다메즈마(27)는 예비군으로서 훈련을 합니다.
다메즈마(27)는 학생으로서 물건을 삽니다.
"""

end # module test_julia_multiple_dispatch_dc