#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


module GDB

using Distributed

# @everywhere begin
using DistributedArrays
using SharedArrays
using LinearAlgebra
using IterativeSolvers
using SparseArrays
using HDF5
using ForwardDiff
using RecipesBase
using Logging
using ProgressMeter

include("./domain/grid.jl")
include("./domain/barrier.jl")
include("./domain/interpolation.jl")
include("./domain/operators.jl")
include("./domain/ghosts.jl")

include("./linsolve/gmres.jl")
include("./linsolve/relaxation.jl")
include("./linsolve/solvers.jl")

include("./magflux/trace.jl")
include("./magflux/maps.jl")
include("./magflux/annalus.jl")
include("./magflux/solovev.jl")

include("./physics/formulas.jl")
include("./physics/timestep.jl")

# end

end


convertStringToVisualMaterial(visualMaterial::Shapes.VisualMaterial) = visualMaterial

convertStringToVisualMaterial(visualMaterial::AbstractString) = Shapes.visualMaterialPalette[visualMaterial]

"""
    Visual(;
        shape = Sphere(diameter = 0.5),
        visualMaterial = VisualMaterial(color = "Green"))

Generate a new visual shape with a visual material. It is only for visualizing, without physical behavior. In this case it is a greenish sphere.

# Arguments
- `shape`: it is possible to use all [Shapes](@ref) from both sets visual and solid, e.g. [`Sphere`](@ref), [`Box`](@ref), [`CoordinateSystem`](@ref). The default value is `nothing`, this means it will not be visualized.
- `visualMaterial`: use a predefined [Visual material](@ref) defined in palettes/visualMaterials.json, or define it with `VisualMaterial(color = "Blue")`. The default value is `VisualMaterial()`.
"""
struct Visual <: Modia3D.AbstractObject3DFeature
    shape::Union{Modia3D.AbstractShape,Nothing}
    visualMaterial::Shapes.VisualMaterial

    function Visual(;
        shape::Union{Modia3D.AbstractShape,Nothing} = nothing,
        visualMaterial::Union{Shapes.VisualMaterial,AbstractString,Nothing} = Shapes.VisualMaterial())
        new(shape, convertStringToVisualMaterial(visualMaterial) )
    end
end


module RandomBattles

if isdefined(Base, :Experimental) && isdefined(Base.Experimental,
      Symbol("@optlevel"))
      @eval Base.Experimental.@optlevel 3
end

# grab data and Pokemon structs so that it can be used
include("data.jl")
include("stats.jl")
include("moves.jl")

# team battles
include("team/pokemon.jl")
include("team/team.jl")
include("team/state.jl")
include("team/mechanics.jl")
include("team/decisions.jl")
include("team/battle_logic.jl")
include("team/diff.jl")
include("team/strategy.jl")
include("team/nash.jl")
include("team/mcts.jl")
#include("team/alpha_beta_pruning.jl")

# individual battles
include("individual/state.jl")
include("individual/mechanics.jl")
include("individual/decisions.jl")
include("individual/battle_logic.jl")
#include("individual/diff.jl")
include("individual/strategy.jl")
include("individual/nash.jl")

# higher level abstractions
include("pokemon_meta.jl")
include("rankings.jl")

# export useful variables and functions
export DynamicState, StaticState, DynamicTeam, StaticTeam, DynamicPokemon,
      StaticPokemon, DynamicIndividualState, StaticIndividualState,
      get_possible_decisions, battle_score, battle_scores, play_turn,
      play_battle, gamemaster,
      diff, reflect,
      Strategy, IndividualStrategy, plot_strategy, rank,
      PokemonMeta,
      SM, solve_battle,
      MCTSNode, update_tree_MCTS!, select_decisions_MCTS

end # module


# https://adventofcode.com/2021/day/8
using AdventOfCode
using Combinatorics
input = split.(readlines("2021/data/day_8.txt"), " | ") .|> (x->split.(x," "))
nseg = Dict((0:9) .=> (6,2,5,5,4,5,6,3,7,6))

defsegs = ("abcefg","cf","acdeg","acdfg","bcdf","abdfg","abdefg","acf","abcdefg","abcdfg")
defsegdict=Dict(Set.(collect.(defsegs)).=> 0:9)

input = (patterns=first.(input), outputs = last.(input), nseg=nseg, defsegdict=defsegdict)
function iteratedict(v::AbstractVector)
    (Dict(v.=> permv) for permv in permutations(v))
end
vec2nseg(v,input)=getindex.(Ref(input.nseg),v)
translateset(transdict, set)=getindex.(Ref(transdict),set)|>Set

function part_1(input)
    outputnsegs=(length.(x) for x in input.outputs) # vector whose elements [w, x, y, z] are 4 segment numbers of each output
    sum(sum(Base.Fix2(in, getindex.(Ref(input.nseg),(1,4,7,8))) ,x) for x in outputnsegs)
end

function part_2(input)
    sum(begin
        matchdict=iteratedict(collect("abcdefg"))|>collect # all possible translations
        translator=filter(dict->translateset.(Ref(dict),
            Set.([collect.(pattern);collect.(output)])) .|>
            Base.Fix2(in, keys(input.defsegdict)) |> all,matchdict) |>first # filter those making sense
        translateset.(Ref(translator), Set.(collect.(output))) .|>
            Base.Fix1(getindex,input.defsegdict) |> join |> parse # translate output
    end
        for (pattern,output) in zip(input.patterns, input.outputs))
end


print("Loading packages... ")
using Images, FileIO, StatsBase, ImageCore
println("Done!")
###############################################################################
# DATA DEFINITIONS                                                            #
# {red|green|blue}map = current value of each pixel                           #
# {red|green|blue}counts = pixels in each bin                                 #
# {red|green|blue}vars: red[1000] = what 1000 mapped to with this transform   #
###############################################################################
########################## defaults ##############################################

obase = "a-stretched-image"  # default output base file name

rootpower    = 6.0      # default power factor: 1/rootpower for all iterations
rootiter     = 1        # number of iterations on applying rootpower - sky
rootpower2   = 1.0      # if rootiter =2 use this power on iteration 2 if the user sets it >1
pcntclip     = 0.005    # default percent clip level = total pixels * pcntclip/100

colorcorrect = 1        # default is to docolor correction.  (turn off to see the difference)

colorenhance = 1.0      # default enhancement value.
tonecurve    = 0        # no application odf a tone curve
specprhist   = 0        # no output histogram to specpr file
cumstretch   = 0        # do not do cumulative histogram stretch after rootpower
skylevelfactor = 0.06   # sky level relative to the histogram peak  (was 0.03 v0.88 and before)
rgbskylevel   = 1024.0    # desired  on a 16-bit  0 to 65535 scale
rgbskylevelrs = 4096.0    # desired  on output root stretched image 16-bit  0 to 65535 scale
zeroskyred    = 4096.0    # desired zero point on sky, red   channel
zeroskygreen  = 4096.0    # desired zero point on sky, green channel
zeroskyblue   = 4096.0    # desired zero point on sky, blue  channel
scurve        = 1         # no s-curve application
setmin        = 1         # no modification of the minimum
setminr = 0.0  # minimum for red
setming = 0.0  # minimum for green
setminb = 0.0  # minimum for blue
idisplay = 0   # do not display the final image
jpegonly = 0   # do not do jpeg only (will do jpeg + 16-bit png)
saveinputminussky  = 0  #  save input image - sky

debuga  = 0  # set to 1 for debugging, or 0 for none
doplots = 0  # show plots of histograms

if doplots != 0
  print("Loading Plots package... ")
  using Plots
  pyplot()
  println("Done!")
else # give plot() a definition
  function plot(args...)
  end
end

redvals = 0
greenvals = 0
bluevals = 0 # to put these in the global namespace

function plotter(title)
  #=edges = (1:256) .* 256
  r = map(x -> log(x + 1), fit(Histogram, redcounts, edges).weights)
  g = map(x -> log(x + 1), fit(Histogram, greencounts, edges).weights)
  b = map(x -> log(x + 1), fit(Histogram, bluecounts, edges).weights)
  print(size(r))
  for i = 1:255
    println(i, ": ", r[i], ", ", g[i], " ", b[i])
  end
#  a = cat(2, r, g, b)
#  plot(a,legend = nothing, color=reshape([:red, :green, :blue], 1, 3), title = title)
#  gui()=#
end

function checksum(v)
  result = reduce($, 0, map(Int,map(floor, v)))
  return result
end

function remap()
  # This function needs to be optimized!
  # if redmap[1000] used to be 987, it becomes where redvals[987] is going
  println("Remapping.... ")
  global redcounts, redmap, greencounts, greenmap, bluecounts, bluemap
  global redvals, greenvals, bluevals
  redvals = map(x -> floor(UInt64, x), redvals)
  greenvals = map(x -> floor(UInt64, x), greenvals)
  bluevals = map(x -> floor(UInt64, x), bluevals)
  redmap = map(x -> redvals[x] + 1, redmap)
  greenmap = map(x -> greenvals[x] + 1, greenmap)
  bluemap = map(x -> bluevals[x] + 1, bluemap)
  println("recounting... ")
  newred = zeros(65536)
  newgreen = zeros(65536)
  newblue = zeros(65536)

  for i = 1:65536
    newred[redvals[i] + 1] += redcounts[i]
    newgreen[greenvals[i] + 1] += greencounts[i]
    newblue[bluevals[i] + 1] += bluecounts[i]
  end

  redcounts = newred
  greencounts = newgreen
  bluecounts = newblue
  println("Done!")
end

function showmeans()
  v = collect(0:65535)

  redmin = reduce(min, npix, v[redcounts .>0])
  greenmin = reduce(min, npix, v[greencounts .>0])
  bluemin = reduce(min, npix, v[bluecounts .>0])

  redmax = reduce(max, 0, v[redcounts .>0])
  greenmax = reduce(max, 0, v[greencounts .>0])
  bluemax = reduce(max, 0, v[bluecounts .>0])

  redmean = sum((0:65535).* redcounts) / npix
  greenmean = sum((0:65535) .* greencounts) / npix
  bluemean = sum((0:65535) .* bluecounts) / npix

  @printf "  red: min = %5d max = %5d mean = %7.1f\n" redmin redmax redmean
  @printf "green: min = %5d max = %5d mean = %7.1f\n" greenmin greenmax greenmean
  @printf " blue: min = %5d max = %5d mean = %7.1f\n" bluemin bluemax bluemean
end

function smoothsky()
  global redvals, greenvals, bluevals
  # histogram plotting disabled
  # now find the sky level on the left side of the histogram

  println("Calculating sky smoothing peaks")
  redacc = cumsum(redcounts)
  greenacc = cumsum(greencounts)
  blueacc = cumsum(bluecounts)
  smoothred = zeros(65536)
  smoothgreen = zeros(65536)
  smoothblue = zeros(65536)
  width = 300          # include this many bins on either side
  redpeak = 0
  maxred = 0
  greenpeak = 0
  maxgreen = 0
  bluepeak = 0
  maxblue = 0
  for i = 400:65500 # ends are left off to avoid clipping artifacts
    low = i - width - 1 < 1 ? 1 : i - width - 1
    high = i + width > 65536 ? 65536 : i + width
    smoothred[i] = (redacc[high] - redacc[low]) / (high - low)
    smoothgreen[i] = (greenacc[high] - greenacc[low]) / (high - low)
    smoothblue[i] = (blueacc[high] - blueacc[low]) / (high - low)
    if smoothred[i] > maxred
      maxred = smoothred[i]
      redpeak = i
    end
    if smoothgreen[i] > maxgreen
      maxgreen = smoothgreen[i]
      greenpeak = i
    end
    if smoothblue[i] > maxblue
      maxblue = smoothblue[i]
      bluepeak = i
    end
  end
  if redpeak == 0 || greenpeak == 0 || bluepeak == 0
    @printf "Max search failed red %d green %d blue %d\n" redpeak greenpeak bluepeak
    exit(-1)
  end
  println("\n     image histogram:    max bin    pixel count")
  @printf "                  red:   %d         %f\n" redpeak maxred
  @printf "                green:   %d         %f\n" greenpeak maxgreen
  @printf "                 blue:   %d         %f\n" bluepeak maxblue

  redsky = maxred * skylevelfactor
  greensky = maxgreen * skylevelfactor
  bluesky = maxblue * skylevelfactor
  redskybin = 0
  greenskybin = 0
  blueskybin = 0

  #find where redvals[i] >= greenske and redlevel[i-1] <= greensky
  for i = redpeak:-1:2
    if smoothred[i] >= greensky && smoothred[i - 1] <= greensky
      redskybin = i + 1
      break
    end
  end

  for i = greenpeak:-1:2
    if smoothgreen[i] >= greensky && smoothgreen[i - 1] <= greensky
      greenskybin = i + 1
      break
    end
  end

  for i = bluepeak:-1:2
    if smoothblue[i] >= greensky && smoothblue[i - 1] <= greensky
      blueskybin = i + 1
      break
    end
  end

  if redskybin == 0 || greenskybin == 0 || blueskybin == 0
    @printf "histogram sky level %f not found:\n" skylevelfactor
    @printf "     channels: red %d   green %d   blue %d\n" redskybin greenskybin blueskybin
    println("          Try increasing -rgbskyzero values\n")
    exit(-1)
  end

  println("Histogram dark sky level")
  println("\n                  bin      Number of pixels in bin")     # need to compute cumulative hstogram
  @printf "           red:   %d         %d\n" redskybin   smoothred[redskybin]
  @printf "         green:   %d         %d\n" greenskybin smoothgreen[greenskybin]
  @printf "          blue:   %d         %d\n" blueskybin  smoothblue[blueskybin]

  redskysub1 = redskybin - zeroskyred
  greenskysub1 = greenskybin - zeroskygreen
  blueskysub1 = blueskybin - zeroskyblue

	@printf "\nsubtract %7d from red   to make red   sky align with zero reference sky: %7.1f\n" redskysub1   zeroskyred
	@printf "subtract %7d from green to make green sky align with zero reference sky: %7.1f\n" greenskysub1 zeroskygreen
	@printf "subtract %7d from blue  to make blue  sky align with zero reference sky: %7.1f\n" blueskysub1  zeroskyblue

	#@printf "\nnow set the RGB sky zero level to %7.1f  %7.1f  %7.1f bin out of 65535\n", zeroskyred zeroskygreen zeroskyblue

  redvals = greenvals = bluevals = collect(0:65535)
  redvals = (redvals .- redskysub1) .* (65535 / (65535 - redskysub1))
  greenvals = (greenvals .- greenskysub1) .* (65535 / (65535 - greenskysub1))
  bluevals = (bluevals .- blueskysub1) .* (65535 / (65535 - blueskysub1))
  redvals[redvals .< 0] = 0
  greenvals[greenvals .< 0] = 0
  bluevals[bluevals .< 0] = 0
  redvals = floor(redvals)
  greenvals = floor(greenvals)
  bluevals = floor(bluevals)
  remap()

  println("sky adjusted")
  showmeans()
end
###### end of function smoothsky

function setmins()
  global redvals, greenvals, bluevals
  println("\nsetting minimum levels.")
  zx = 0.2
  redvals = collect(0:65535)
  screen = redvals .< setminr
  redvals[screen] = redvals[screen] .* zx .+ setminr
  greenvals = collect(0:65535)
  screen = greenvals .< setming
  greenvals[screen] = greenvals[screen] .* zx .+ setming
  bluevals = collect(0:65535)
  screen = bluevals .< setminb
  bluevals[screen] = bluevals[screen] .* zx .+ setminb
  remap()

end

filename = "BaseOrion.tif"
print("Loading $filename... ")
tic()
img = load(filename)
save("save1.png", img)
if img == nothing
  println("Image failed to load.")
  exit(-1)
end
println("Done!")
println("Image loading time: ", toq())
println("Loading processing vectors and histograms...")
tic()
sz = size(img)
px = sz[1]
py = sz[2]
npix = px * py
npixm = floor(npix * pcntclip / 100.0)
npixm = npixm < 1 ? 1 : npixm
@printf "\nImage size: %d x %d pixels, cut level = %d\n" px py npixm
img = reinterpret(UInt16, img)
imgr = reshape(img[1, :, :], npix)
imgg = reshape(img[2, :, :], npix)
imgb = reshape(img[3, :, :], npix)
img = 0
#=
redcounts = fit(Histogram, imgr, 0:65536).weights
greencounts = fit(Histogram, imgg, 0:65536).weights
bluecounts = fit(Histogram, imgb, 0:65536).weights
=#
# still can't make this work right :(
redcounts = zeros(UInt32, 65536)
greencounts = zeros(UInt32, 65536)
bluecounts = zeros(UInt32, 65536)
for i = 1:npix
  redcounts[imgr[i] + 1] += 1
  greencounts[imgg[i] + 1] += 1
  bluecounts[imgb[i] + 1] += 1
end

println("Histogram loading time: ", toq())

plotter("initial")
if doplots == 1
  plotter("Initial")
end

# pixel values are off by one throughout calculations
# because of 1-base arrays

redmap = collect(1:65536)
greenmap = collect(1:65536)
bluemap = collect(1:65536)

print("\n\nInitial bounds and means:\n")
showmeans()

if tonecurve == 1
  print("\n\nApplying tone curve... ")
  b = 12
  c = 65535
  d = 12
  redvals = greenvals = bluevals = collect(0:65535)
  # factor = b * ((1 / d)) ^ ((val / c) ^ 0.4)
  # could probably be done in one statement but since it's only done once,
  # it's not time-critical and not worth the time to figure out
  f = (map( x-> x ^ 0.4, (collect(0:65535) ./ c)))
  f = map( x -> b * (1 / d) ^ x, f)
  g = summarystats(f)
  print(g)
  redvals = floor(redvals .* f)
  greenvals = floor(greenvals .* f)
  bluevals = floor(bluevals .* f)
  println("Done!")
  remap()
  print("after tone curve:\n")
  showmeans()
else
  print("\nNot doing tone curve\n")
end

alowred   = 0  # find the DN level which is the cut level
alowgreen = 0
alowblue  = 0
nred   = redcounts[1]     # number of pixels in histogram bin
ngreen = greencounts[1]
nblue  = bluecounts[1]

for i = 2:32000   # upper limit is more or less pulled out of the air
                  # don't need to look at the first element because it's already zero
  if alowred == 0
    nred = nred + redcounts[i]
    if nred > npixm
      alowred = i
    end
  end
  if alowgreen == 0
    ngreen = ngreen + greencounts[i]
    if ngreen > npixm
      alowgreen = i
    end
  end
  if alowblue == 0
    nblue = nblue + bluecounts[i]
    if nblue > npixm
      alowblue = i
    end
  end
  if alowred > 0 && alowgreen > 0 && alowblue > 0
    break
  end
end

@printf "\n\ncumulative histogram cut level (%f %% = %d pixels):\n" pcntclip floor(npixm)
@printf "           red:   %d  (%d pixels)\n" alowred nred
@printf "           green: %d  (%d pixels)\n" alowgreen ngreen
@printf "           blue:  %d  (%d pixels)\n" alowblue nblue

# reset dark sky
for pass = 1:2
  @printf "\n\ndark sky subtraction pass %d\n" pass
#  println("r g b checksums before ", checksum(redvals), " ", checksum(greenvals), " ", checksum(bluevals))
  smoothsky()
end

afr = map(x -> redmap[x + 1], imgr)
afg = map(x -> greenmap[x + 1], imgg)
afb = map(x -> bluemap[x + 1], imgb)
# assemble sky-subtracted image here
g = summarystats(afr)
println(">>> afr summary")
print(g)
println("============= computing root stretch =============")
for pass = 1:rootiter
  if pass == 2
    x = 1 / rootpower2
  else
    x = 1 / rootpower
  end

  println("Root stretch pass $pass of $rootiter")
  redfloat = collect(1:65536) / 65536 # has a +1 built in
  greenfloat = collect(1:65536) / 65536
  bluefloat = collect(1:65536) / 65536
  redfloat = map(z -> z ^ x, redfloat)
  greenfloat = map(z -> z ^ x, greenfloat)
  bluefloat = map(z -> z ^ x, bluefloat)
  redfloat = redfloat .* 65535
  greenfloat = greenfloat .* 65535
  bluefloat = bluefloat .* 65535
#  remap()
  redmin = reduce(min, 100000, redfloat[redcounts .> 0])
  greenmin = reduce(min, 100000, greenfloat[greencounts .> 0])
  bluemin = reduce(min, 100000, bluefloat[bluecounts .> 0])
  bmin = min(redmin, greenmin, bluemin)
  println("redmin $redmin greenmin $greenmin bluemin $bluemin bmin $bmin")
  if bmin < 4096
    bminz = 0
  else
    bminz = bmin - 4096
  end
  println("Subtracting $bminz from root stretched image")
  redfloat = (redfloat .- bminz) / (65535 - bminz) * 65535
  greenfloat = (greenfloat .- bminz) / (65535 - bminz) * 65535
  bluefloat = (bluefloat .- bminz) / (65535 - bminz) * 65535
  redfloat[redcounts .== 0] = 0
  greenfloat[greencounts .== 0] = 0
  bluefloat[bluecounts .== 0] = 0
  redvals = map(x -> floor(UInt16, x), redfloat)
  greenvals = map(x -> floor(UInt16, x), greenfloat)
  bluevals = map(x -> floor(UInt16, x), bluefloat)
  remap()

  showmeans()

  for pass2 = 1:(rootpower > 60 ? 3 : 2)
    println("Smooth sky pass $pass2 within root stretch pass $pass")
    smoothsky()
  end
end

if scurve > 0
  println("========== Doing s-curve stretching ==========")
  for i = 1:scurve
    if i == 2 || i == 4
      xfactor = 3
      xoffset = 0.22
    else
      xfactor = 5
      xoffset = 0.42
    end

    scurvemin = (xfactor / (1.0 + exp(-1.0 * (0 - xoffset) * xfactor))) - (1.0 - xoffset)
    scurvemax = (xfactor / (1.0 + exp(-1.0 * (1 - xoffset) * xfactor))) - (1.0 - xoffset)
    scurveminsc = scurvemin / scurvemax

    println("s-curve pass $i")
    println("          xfactor     = $xfactor")
    println("          xoffset     = $xoffset")
    @printf "          scurvemin   = %8.4f\n" scurvemin
    @printf "          scurvemax   = %8.4f\n" scurvemax
    @printf "          scurveminsc = %8.4f\n" scurveminsc

    xo = 1 - xoffset

    redsc = (collect(0:65535) ./ 65535 .- xoffset) .* -xfactor
    redsc = map(exp, redsc)
    redsc = redsc .+ 1
    redsc = (xfactor ./ redsc) .- xo
    redsc = redsc ./ scurvemax .- scurveminsc
    redsc = 65535 .* redsc
    redvals = map(x -> floor(UInt64, x), redsc ./ (1 - scurveminsc))

    greensc = (collect(0:65535) ./ 65535 .- xoffset) .* -xfactor
    greensc = map(exp, greensc)
    greensc = greensc .+ 1
    greensc = (xfactor ./ greensc) .- xo
    greensc = greensc ./ scurvemax .- scurveminsc
    greensc = 65535 .* greensc
    greenvals = map(x -> floor(UInt64, x), greensc ./ (1 - scurveminsc))

    bluesc = (collect(0:65535) ./ 65535 .- xoffset) .* -xfactor
    bluesc = map(exp, bluesc)
    bluesc = bluesc .+ 1
    bluesc = (xfactor ./ bluesc) .- xo
    bluesc = bluesc ./ scurvemax .- scurveminsc
    bluesc = 65535 .* bluesc
    bluevals = map(x -> floor(UInt64, x), bluesc ./ (1 - scurveminsc))

    remap()

    println("\n")
    showmeans()
  end

  for j = 1:2
    println("Doing sky offset pass $j")
    smoothsky()
  end
end

if setmin > 0
  setmins()
end

#plotter("Before final load")
redvals = 0
greenvals = 0
bluevals = 0
newred = 0
newgreen = 0
newblue = 0

cr = map(x -> redmap[x + 1] - 1, imgr)
cg = map(x -> greenmap[x + 1] - 1, imgg)
cb = map(x -> bluemap[x + 1] - 1, imgb)
cr[cr .== 0] = 1
cg[cg .== 0] = 1
cb[cb .== 0] = 1

redmap = 0
greenmap = 0
bluemap = 0
imgr = 0
imgg = 0
imgb = 0

crmin = reduce(min, 65537, cr)
crmax = reduce(max, 0, cr)
crmean = mean(cr)

cgmin = reduce(min, 65537, cg)
cgmax = reduce(max, 0, cg)
cgmean = mean(cg)

cbmin = reduce(min, 65537, cb)
cbmax = reduce(max, 0, cb)
cbmean = mean(cb)

println("Image reloaded")
@printf "red   min = %d max = %d mean = %7.1f\n" crmin crmax crmean
@printf "green min = %d max = %d mean = %7.1f\n" cgmin cgmax cgmean
@printf "blue  min = %d max = %d mean = %7.1f\n" cbmin cbmax cbmean

if colorcorrect > 0
  println("\n ========== final color adjustment ==========")
  afr = afr .- zeroskyred
  afr[afr .< 10] = 10
  afg = afg .- zeroskygreen
  afg[afg .< 10] = 10
  afb = afb .- zeroskyblue
  afb[afb .< 10] = 10

  println("color adjustments")

  grratio = (afg ./ afr) ./ (cg ./ cr)
  brratio = (afb ./ afr) ./ (cb ./ cr)
  rgratio = (afr ./ afg) ./ (cr ./ cg)
  bgratio = (afb ./ afg) ./ (cb ./ cg)
  gbratio = (afg ./ afb) ./ (cg ./ cb)
  rbratio = (afr ./ afb) ./ (cr ./ cb)

  zmin = 0.2
  zmax = 1.0

  grratio[grratio .< zmin] = zmin
  grratio[grratio .> zmax] = zmax

  brratio[brratio .< zmin] = zmin
  brratio[brratio .> zmax] = zmax

  rgratio[rgratio .< zmin] = zmin
  rgratio[rgratio .> zmax] = zmax

  bgratio[bgratio .< zmin] = zmin
  bgratio[bgratio .> zmax] = zmax

  gbratio[gbratio .< zmin] = zmin
  gbratio[gbratio .> zmax] = zmax

  rbratio[rbratio .< zmin] = zmin
  rbratio[rbratio .> zmax] = zmax

  println("\nColor ratio images after constraint to interval ($zmin, $zmax)")
  g = summarystats(grratio)
  print(g)
  g = summarystats(brratio)
  print(g)
  g = summarystats(rgratio)
  print(g)
  g = summarystats(bgratio)
  print(g)
  g = summarystats(gbratio)
  print(g)
  g = summarystats(rbratio)
  print(g)

  avgn = (cr .+ cg .+ cb) ./ (3 * 65535)
  avgn[avgn .< 0] = 0
  cmax = reduce(max, 0, avgn)
  if cmax < 1
    avgn ./= cmax
  end

  avgn .^= 0.2
  avgn .+= 0.3
  avgn ./= 1.3
  cmin = reduce(min, 65537, avgn)
  cmax = reduce(max, 0, avgn)
  cmean = mean(avgn)

  @printf "Color correction intensity min = %7.5f max = %7.5f mean = %7.5f\n" cmin cmax cmean

  cfactor = 1.2
  cfe = avgn .* (cfactor * colorenhance)
  grratio = 1 .+ cfe .* (grratio .- 1)
  brratio = 1 .+ cfe .* (brratio .- 1)
  rgratio = 1 .+ cfe .* (rgratio .- 1)
  bgratio = 1 .+ cfe .* (bgratio .- 1)
  gbratio = 1 .+ cfe .* (gbratio .- 1)
  rbratio = 1 .+ cfe .* (rbratio .- 1)

  println("Color ratio images after factors applied")
  println("grratio:")
  describe(grratio)
  println("brratio:")
  describe(brratio)
  println("rgratio:")
  describe(rgratio)
  println("bgratio:")
  describe(bgratio)
  println("gbratio:")
  describe(gbratio)
  println("rbratio:")
  describe(rbratio)

  cmin = reduce(min, 65537, cfe)
  cmax = reduce(max, 0, cfe)
  cmean = mean(cfe)
  @printf "Color correction cfe min = %7.5f max = %7.5f mean = %7.5f\n" cmin cmax cmean

  println("Computing possible image sets")
  c2gr = cg .* grratio
  c3br = cb .* brratio
  c1rg = cr .* rgratio
  c3bg = cb .* bgratio
  c1rb = cr .* rbratio
  c2gb = cg .* gbratio

  println("Starting signal dependent color recovery")
  println("Processing $npix pixels....")
  decile = floor(npix / 10)
  pixelcount = decile
  for i = 1:npix
    pixelcount -= 1
    if pixelcount == 0
      println("Processing pixel $i of $npix...")
      pixelcount = decile
    end
    # I think this could be done with stacked ternary operators
    maxv = cr[i]
    maxch = 1
    if cg[i] > maxv
      maxv = cg[i]
      maxch = 2
    end
    if cb[i] > maxv
      # don't need the max value!
      maxch = 3
    end
    if maxch == 1
      cg[i] = floor(UInt16, c2gr[i])
      cb[i] = floor(UInt16, c3br[i])
    elseif maxch == 2
      cr[i] = floor(UInt16, c1rg[i])
      cb[i] = floor(UInt16, c3bg[i])
    else
      cr[i] = floor(UInt16, c1rb[i])
      cg[i] = floor(UInt16, c2gb[i])
    end
  end

  println("Color recovery complete")
end

if setmin > 0
  setmins()
end

smoothsky()
#=if cumstretch == 1
  println("Performing cumulative histogram strech")
  hist = fit(Histogram, [cr, cg, cb], 0:65536).weights
  #redcounts = fit(Histogram, imgr, 0:65536).weights
  cs = chist
  gap = 15000
  for i=1:65535
    ilow = i <= gap ? 1 : i - gap
    ihigh = i >= 65535 - gap ? 65535 : i + gap
    =#

cr[cr .< 0] = 0
cr[cr .> 65535] = 65535
cg[cg .< 0] = 0
cg[cg .> 65535] = 65535
cb[cb .< 0] = 0
cb[cb .> 65535] = 65535

finalr = reshape(map(x -> floor(UInt16, x), cr), px, py)
finalg = reshape(map(x -> floor(UInt16, x), cg), px, py)
finalb = reshape(map(x -> floor(UInt16, x), cb), px, py)
final = zeros(UInt16, 3, px, py)
final[1, :, :] = finalr
final[2, :, :] = finalg
final[3, :, :] = finalb

final = colorview(RGB{N0f16}, final)

#final = reinterpret()
oname = "$obase.png"
oname = "neworionout.png"
print("Saving $oname... ")
save(oname, final)
print("Done!")


# Important information about oldTraj:
# ======================================
# oldTraj.oldTraj contains all state information of one lap. It is structured like follows:
# The first <prebuf> values are the end of one lap before the next lap starts (necessary for inter-lap-system-ID)
# The last value of <prebuf> ends with the last value *before* the finish line (s < s_target)
# The s-values within <prebuf> need to be below zero (-> before the finish line). Otherwise they can be mistaken as the end of the current (not previous) lap.
# After <prebuf> follow the recorded states of the trajectory (as many as there were during one lap)
# The number of recorded states during one trajectory (0 <= s < s_target) is equal to <oldCost>.
# After the recorded trajectory, the rest of the vector (until <buffersize>) is filled up with constant values

function saveOldTraj(oldTraj::OldTrajectory,zCurr::Array{Float64},uCurr::Array{Float64},lapStatus::LapStatus,posInfo::PosInfo,buffersize::Int64)
                println("Starting function")
                i               = lapStatus.currentIt-1         # i = number of points for 0 <= s < s_target (= cost of this lap)
                prebuf          = oldTraj.prebuf                # so many points of the end of the previous old traj will be attached to the beginning
                zCurr_export    = zeros(buffersize,6)
                uCurr_export    = zeros(buffersize,2)

                println("Creating exports")
                zCurr_export    = cat(1,oldTraj.oldTraj[oldTraj.oldCost[1]+1:oldTraj.oldCost[1]+prebuf,:,1],
                                        zCurr[1:i,:], NaN*ones(buffersize-i-prebuf,6))#[ones(buffersize-i-prebuf,1)*zCurr[i,1:5] zCurr[i,6]+collect(1:buffersize-i-prebuf)*dt*zCurr[i,1]])
                uCurr_export    = cat(1,oldTraj.oldInput[oldTraj.oldCost[1]+1:oldTraj.oldCost[1]+prebuf,:,1],
                                        uCurr[1:i,:], NaN*ones(buffersize-i-prebuf,2))#zeros(buffersize-i-prebuf,2))

                zCurr_export[1:prebuf,6] -= posInfo.s_target       # make the prebuf-values below zero
                costLap                   = i                      # the cost of the current lap is the time it took to reach the finish line
                println("Saving")
                # Save all data in oldTrajectory:
                if lapStatus.currentLap <= 2                        # if it's the first or second lap
                    oldTraj.oldTraj[:,:,1]  = zCurr_export          # ... just save everything
                    oldTraj.oldInput[:,:,1] = uCurr_export
                    oldTraj.oldTraj[:,:,2]  = zCurr_export
                    oldTraj.oldInput[:,:,2] = uCurr_export
                    oldTraj.oldCost = [costLap,costLap]
                else                                                # idea: always copy the new trajectory in the first array!
                    if oldTraj.oldCost[1] < oldTraj.oldCost[2]      # if the first old traj is better than the second
                        oldTraj.oldTraj[:,:,2]  = oldTraj.oldTraj[:,:,1]    # ... copy the first in the second
                        oldTraj.oldInput[:,:,2] = oldTraj.oldInput[:,:,1]   # ... same for the input
                        oldTraj.oldCost[2] = oldTraj.oldCost[1]
                    end
                    oldTraj.oldTraj[:,:,1]  = zCurr_export                 # ... and write the new traj in the first
                    oldTraj.oldInput[:,:,1] = uCurr_export
                    oldTraj.oldCost[1] = costLap
                end
end

function InitializeParameters(mpcParams::MpcParams,mpcParams_pF::MpcParams,trackCoeff::TrackCoeff,modelParams::ModelParams,
                              posInfo::PosInfo,oldTraj::OldTrajectory,mpcCoeff::MpcCoeff,lapStatus::LapStatus,buffersize::Int64,
                              obstacle::Obstacle,selectedStates::SelectedStates,oldSS::SafeSetData)
    selectedStates.simulator = false     # set this to TRUE if SIMULATOR is in use, set this to FALSE if BARC is in use

    if selectedStates.simulator == false   # if the BARC is in use

        selectedStates.Np           = 15                           # Number of points to take from each previous trajectory to build the convex hull
        selectedStates.Nl           = 2                             # Number of previous laps to include in the convex hull
        selectedStates.shift        = 8
        Nl                          = selectedStates.Nl
        selectedStates.selStates    = zeros(Nl*selectedStates.Np,6)
        selectedStates.statesCost   = zeros(Nl*selectedStates.Np)
        selectedStates.version      = false

        mpcParams.N                 = 14
        mpcParams.Q                 = [5.0,0.0,0.0,0.1,50.0,0.0]   # Q (only for path following mode)
        mpcParams.vPathFollowing    = 1                           # reference speed for first lap of path following
        mpcParams.Q_term            = 1.0*[20.0,1.0,10.0,20.0,50.0]   # weights for terminal constraints (LMPC, for xDot,yDot,psiDot,ePsi,eY).Not used if using convex hull
        mpcParams.R                 = 0*[10.0,10.0]                 # put weights on a and d_f
        mpcParams.QderivZ           = 10.0*[1,1,1,1,1,1]             # cost matrix for derivative cost of states
        mpcParams.QderivU           = 1.0*[5.0,1.0] #NOTE Set this to [5.0, 0/40.0]              # cost matrix for derivative cost of inputs
        mpcParams.Q_term_cost       = 3                        # scaling of Q-function
        mpcParams.delay_df          = 3                             # steering delay
        mpcParams.delay_a           = 1                             # acceleration delay
        mpcParams.Q_lane            = 1                      # weight on the soft constraint for the lane
        mpcParams.Q_vel             = 1                    # weight on the soft constraint for the maximum velocity
        mpcParams.Q_slack           = 1*[5*20.0,20.0,10.0,30.0,80.0,50.0]#[20.0,10.0,10.0,30.0,80.0,50.0]  #vx,vy,psiDot,ePsi,eY,s
        mpcParams.Q_obs             = ones(Nl*selectedStates.Np)# weight to esclude some of the old trajectories

    elseif selectedStates.simulator == true  # if the simulator is in use

        selectedStates.Np           = 15                           # Number of points to take from each previous trajectory to build the convex hull
        selectedStates.Nl           = 2                             # Number of previous laps to include in the convex hull
        selectedStates.shift        = 5
        Nl                          = selectedStates.Nl
        selectedStates.selStates    = zeros(Nl*selectedStates.Np,6)
        selectedStates.statesCost   = zeros(Nl*selectedStates.Np)
        selectedStates.version      = false

        mpcParams.N                 = 13
        mpcParams.Q                 = [5.0,0.0,0.0,0.1,50.0,0.0]   # Q (only for path following mode)
        mpcParams.vPathFollowing    = 1                           # reference speed for first lap of path following
        mpcParams.Q_term            = 1.0*[20.0,1.0,10.0,20.0,50.0]   # weights for terminal constraints (LMPC, for xDot,yDot,psiDot,ePsi,eY).Not used if using convex hull
        mpcParams.R                 = 0*[10.0,10.0]                 # put weights on a and d_f
        mpcParams.QderivZ           = 10.0*[1,1,1,1,1,1]             # cost matrix for derivative cost of states
        mpcParams.QderivU           = 1.0*[1,1.0] #NOTE Set this to [5.0, 0/40.0]              # cost matrix for derivative cost of inputs
        mpcParams.Q_term_cost       = 1                        # scaling of Q-function
        mpcParams.delay_df          = 3                             # steering delay
        mpcParams.delay_a           = 1                             # acceleration delay
        mpcParams.Q_lane            = 3                      # weight on the soft constraint for the lane
        mpcParams.Q_vel             = 1                    # weight on the soft constraint for the maximum velocity
        mpcParams.Q_slack           = 1*[20.0,1.0,10.0,30.0,80.0,50.0]#[20.0,10.0,10.0,30.0,80.0,50.0]  #vx,vy,psiDot,ePsi,eY,s
        mpcParams.Q_obs             = ones(Nl*selectedStates.Np)# weight to esclude some of the old trajectories
    end

    mpcParams_pF.N              = 16
    mpcParams_pF.Q              = [0.0,50.0,1.0,10.0]
    mpcParams_pF.R              = 0*[1.0,1.0]               # put weights on a and d_f
    mpcParams_pF.QderivZ        = 0.0*[0.0,0,1.0,0]           # cost matrix for derivative cost of states
    mpcParams_pF.QderivU        = 10*[1,1]                # cost matrix for derivative cost of inputs
    mpcParams_pF.vPathFollowing = 1                       # reference speed for first lap of path following
    mpcParams_pF.delay_df       = 3                         # steering delay (number of steps)
    mpcParams_pF.delay_a        = 1                         # acceleration delay

    trackCoeff.nPolyCurvature   = 8                         # 4th order polynomial for curvature approximation
    trackCoeff.coeffCurvature   = zeros(trackCoeff.nPolyCurvature+1)         # polynomial coefficients for curvature approximation (zeros for straight line)
    trackCoeff.width            = 0.6                       # width of the track (60cm)

    modelParams.l_A             = 0.125
    modelParams.l_B             = 0.125
    modelParams.dt              = 0.1                   # sampling time, also controls the control loop, affects delay_df and Qderiv
    modelParams.m               = 1.98
    modelParams.I_z             = 0.03
    modelParams.c_f             = 0.5                   # friction coefficient: xDot = - c_f*xDot (aerodynamic+tire)

    posInfo.s_target            = 19.11

    oldTraj.oldTraj             = NaN*ones(buffersize,7,30)
    oldTraj.oldInput            = zeros(buffersize,2,30)
    oldTraj.oldTimes            = NaN*ones(buffersize,30)
    oldTraj.count               = ones(30)*2
    oldTraj.oldCost             = ones(Int64,30)                   # dummies for initialization
    oldTraj.prebuf              = 30
    oldTraj.postbuf             = 30
    oldTraj.idx_start           = zeros(30)
    oldTraj.idx_end             = zeros(30)

    oldSS.oldSS                 = NaN*ones(buffersize,7,30)          # contains data from previous laps usefull to build the safe set
    oldSS.oldSS_xy              = NaN*ones(buffersize,4,30)
    oldSS.cost2target           = zeros(buffersize,30)     # cost to arrive at the target, i.e. how many iterations from the start to the end of the lap
    oldSS.oldCost               = ones(Int64,30)              # contains costs of laps
    oldSS.count                 = ones(30)*2               # contains the counter for each lap
    oldSS.prebuff                = 30
    oldSS.postbuff               = 50
    oldSS.idx_start             = ones(30)            # index of the first measurement with s > 0
    oldSS.idx_end               = zeros(30)              # index of the last measurement with s < s_target

    mpcCoeff.order              = 5
    mpcCoeff.coeffCost          = zeros(mpcCoeff.order+1,2)
    mpcCoeff.coeffConst         = zeros(mpcCoeff.order+1,2,5)
    mpcCoeff.pLength            = 5*2*mpcParams.N        # small values here may lead to numerical problems since the functions are only approximated in a short horizon
    mpcCoeff.c_Vx               = zeros(3)
    mpcCoeff.c_Vy               = zeros(4)
    mpcCoeff.c_Psi              = zeros(3)

    lapStatus.currentLap        = 1         # initialize lap number
    lapStatus.currentIt         = 0         # current iteration in lap

    obstacle.obstacle_active    = false     # true if we have to consider the obstacles in the optimization problem
    obstacle.lap_active         = 1000         # number of the first lap in which the obstacles are used
    obstacle.obs_detect         = 10         # maximum distance at which we can detect obstacles (in terms of s!!)
    obstacle.n_obs              = 1         # number of obstacles
    obstacle.s_obs_init         = [14]    # initial s coordinate of each obstacle
    obstacle.ey_obs_init        = [-0.3]       # initial ey coordinate of each obstacle
    obstacle.v_obs_init         = [0]       # initial velocity of each obstacles
    obstacle.r_s                = 0.2#0.5
    obstacle.r_ey               = 0.1#0.2
    obstacle.inv_step           = 0         # number of step of invariance required for the safe set
end

# Use this function to smooth data (moving average)
function smooth(x,n)
    y = zeros(size(x))
    for i=1:size(x,1)
        start = max(1,i-n)
        fin = min(size(x,1),start + 2*n)
        y[i,:] = mean(x[start:fin,:],1)
    end
    return y
end


include(joinpath("deps", "ReadMolecule.jl"))
include(joinpath("deps", "Shapes.jl"))
import .ReadMoelcule
import .Shapes
import SimpleMolecule: Molecule, obtain
import MolecularPDB

function main(argd)
    systems = ReadMolecule.read(convert(Array{String,1}, argd["inputs"]))

    fn_assemble = Shapes.choose(Symbol(argd["shape"]))

    @time molecules = fn_assemble(systems, argd["params"])
    new_system = Molecule(vcat([obtain(mol, :atoms) for mol in molecules]...))

    for file_output in argd["outputs"]
        @time MolecularPDB.save(
            MolecularPDB.VMD,
            file_output,
            new_system
        )
    end
end


@testset "URDF parse" begin
    @testset "joint bounds" begin
        acrobot = parse_urdf(joinpath(@__DIR__, "urdf", "Acrobot.urdf"), remove_fixed_tree_joints=false)
        @test position_bounds(findjoint(acrobot, "shoulder")) == [RigidBodyDynamics.Bounds(-Inf, Inf)]
        @test velocity_bounds(findjoint(acrobot, "shoulder")) == [RigidBodyDynamics.Bounds(-Inf, Inf)]
        @test effort_bounds(findjoint(acrobot, "shoulder")) == [RigidBodyDynamics.Bounds(-Inf, Inf)]
        @test position_bounds(findjoint(acrobot, "elbow")) == [RigidBodyDynamics.Bounds(-Inf, Inf)]
        @test velocity_bounds(findjoint(acrobot, "elbow")) == [RigidBodyDynamics.Bounds(-Inf, Inf)]
        @test effort_bounds(findjoint(acrobot, "elbow")) == [RigidBodyDynamics.Bounds(-Inf, Inf)]

        acrobot_with_limits = parse_urdf(joinpath(@__DIR__, "urdf", "Acrobot_with_limits.urdf"), remove_fixed_tree_joints=false)
        @test position_bounds(findjoint(acrobot_with_limits, "shoulder")) == [RigidBodyDynamics.Bounds(-6.28, 6.28)]
        @test velocity_bounds(findjoint(acrobot_with_limits, "shoulder")) == [RigidBodyDynamics.Bounds(-10, 10)]
        @test effort_bounds(findjoint(acrobot_with_limits, "shoulder")) == [RigidBodyDynamics.Bounds(0, 0)]
        @test position_bounds(findjoint(acrobot_with_limits, "elbow")) == [RigidBodyDynamics.Bounds(-6.28, 6.28)]
        @test velocity_bounds(findjoint(acrobot_with_limits, "elbow")) == [RigidBodyDynamics.Bounds(-10, 10)]
        @test effort_bounds(findjoint(acrobot_with_limits, "elbow")) == [RigidBodyDynamics.Bounds(-5, 5)]
    end

    @testset "planar joints" begin
        robot = parse_urdf(joinpath(@__DIR__, "urdf", "planar_slider.urdf"), remove_fixed_tree_joints=false)
        # the first joint is the fixed joint between the world and the base link
        jt = joint_type(joints(robot)[1])
        @test jt isa Fixed

        # the 2nd joint is planar with <axis xyz="1 0 0"/>
        jt = joint_type(joints(robot)[2])
        @test jt isa Planar
        @test jt.x_axis ≈ SVector(0, 0, -1)
        @test jt.y_axis ≈ SVector(0, 1, 0)

        # the 3rd joint is planar with <axis xyz="0 1 0"/>
        jt = joint_type(joints(robot)[3])
        @test jt isa Planar
        @test jt.x_axis ≈ SVector(1, 0, 0)
        @test jt.y_axis ≈ SVector(0, 0, -1)

        # the 4th joint is planar with <axis xyz="0 0 1"/>
        jt = joint_type(joints(robot)[4])
        @test jt isa Planar
        @test jt.x_axis ≈ SVector(1, 0, 0)
        @test jt.y_axis ≈ SVector(0, 1, 0)
    end

    @testset "rotation handling" begin
        # https://github.com/JuliaRobotics/RigidBodyDynamics.jl/issues/429
        xml = """
            <origin rpy="0 0 0"/>
        """
        xpose = LightXML.root(LightXML.parse_string(xml))
        rot, trans = RigidBodyDynamics.parse_pose(Float64, xpose)
        @test rot ≈ SDiagonal(1, 1, 1)
        @test trans ≈ SVector(0, 0, 0)

        xml = """
            <origin rpy="0.3 0 0"/>
        """
        xpose = LightXML.root(LightXML.parse_string(xml))
        rot, trans = RigidBodyDynamics.parse_pose(Float64, xpose)
        @test rot ≈ RotMatrix(RotX(0.3))
        @test trans ≈ SVector(0, 0, 0)

        xml = """
            <origin rpy="0 0.2 0"/>
        """
        xpose = LightXML.root(LightXML.parse_string(xml))
        rot, trans = RigidBodyDynamics.parse_pose(Float64, xpose)
        @test rot ≈ RotMatrix(RotY(0.2))
        @test trans ≈ SVector(0, 0, 0)

        xml = """
            <origin rpy="0 0 0.1"/>
        """
        xpose = LightXML.root(LightXML.parse_string(xml))
        rot, trans = RigidBodyDynamics.parse_pose(Float64, xpose)
        @test rot ≈ RotMatrix(RotZ(0.1))
        @test trans ≈ SVector(0, 0, 0)

        # Comparison against ROS's tf.transformations.quaternion_from_euler
        xml = """
            <origin rpy="1 2 3"/>
        """
        xpose = LightXML.root(LightXML.parse_string(xml))
        rot, trans = RigidBodyDynamics.parse_pose(Float64, xpose)
        @test rot ≈ [0.41198225 -0.83373765 -0.36763046;
                    -0.05872664 -0.42691762  0.90238159;
                    -0.90929743 -0.35017549 -0.2248451] atol=1e-7
        @test rot ≈ RotZ(3) * RotY(2) * RotX(1)
        @test trans ≈ SVector(0, 0, 0)

        # Comparison against ROS's tf.transformations.quaternion_from_euler
        xml = """
            <origin rpy="0.5 0.1 0.2"/>
        """
        xpose = LightXML.root(LightXML.parse_string(xml))
        rot, trans = RigidBodyDynamics.parse_pose(Float64, xpose)
        @test rot ≈ [0.97517033 -0.12744012  0.18111281;
                     0.19767681  0.86959819 -0.45246312;
                    -0.09983342  0.47703041  0.8731983] atol=1e-7
        @test rot ≈ RotZ(0.2) * RotY(0.1) * RotX(0.5)
        @test trans ≈ SVector(0, 0, 0)
    end
end

function test_inverse_dynamics_match(mechanism1::Mechanism, mechanism2::Mechanism; atol::Real=0, rtol::Real=atol>0 ? 0 : √eps)
    state1 = MechanismState(mechanism1)
    v̇_vec = rand(num_velocities(mechanism1))
    rand!(state1)
    v̇1 = copyto!(similar(velocity(state1)), v̇_vec)
    τ1 = inverse_dynamics(state1, v̇1)

    state2 = MechanismState(mechanism2)
    copyto!(state2, Vector(state1))
    v̇2 = copyto!(similar(velocity(state2)), v̇_vec)
    τ2 = inverse_dynamics(state2, v̇2)

    @test τ1 ≈ τ2 rtol=rtol atol=atol
end

function test_kinematic_graph_layout_match(mechanism1::Mechanism, mechanism2::Mechanism)
    for (j1, j2) in zip(joints(mechanism1), joints(mechanism2))
        @test string(j1) == string(j2)
        @test string(predecessor(j1, mechanism1)) == string(predecessor(j2, mechanism2))
        @test string(successor(j1, mechanism1)) == string(successor(j2, mechanism2))
        @test joint_type(j1) == joint_type(j2)
    end
end

function test_urdf_serialize_deserialize(mechanism1::Mechanism; remove_fixed_tree_joints::Bool)
    state1 = MechanismState(mechanism1)

    mktempdir() do dir
        urdf2 = joinpath(dir, "test.urdf")
        write_urdf(urdf2, mechanism1, robot_name="test")
        mechanism2 = parse_urdf(urdf2, remove_fixed_tree_joints=remove_fixed_tree_joints)
        test_inverse_dynamics_match(mechanism1, mechanism2, atol=1e-10)

        for joint1 in tree_joints(mechanism1)
            if remove_fixed_tree_joints && joint_type(joint1) isa Fixed
                continue
            else
                joint2 = findjoint(mechanism2, string(joint1))
                @test position_bounds(joint1) == position_bounds(joint2)
                @test velocity_bounds(joint1) == velocity_bounds(joint2)
                @test effort_bounds(joint1) == effort_bounds(joint2)
            end
        end
    end
end

@testset "URDF write" begin
    @testset "Basics" begin
        Random.seed!(124)
        urdfdir = joinpath(@__DIR__, "urdf")
        for basename in readdir(urdfdir)
            last(splitext(basename)) == ".urdf" || continue
            urdf = joinpath(urdfdir, basename)
            for floating in (true, false)
                for remove_fixed_joints_before in (true, false)
                    mechanism = parse_urdf(urdf, remove_fixed_tree_joints=remove_fixed_joints_before, floating=floating)
                    for remove_fixed_joints_after in (true, false)
                        test_urdf_serialize_deserialize(mechanism, remove_fixed_tree_joints=remove_fixed_joints_after)
                    end
                end
            end
        end
    end

    @testset "include_root" begin
        Random.seed!(1245)
        urdfdir = joinpath(@__DIR__, "urdf")
        for basename in readdir(urdfdir)
            last(splitext(basename)) == ".urdf" || continue
            urdf = joinpath(urdfdir, basename)

            unmodified_mechanism = parse_urdf(urdf, remove_fixed_tree_joints=false)
            floating_mechanism = parse_urdf(urdf, remove_fixed_tree_joints=false, floating=true)

            # Write unmodified mechanism to URDF using `include_root=false`, then parse
            # using the default `Fixed` joint type as the root joint type. Ensure we get
            # the same kinematic tree.
            mktempdir() do dir
                urdf2 = joinpath(dir, "test.urdf")
                write_urdf(urdf2, unmodified_mechanism, include_root=false)
                unmodified_mechanism_back = parse_urdf(urdf2, remove_fixed_tree_joints=false)
                test_kinematic_graph_layout_match(unmodified_mechanism, unmodified_mechanism_back)
                test_inverse_dynamics_match(unmodified_mechanism, unmodified_mechanism_back, atol=1e-10)
            end

            # Write floating-base mechanism to URDF using `include_root=true` (default), then parse
            # using the default `Fixed` joint type as the root joint type:
            mktempdir() do dir
                urdf2 = joinpath(dir, "test.urdf")
                write_urdf(urdf2, floating_mechanism)
                floating_mechanism_back = parse_urdf(urdf2)
                # Note: floating_mechanism_back has an extra fixed joint at the root, compared to floating_mechanism,
                # but is dynamically equivalent.
                test_inverse_dynamics_match(floating_mechanism, floating_mechanism_back, atol=1e-10)
            end

            # Write floating-base mechanism to URDF using `include_root=false`, then parse and compare
            # to unmodified mechanism:
            mktempdir() do dir
                urdf2 = joinpath(dir, "test.urdf")
                write_urdf(urdf2, floating_mechanism, include_root=false)
                unmodified_mechanism_back = parse_urdf(urdf2, remove_fixed_tree_joints=false)
                test_kinematic_graph_layout_match(unmodified_mechanism, unmodified_mechanism_back)
                test_inverse_dynamics_match(unmodified_mechanism, unmodified_mechanism_back, atol=1e-10)
            end
        end
    end
end


"""
# module CAPIs

# Examples

```jldoctest
julia>
```
"""
module CAPIs

using CoordinateTransformations
using DataStructures: counter
using Setfield: @set

using JuSpglib.DataStructure: Cell

export get_symmetry, get_international, get_schoenflies,
    standardize_cell, find_primitive, refine_cell,
    niggli_reduce, delaunay_reduce

include(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))

function getfields(obj, fields...)
    Tuple(getfield(obj, name) for name in fields)
end

function get_ccell(cell::Cell)::Cell
    lattice, positions, types = getfields(cell, :lattice, :positions, :numbers)
    clattice = convert(Matrix{Cdouble}, lattice)
    cpositions = convert(Matrix{Cdouble}, positions)
    cnumbers = convert(Vector{Cint}, [repeat([i], v) for (i, v) in (enumerate ∘ values ∘ counter)(types)] |> Iterators.flatten |> collect)
    return Cell(clattice, cpositions, cnumbers)
end

cchars_to_string(s::Vector{Cchar}) = map(Char, s) |> join |> x -> split(x, "\0") |> first

function get_symmetry(cell::Cell; symprec::Real = 1e-8)
    size(cell.positions, 2) != length(cell.numbers) && throw(DimensionMismatch("The number of positions and atomic types do not match!"))
    size(cell.positions, 1) != 3 && error("Operations in 3D space is supported here!")

    maxsize = 52
    rotations = Array{Cint}(undef, 3, 3, maxsize)
    translations = Array{Cdouble}(undef, 3, maxsize)

    ccell = get_ccell(cell)
    clattice, cpositions, cnumbers = getfields(ccell, :lattice, :positions, :numbers)

    numops = ccall((:spg_get_symmetry, spglib), Cint,
        (Ptr{Cint}, Ptr{Cdouble}, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Cint, Cdouble),
        rotations, translations, maxsize, clattice, cpositions, cnumbers, length(cnumbers), symprec)
    numops == 0 && error("Could not determine symmetries!")

    [AffineMap(transpose(rotations[:, :, i]), translations[:, i]) for i in 1:numops]
end

function get_international(cell::Cell; symprec::Real = 1e-8)
    result = zeros(Cchar, 11)

    ccell = get_ccell(cell)
    clattice, cpositions, cnumbers = getfields(ccell, :lattice, :positions, :numbers)

    numops = ccall((:spg_get_international, spglib), Cint,
        (Ptr{Cchar}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Cint, Cdouble),
        result, clattice, cpositions, cnumbers, length(cnumbers), symprec)
    numops == 0 && error("Could not determine the international symbol!")

    cchars_to_string(result)
end

function get_schoenflies(cell::Cell; symprec::Real = 1e-8)
    result = zeros(Cchar, 11)

    ccell = get_ccell(cell)
    clattice, cpositions, cnumbers = getfields(ccell, :lattice, :positions, :numbers)

    numops = ccall((:spg_get_schoenflies, spglib), Cint,
        (Ptr{Cchar}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Cint, Cdouble),
        result, clattice, cpositions, cnumbers, length(cnumbers), symprec)
    numops == 0 && error("Could not determine the Schoenflies symbol!")

    cchars_to_string(result)
end

function standardize_cell(cell::Cell, to_primitive::Bool = false, no_idealize::Bool = false, symprec::Real = 1e-5)
    ccell = get_ccell(cell)
    clattice, cpositions, cnumbers = getfields(ccell, :lattice, :positions, :numbers)
    to_primitive, no_idealize = map(x -> convert(Cint, x), (to_primitive, no_idealize))

    atoms_amount = ccall((:spg_standardize_cell, spglib), Cint,
        (Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Cint, Cint, Cint, Cdouble),
        clattice, cpositions, cnumbers, length(cnumbers), to_primitive, no_idealize, symprec)
    atoms_amount == 0 && error("Standardizing cell failed!")

    Cell(clattice, cpositions, cnumbers)
end

find_primitive(cell::Cell; symprec::Real = 1e-5) = standardize_cell(cell; to_primitive = true, no_idealize = false, symprec = symprec)

refine_cell(cell::Cell; symprec::Real = 1e-5) = standardize_cell(cell; to_primitive = false, no_idealize = false, symprec = symprec)

function niggli_reduce(cell::Cell, symprec::Real = 1e-5)
    ccell = get_ccell(cell)
    clattice = getfield(ccell, :lattice)

    ret = ccall((:spg_niggli_reduce, spglib), Cint, (Ptr{Cdouble}, Cdouble), clattice, symprec)
    ret == 0 && error("Niggli reduce failed!")

    @set cell.lattice = clattice
end

function delaunay_reduce(cell::Cell, symprec::Real = 1e-5)
    ccell = get_ccell(cell)
    clattice = getfield(ccell, :lattice)

    ret = ccall((:spg_niggli_reduce, spglib), Cint, (Ptr{Cdouble}, Cdouble), clattice, symprec)
    ret == 0 && error("Delaunay reduce failed!")

    @set cell.lattice = clattice
end

function get_ir_reciprocal_mesh(cell::Cell,
                                grid::Vector{T},
                                shift::Vector{T} = [0, 0, 0],
                                is_time_reversal::Bool = true,
                                symprec::Real = 1e-5) where {T <: Integer}
    all(x -> x in (zero(T), one(T)), shift) || throw(ArgumentError("The shift can be only a vector of ones or zeros!"))

    qpoints_amount = prod(grid)
    grid_address = Array{Cint}(undef, qpoints_amount, 3)
    mapping = Array{Cint}(undef, qpoints_amount)
    ccell = get_ccell(cell)
    clattice, cpositions, cnumbers = getfields(ccell, :lattice, :positions, :numbers)

    ret = ccall((:spg_get_ir_reciprocal_mesh, spglib), Cint,
        (Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}, Cint, Cdouble),
        grid_address, mapping, grid, shift, is_time_reversal, clattice, cpositions, cnumbers, length(cnumbers), symprec)
    ret != qpoints_amount && error("Something wrong happens when finding mesh!")

    mapping, grid_address
end

function get_stabilized_reciprocal_mesh(rotations::Vector{Matrix{T}},
                                        grid::Vector{T},
                                        shift::Vector{T} = [0, 0, 0],
                                        qpoints::Vector{} = nothing,
                                        is_time_reversal::Bool = true) where {T <: Integer}
    all(x -> x in (zero(T), one(T)), shift) || throw(ArgumentError("The shift can be only a vector of ones or zeros!"))

    qpoints_amount = prod(grid)
    grid_address = Array{Cint}(undef, qpoints_amount, 3)
    mapping_table = Array{Cint}(undef, qpoints_amount)
    qpoints == nothing ? qpoints = Float64[0, 0, 0] : qpoints = Vector(qpoints)

    ret = ccall((:spg_get_stabilized_reciprocal_mesh, spglib), Cint,
        (Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Cint, Cint, Ptr{Cint}, Cint, Ptr{Cint}),
        grid_address, mapping_table, grid, shift, is_time_reversal, length(rotations), rotations, length(qpoints), qpoints)
    ret != qpoints_amount && error("Something wrong happens when finding mesh!")

    mapping_table, grid_address
end

end


import Base

"""
    Base.:/(df1::DataFrame, df2::DataFrame)
Extends / to operate on 2 DataFrames, each with one column of AbstractFloat type.
The operation will join the DataFrames on their descriptive columns to ensure the operation
is performed for related values. Values that are "missing" after joining are set to 0.
"""
function Base.:/(df1::DataFrame, df2::DataFrame)
    df = indexjoin(copy.([df1,df2]))
    out = findvalue(df)

    df[!,:value] .= df[:, out[1]] ./ df[:, out[2]]

    return df[:, unique([setdiff(propertynames(df),out);:value])]
end

"""
    Base.:+(df::Vararg{DataFrame})
Extends + to operate on 2+ DataFrames, each with one column of AbstractFloat type.
The operation will join the DataFrames on their descriptive columns to ensure the operation
is performed for related values. Values that are "missing" after joining are set to 0.
"""
function Base.:+(df::Vararg{DataFrame})
    N = length(df)

    if length(findvalue.(df)) > N
        @error("Can only add DataFrames with one value column EACH.")
    end

    df = indexjoin(ensurearray(df))
    out = findvalue(df)

    df[!,:value] .= df[:, out[1]];
    [df[!,:value] += df[:, out[ii]] for ii in 2:N]

    return df[:, unique([setdiff(propertynames(df),out);:value])]
end

"""
    Base.:-(df::Vararg{DataFrame})
Extends - to operate on 2+ DataFrames, each with one column of AbstractFloat type.
The operation will join the DataFrames on their descriptive columns to ensure the operation
is performed for related values. Values that are "missing" after joining are set to 0.
"""
function Base.:-(df::Vararg{DataFrame})
    N = length(df)

    if length(findvalue.(df)) > N
        @error("Can only subtract DataFrames with one value column EACH.")
    end

    df = indexjoin(ensurearray(df))
    out = findvalue(df)

    df[!,:value] .= df[:, out[1]];
    [df[!,:value] -= df[:, out[ii]] for ii in 2:N]

    return df[:, unique([setdiff(propertynames(df),out);:value])]
end

"""
    Base.:*(df::Vararg{DataFrame})
Extends * to operate on 2+ DataFrames, each with one column of AbstractFloat type.
The operation will join the DataFrames on their descriptive columns to ensure the operation
is performed for related values. Values that are "missing" after joining are set to 0.
"""
function Base.:*(df::Vararg{DataFrame})
    N = length(df)

    if length(findvalue.(df)) > N
        @error("Can only multiply DataFrames with one value column EACH.")
    end

    df = indexjoin(ensurearray(df))
    out = findvalue(df)

    df[!,:value] .= df[:, out[1]];
    [df[!,:value] .*= df[:, out[ii]] for ii in 2:N]

    return df[:, unique([setdiff(propertynames(df),out);:value])]
end

"""
    combine_over(df::DataFrame, col::Array{Symbol,1}; operation::Function = sum)
    combine_over(df::DataFrame, col::Symbol; operation::Function = sum)
This function applies [`combine`](https://juliadata.github.io/DataFrames.jl/stable/lib/functions/#DataFrames.combine)
to the input DataFrame `df` over the input column(s) `col`.

# Arguments
- `df::DataFrame`: DataFrame on which to operate.
- `col::Symbol` or `col::Array{Symbol,1}`: column(s) over which to operate.

# Keywords
- `operation::Function = sum`: Operation to perform over the DataFrame columns. By default,
    the function will return a summation. Other standard summary functions include: `sum`,
    `prod`, `minimum`, `maximum`, `mean`, `var`, `std`, `first`, `last` and `length`.

# Returns
- `df::DataFrame` WITHOUT the specified column(s) argument. The resulting DataFrame will be
    'shorter' than the input DataFrame.
"""
function combine_over(df::DataFrame, col::Array{Symbol,1}; fun::Function = sum)
    # !!!! add kwarg to findvalue to indicate whether to include integers as values
    val = findvalue(df)
    idx_by = setdiff(propertynames(df), [col; val])

    df = combine(groupby(df, idx_by), val .=> fun .=> val)

    [df[!,ii] .= round.(df[:,ii]; digits = DEFAULT_ROUND_DIGITS)
        for ii in find_oftype(df[:,val], AbstractFloat)]

    # !!!! See where we actually want to convert boolean sums to integers. I think it's just
    # in some labor sharing. We can probably keep summed booleans as integers. This seems less confusing.
    [df[!,ii] .= convert_type.(Float64, df[:,ii]) for ii in find_oftype(df[:,val], Int)]
    return df
end

function combine_over(df::DataFrame, col::Symbol; fun::Function = sum)
    return combine_over(df, ensurearray(col); fun = fun)
end

"""
    transform_over(df::DataFrame, col::Array{Symbol,1}; operation::Function = sum)
    transform_over(df::DataFrame, col::Symbol; operation::Function = sum)
This function applies [`transform`](https://juliadata.github.io/DataFrames.jl/stable/lib/functions/#DataFrames.transform)
to the input DataFrame `df` over the input column(s) `col`.

# Arguments
- `df::DataFrame`: DataFrame on which to operate.
- `col::Symbol` or `col::Array{Symbol,1}`: column(s) over which to operate.

# Keywords
- `operation::Function = sum`: Operation to perform over the DataFrame columns. By default,
    the function will return a summation. Other standard summary functions include: `sum`,
    `prod`, `minimum`, `maximum`, `mean`, `var`, `std`, `first`, `last` and `length`.

# Returns
- `df::DataFrame` WITH the specified column(s) argument. The resulting DataFrame will be
    the same length as the input DataFrame.
"""
function transform_over(df::DataFrame, col::Array{Symbol,1}; fun::Function = sum)
    val = findvalue(df)
    idx_by = setdiff(propertynames(df), [col; val])

    df = transform(groupby(df, idx_by), val .=> fun .=> val)

    [df[!,ii] .= round.(df[:,ii]; digits = DEFAULT_ROUND_DIGITS)
        for ii in find_oftype(df[:,val], AbstractFloat)]
    [df[!,ii] .= convert_type.(Float64, df[:,ii]) for ii in find_oftype(df[:,val], Int)]
    return df
end

function transform_over(df::DataFrame, col::Symbol; fun::Function = sum)
    return transform_over(df, ensurearray(col); fun = fun)
end

"""
Round either the specified columns or all columns of type `AbstractFloat` to the specified
number of digits.

# Arguments
- `df::DataFrame` of values in need of rounding
- `col::Symbol` or `col::Array{Symbol,1}` to round. If no columns are specified,
    all columns of type `AbstractFloat` will be rounded.

# Keywords
- `digits::Int = 10`: Number of decimal places to keep when rounding
"""
function round!(df::DataFrame; digits::Int = DEFAULT_ROUND_DIGITS)
    return round!(df, find_oftype(df, AbstractFloat); digits = digits)
end

function round!(df::DataFrame, col::Union{Symbol,Array{Symbol,1}}; digits::Int = DEFAULT_ROUND_DIGITS)
    df[!,col] .= round.(df[:,col]; digits = digits)
    return df
end


using Test
using ScottishTaxBenefitModel
using .Utils
using .Definitions

@enum A a b c d e

@testset "tests of miscellaneous functions" begin
    gr = Dict{A,Real}( a=>1.0, b=>2, e=> 4.1 )
    ne = Dict{A,Real}( c=>2, b=>9 )
    inc = Dict{A,Real}( a=>1, b=>1.0, c=>1, d =>1 )
    m = mult( data=gr, calculated=ne, included=inc )
    @test m == 12

    pid = 120181906404
    p = from_pid( pid )
    @test p.datasource == FRS
    @test p.year == 2018
    @test p.hid == 19064
    @test p.pno == 4

end


@testset "Parameter" begin
    p = GenericPara(diagType = Ver4Diag, innerLoopNum = 1)
    q = GenericPara(diagType = Ver4Diag, innerLoopNum = 2)
    a = GenericPara(diagType = Ver4Diag, innerLoopNum = 2)

    @test p != q
    @test q == a

    aa = reconstruct(a, transferLoop = [0.0, 0.0, 0.0])
    @test a != aa

    # reconstruct with the same diagType but different interaction
    aaa = reconstruct(a, interaction = [])
    @test a != aaa

    # reconstruct with different diagram type leads to different parameter
    #reconstructed GenericPara uses the old parameters such as firstLoopIdx, totalLoopNum etc., which are different between diagTypes
    s = reconstruct(a, diagType = SigmaDiag)
    ss = GenericPara(diagType = SigmaDiag, innerLoopNum = 2)
    @test s != ss

end


export getCharge, ParamList

const ElementNames =
[
"H",
"He",
"Li",
"Be",
"B",
"C",
"N",
"O",
"F",
"Ne",
"Na",
"Mg",
"Al",
"Si",
"P",
"S",
"Cl",
"Ar",
"K",
"Ca"
]

const SubshellNames =
[
"S",
"P",
"D",
"F",
"G",
"H",
"I"
]

const SubshellAngularMomentums =
[
["X⁰Y⁰Z⁰"],
["X¹Y⁰Z⁰", "X⁰Y¹Z⁰", "X⁰Y⁰Z¹"],
["X²Y⁰Z⁰", "X¹Y¹Z⁰", "X¹Y⁰Z¹", "X⁰Y²Z⁰", "X⁰Y¹Z¹", "X⁰Y⁰Z²"],
["X³Y⁰Z⁰", "X²Y¹Z⁰", "X²Y⁰Z¹", "X¹Y²Z⁰", "X¹Y¹Z¹", "X¹Y⁰Z²", "X⁰Y³Z⁰", "X⁰Y²Z¹", "X⁰Y¹Z²", "X⁰Y⁰Z³"],
["X⁴Y⁰Z⁰", "X³Y¹Z⁰", "X³Y⁰Z¹", "X²Y²Z⁰", "X²Y¹Z¹", "X²Y⁰Z²", "X¹Y³Z⁰", "X¹Y²Z¹", "X¹Y¹Z²", "X¹Y⁰Z³", "X⁰Y⁴Z⁰", "X⁰Y³Z¹", "X⁰Y²Z²", "X⁰Y¹Z³", "X⁰Y⁰Z⁴"],
["X⁵Y⁰Z⁰", "X⁴Y¹Z⁰", "X⁴Y⁰Z¹", "X³Y²Z⁰", "X³Y¹Z¹", "X³Y⁰Z²", "X²Y³Z⁰", "X²Y²Z¹", "X²Y¹Z²", "X²Y⁰Z³", "X¹Y⁴Z⁰", "X¹Y³Z¹", "X¹Y²Z²", "X¹Y¹Z³", "X¹Y⁰Z⁴", "X⁰Y⁵Z⁰", "X⁰Y⁴Z¹", "X⁰Y³Z²", "X⁰Y²Z³", "X⁰Y¹Z⁴", "X⁰Y⁰Z⁵"],
["X⁶Y⁰Z⁰", "X⁵Y¹Z⁰", "X⁵Y⁰Z¹", "X⁴Y²Z⁰", "X⁴Y¹Z¹", "X⁴Y⁰Z²", "X³Y³Z⁰", "X³Y²Z¹", "X³Y¹Z²", "X³Y⁰Z³", "X²Y⁴Z⁰", "X²Y³Z¹", "X²Y²Z²", "X²Y¹Z³", "X²Y⁰Z⁴", "X¹Y⁵Z⁰", "X¹Y⁴Z¹", "X¹Y³Z²", "X¹Y²Z³", "X¹Y¹Z⁴", "X¹Y⁰Z⁵", "X⁰Y⁶Z⁰", "X⁰Y⁵Z¹", "X⁰Y⁴Z²", "X⁰Y³Z³", "X⁰Y²Z⁴", "X⁰Y¹Z⁵", "X⁰Y⁰Z⁶"]
]

const SubshellOrbitals =
[ # Every XYZs must start with [l, 0, 0]
[[0,0,0]],
[[1,0,0], [0,1,0], [0,0,1]],
[[2,0,0], [1,1,0], [1,0,1], [0,2,0], [0,1,1], [0,0,2]],
[[3,0,0], [2,1,0], [2,0,1], [1,2,0], [1,1,1], [1,0,2], [0,3,0], [0,2,1], [0,1,2], [0,0,3]],
[[4,0,0], [3,1,0], [3,0,1], [2,2,0], [2,1,1], [2,0,2], [1,3,0], [1,2,1], [1,1,2], [1,0,3], [0,4,0], [0,3,1], [0,2,2], [0,1,3], [0,0,4]],
[[5,0,0], [4,1,0], [4,0,1], [3,2,0], [3,1,1], [3,0,2], [2,3,0], [2,2,1], [2,1,2], [2,0,3], [1,4,0], [1,3,1], [1,2,2], [1,1,3], [1,0,4], [0,5,0], [0,4,1], [0,3,2], [0,2,3], [0,1,4], [0,0,5]],
[[6,0,0], [5,1,0], [5,0,1], [4,2,0], [4,1,1], [4,0,2], [3,3,0], [3,2,1], [3,1,2], [3,0,3], [2,4,0], [2,3,1], [2,2,2], [2,1,3], [2,0,4], [1,5,0], [1,4,1], [1,3,2], [1,2,3], [1,1,4], [1,0,5], [0,6,0], [0,5,1], [0,4,2], [0,3,3], [0,2,4], [0,1,5], [0,0,6]]
]

const BasisFuncNames =
[
"STO-2G",
"STO-3G",
"STO-6G",
"3-21G",
"6-31G",
"cc-pVDZ",
"cc-pVTZ",
"cc-pVQZ"
]

const SciNotMarker = r"D(?=[\+\-])"
const sciNotReplace = (txt)->replace(txt, SciNotMarker => "e")
const BStextEndingMarker = "****"
const BasisSetList = Dict(BasisFuncNames .=> BasisFuncTexts)
const AtomicNumberList = Dict(ElementNames .=> collect(1 : length(ElementNames)))
const AngularMomentumList = Dict(SubshellNames .=> collect(0 : length(SubshellNames)-1))
const SubshellSuborderList = Dict(SubshellNames .=> SubshellOrbitals)
const ijkOrderList = Dict(SubshellNames .=> SubshellAngularMomentums)
const ijkIndexList = Dict(flatten(SubshellAngularMomentums) .=> flatten([collect(1:length(i)) for i in SubshellAngularMomentums]))
const ijkOrbitalList = Dict(flatten(SubshellAngularMomentums)  .=> flatten(SubshellOrbitals))
const ijkStringList = Dict(flatten(SubshellOrbitals) .=> flatten(SubshellAngularMomentums))
const SubshellDimList = Dict(SubshellNames .=> length.(SubshellOrbitals))
# const ParamNames = [:𝑋, :𝑌, :𝑍, :𝑑, :𝛼, :𝐿]
const ParamNames = [:X, :Y, :Z, :d, :α, :L]
const ParamSymbols = [:X, :Y, :Z, :con, :xpn, :spacing]
const ParamList = Dict(ParamSymbols .=> ParamNames)

getCharge(nucs::Vector{String}) = getCharge.(nucs) |> sum

getCharge(nucStr::String) = AtomicNumberList[nucStr]

function checkBSList(;printInfo::Bool=false)
    texts = vcat(BasisFuncTexts...)
    for text in texts
        if text !== nothing
            sBool = startswith(text, r"[A-Z][a-z]?     [0-9]\n")
            eBool = endswith(text, BStextEndingMarker*"\n")
            @assert (sBool && eBool) """\n
            The format of 'Basis functions' is NOT CORRECT!
            "[A-Z][a-z]?     [0-9]\\n" (regex) should exit as the 1st line of the string!
            "$(BStextEndingMarker)\\n" should exit as the last line of the string!
            The incorrect text content is:\n$(text)
            """
        end
    end
    printInfo && println("Basis function list checked.")
end


function playEpisode(env::AbsEnvironment, agent::AbsAgent; verbose=false, randomInitial = false, learn=true, threshold=100000, render=false)
	state = randomInitial ? rand(getAllStates(env)) : getInitialState(env)
	totalRewards = 0.0
	numOfStates = 1

	verbose && println("Initial State: $(state)")

	while !isTerminal(state, env) && numOfStates < threshold
		render && render_env(env)
		action = play(agent, state, env; learn=learn)
		verbose && println("Action: $(action)")
		state, reward = transfer(env, state, action)
		verbose && println("State: $(state)")
		verbose && println("Reward: $(reward)")
		observe(agent, state, reward, env; learn=learn, terminal = (numOfStates + 1 == threshold || isTerminal(state, env)))
		totalRewards += reward
		numOfStates += 1
	end

	return (totalRewards, numOfStates)
end


# ------- IPOPT ------------

"""
    IPOPT(options=Dict())

Use Ipopt as the optimizer

# Arguments
- `options::Dict`: options for Ipopt.  see Ipopt docs.
"""
struct IPOPT{TD} <: AbstractSolver
    options::TD
end

IPOPT() = IPOPT(Dict())

# wrapper for IPOPT
function optimize(solver::IPOPT, cache, x0, lx, ux, lg, ug, rows, cols)

    # initialize
    nx = length(x0)
    ng = length(lg)
    options = solver.options

    # initialize data for caching results since ipopt separates functions out
    xlast = 2*x0
    f = 0.0
    g = zeros(ng)
    df = zeros(nx)
    dg = zeros(length(rows))

    function ipcon(x, g)  # ipopt calls this function first

        f = evaluate!(g, df, dg, x, cache)
        xlast = x

        return nothing
    end

    function ipobj(x)
        if !isequal(x, xlast)
            f = func!(g, x)
        end
        return f
    end

    function ipgrad_obj(x, grad)
        if !isequal(x, xlast)
            f = evaluate!(g, df, dg, x, cache)
            xlast = x
        end
        grad[:] = df
        return nothing
    end

    function ipjac_con(x, mode, r, c, values)
        if mode == :Structure
            r[:] = rows
            c[:] = cols
        else
            if !isequal(x, xlast)
                f = evaluate!(g, df, dg, x, cache)
                xlast = x
            end
            values[:] = dg
        end
        return nothing
    end

    nzJ = length(rows)
    nH = 1  # irrelevant for quasi-newton
    prob = Ipopt.createProblem(nx, lx, ux, ng, lg, ug, nzJ, nH,
        ipobj, ipcon, ipgrad_obj, ipjac_con)

    # set options
    Ipopt.addOption(prob, "hessian_approximation", "limited-memory")
    for (key, value) in options
        Ipopt.addOption(prob, key, value)
    end

    # open output file
    filename = "ipopt.out"
    print_level = 5
    if haskey(options, "output_file")
        filename = options["output_file"]
    end
    if haskey(options, "print_level")
        print_level = options["print_level"]
    end
    Ipopt.openOutputFile(prob, filename, print_level)

    # solve problem
    prob.x = x0
    status = solveProblem(prob)

    return prob.x, prob.obj_val, Ipopt.ApplicationReturnStatus[status], nothing
end


# # Turn to Gold

# ### Turn blocks to gold by hitting them
# Demonstration of pollBlockHits()
# the function pollBlockHits() returns an Array of hit Events
# A hit event is a tuple, say `((5, 4, 5), 2, 17)`
# `(5, 4, 5)` are the coordinates to the block which was hit
# `2` is the face of the block which was hit
# `17` is the entityId of the Entity which was responsible for the hit

using Dates
using PiCraft

clearEvents() # Clear all previous events
t = now() # Mark the current time

post("60 seconds to turn blocks to gold.")

while (now() - t).value/1000 < 60.0 # Continue the loop if 60 seconds has not elapsed
    for hit in pollBlockHits()
        setBlock(hit[1], Block(41)) # hit[1] are the coordinates to the Block which was hit
    end
end

# Post a message to screen when done
post("Time Over!")


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

@recipe function f(domain::AbstractDomain, data::AbstractVector) where {N,T}
  @series begin
    PointSet(coordinates(domain)), data
  end
end

@recipe function f(domain::AbstractDomain)
  @series begin
    PointSet(coordinates(domain))
  end
end


module PlotUtils

# import ..TabularTMazeUtils
using FileIO
using JLD2

function build_tabular_tmaze_heatmap(itm)
    values = load_result(itm, :value_map)
    map_set = @load "../src/data/TTMazeValueSet.jld2" ValueSet

    num_steps = size(values)[1]

    value_heatmap = zeros(num_steps, 9, 7)

    # position_mapping = TabularTMazeUtils.valid_state_mask()
    function valid_state_mask()
        """
        return a mask of valid states that is 9x7
        """
        world = [["G1", "0", "0", "0", "0", "0", "G3"],
                          ["1", "0", "0", "0", "0", "0", "1"],
                          ["1", "1", "1", "1", "1", "1", "1"],
                          ["1", "0", "0", "1", "0", "0", "1"],
                          ["G2", "0", "0", "1", "0", "0", "G4"],
                          ["0", "0", "0", "1", "0", "0", "0"],
                          ["0", "0", "0", "1", "0", "0", "0"],
                          ["0", "0", "0", "1", "0", "0", "0"],
                          ["0", "0", "0", "1", "0", "0", "0"]]

        world = permutedims(hcat(world...))
        valid_states = findall(x-> x!="0", world)
        return valid_states
    end
    position_mapping = valid_state_mask()

    for i in 1:num_steps

        for (set_ind,obs) in enumerate(ValueSet["states"])
            heatmap_ind = position_mapping[Int(obs[1])]
            value_heatmap[i,heatmap_ind] = values[i,set_ind,1]
        end
    end
    return value_heatmap
end

"""
    Given an ItemCollection, returns all logs at the logger_key
"""
function load_results(ic, logger_key; return_type = "tensor")
    num_results = length(ic)
    results = []
    for itm in ic.items
        data = FileIO.load(joinpath(itm.folder_str, "results.jld2"))["results"]
        push!(results,data[logger_key])
    end

    if return_type == "tensor"
        return cat(results..., dims = 3)
    elseif return_type == "array"
        return results
    end
end

function load_result(itm, logger_key)
    data = FileIO.load(joinpath(itm.folder_str, "results.jld2"))["results"]
    return data[logger_key]
end

end