function talusplot(σ2::Vector, k=1:length(σ2)-1, args...; kwargs...)
    talus = -diff(log2.(σ2))/2 # log₂(σₖ)-log₂(σₖ₊₁)

    df = DataFrame(k=k, talus=talus[k])

    plot(layer(df, x=:k, y=:talus, Geom.line()),
         layer(df, x=:k, y=:talus, Geom.point()),
         Guide.xlabel("k (Principal Component)"), Guide.ylabel("log₂(σₖ)-log₂(σₖ₊₁)"),
         args...; kwargs...)
end

function talusplot(X::Matrix, k=[], args...; kwargs...)
    K = size(X,2)<=size(X,1) ? Symmetric(X'X) : Symmetric(X*X') # Work with the smaller of the kernel matrices

    # compute σᵢ² and talus plot
    σ2 = eigvals(K)
    sort!(σ2,rev=true) # largest to smallest
    σ2 = σ2[1:findlast(σ2.>1e-9)] # get rid of eigenvalues at the end that might be 0 (or below, due to numerical problems)

    talusplot(σ2, isempty(k) ? (1:length(σ2)-1) : k, args...; kwargs...)
end
talusplot(X::Matrix, k::Integer, args...; kwargs...) = talusplot(X, 1:k, args...; kwargs...)


## Test pvalue adjustment methods ##
module Test_pval_pi0_adjustment

using MultipleTesting
using Test

@testset "p-Value π₀ adjustment" begin

    pval1 = [0.0, 0.0001, 0.001, 0.01, 0.05, 0.1, 0.2, 0.4, 0.75, 1.0]
    pi0 = 0.4

    @test_throws MethodError adjust()
    @test_throws MethodError adjust(pval1)

    @testset "benjamini_hochberg" begin

        t = BenjaminiHochbergAdaptive

        ref0 = [0.0, 0.0005, 0.003333333, 0.025, 0.1, 0.166666667, 0.285714286, 0.5, 0.833333333, 1.0]
        ref = ref0 .* pi0

        ## no integers as input
        @test_throws MethodError adjust([0, 1], t())

        ## no valid p-values as input
        @test_throws DomainError adjust([-1.0, 0.5], t(pi0))
        @test_throws DomainError adjust([0.5, 1.5], t(pi0))
        @test_throws DomainError adjust([0.5, 0.7], t(-1.0))
        @test_throws DomainError adjust([0.5, 0.7], t(1.5))

        ## single p-value is returned unchanged
        pval = rand(1)
        @test adjust(pval, t(pi0)) == pval .* pi0

        ## compare with reference values
        @test isapprox(adjust(pval1, t(pi0)), ref, atol = 1e-8)

        ## BH Adaptive same as BH for π₀ missing or 1
        @test isapprox(adjust(pval1, t(1.0)), ref0, atol = 1e-8)
        @test isapprox(adjust(pval1, t()), ref0, atol = 1e-8)
        @test isapprox(adjust(pval1, t(0.0)), zero(ref0), atol = 1e-8)

        # adaptive with pi0 estimator == oracle with estimated pi0
        @test adjust(pval1, t(Storey())) == adjust(pval1, t(estimate(pval1, Storey())))
        @test adjust(pval1, t(Storey(0.3))) == adjust(pval1, t(estimate(pval1, Storey(0.3))))
        @test adjust(pval1, t(LeastSlope())) == adjust(pval1, t(estimate(pval1, LeastSlope())))

    end

end

end


using ClutteredEnvPathOpt
using LightGraphs
import Polyhedra
import GLPK
using Pipe
using Plots
using JuMP, Gurobi, PiecewiseLinearOpt

function get_M_A_b_easy(obstacles)
    Ms = []
    As = []
    bs = []
    acc = [0] # accumulator # try using reduce() instead

    _, points, _, _, free_faces = ClutteredEnvPathOpt.construct_graph(obstacles)

    for face in free_faces
        v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)
        polygon = Polyhedra.polyhedron(
            v,
            Polyhedra.DefaultLibrary{Rational{Int64}}(Gurobi.Optimizer)
        )
        halfspaces = Polyhedra.hrep(polygon).halfspaces
        temp = acc[end]
        push!(acc, temp + length(halfspaces))

        # for i in v.points
        #     println("$i $face")
        # end

        for hs in halfspaces
            A = hs.a'
            b = hs.β

            M = maximum(map(x -> A * x, v.points))

            push!(Ms, M)
            push!(As, A)
            push!(bs, b)
        end
    end
    # acc = acc[2:end]

    return vcat(Ms...), vcat(As...), vcat(bs...), acc
end

function get_M_A_b(points, free_faces)
    Ms = []
    As = []
    bs = []
    acc = [0] # try using reduce() instead

    # _, points, _, _, free_faces = ClutteredEnvPathOpt.construct_graph(obstacles)

    # Solve over the unitcell
    u = Polyhedra.convexhull([0//1,0//1],[0//1,1//1],[1//1,0//1],[1//1,1//1])
    unitcell = Polyhedra.polyhedron(u)

    for face in free_faces
        v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)
        polygon = Polyhedra.polyhedron(
            v,
            Polyhedra.DefaultLibrary{Rational{Int64}}(Gurobi.Optimizer)
        )
        halfspaces = Polyhedra.hrep(polygon).halfspaces
        temp = acc[end]
        push!(acc, temp + length(halfspaces))

        for hs in halfspaces
            A = hs.a'
            b = hs.β

            sub_model = JuMP.Model(Gurobi.Optimizer)
            #sub_x = JuMP.@variable(sub_model, [1:2])
            JuMP.@variable(sub_model, sub_x[1:2])
            JuMP.@constraint(sub_model, sub_x in unitcell)
            #JuMP.@constraint(sub_model, sub_x in safe_regions)
            JuMP.@objective(sub_model, Max, A * sub_x)

            JuMP.optimize!(sub_model)

            #return value.(sub_x)
            M = JuMP.objective_value(sub_model)

            push!(Ms, M)
            push!(As, A)
            push!(bs, b)
        end

        # A = vcat([hs.a for hs in halfspaces]...)
        # b = vcat([hs.β for hs in halfspaces]...)

        # M = map(hs ->
        #     begin
        #         A = hs.a
        #         b = hs.β

        #         sub_model = JuMP.Model(Gurobi.Optimizer)
        #         #sub_x = JuMP.@variable(sub_model, [1:2])
        #         JuMP.@variable(sub_model, sub_x[1:2])
        #         JuMP.@constraint(sub_model, sub_x in polygon)   # TODO "no method matching copy" error WHY DOES THIS ONLY WORK THE FIRST TIME
        #         JuMP.@objective(sub_model, Max, A' * sub_x)

        #         stat = JuMP.optimize!(sub_model)

        #         #return value.(sub_x)
        #         return JuMP.objective_value(sub_model)
        #     end,
        #     halfspaces
        # )

        # push!(Ms, M)
        # push!(As, A)
        # push!(bs, b)
    end
    # acc = acc[2:end]

    return vcat(Ms...), vcat(As...), vcat(bs...), acc
    #return Ms, As, bs, acc
end

# obstacles <- obstacles (list of polyhedra)
# N <- max number of steps (scalar)
# f1 <- initial left foot ([x, y, theta])
# f2 <- initial right foot ([x, y, theta])
# g <- goal pose
# Q_g <- pose weights to goal (4x4 mat)
# Q_r <- pose weights between steps (4x4 mat)
# q_t <- weight on unused steps (scalar)
# method <- "merged" for the compact biclique cover, "full" for the original biclique cover, "bigM" for big-M constraints
# d1 = 0.2 <- radius of reference foot circle
# d2 = 0.2 <- radius of moving foot circle
# p1 = [0, 0.07] <- center of reference foot circle
# p2 = [0, -0.27] <- center of moving foot circle
# delta_x_y_max = 0.1 <- max stride norm in space
# delta_θ_max = pi/4 <- max difference in θ
# L = 5 <- number of pieces of pwl sin/cos (scalar)
function solve_deits(obstacles, N, f1, f2, g, Q_g, Q_r, q_t; method="merged", d1=0.2, d2=0.2, p1=[0, 0.07], p2=[0, -0.27])

    _, points, graph, _, free_faces = ClutteredEnvPathOpt.construct_graph(obstacles)
    skeleton = LabeledGraph(graph)

    # model = JuMP.Model(JuMP.optimizer_with_attributes(Gurobi.Optimizer))
    model = JuMP.Model(JuMP.optimizer_with_attributes(Gurobi.Optimizer, "MIPGap" => .01, "TimeLimit" => 180))
    # model = JuMP.Model(JuMP.optimizer_with_attributes(Gurobi.Optimizer, "Heuristics"=> 0, "Cuts"=> 0, "Precrush"=>1, "MIPGap" => .01, "TimeLimit" => 180))

    # model has scalar variables x, y, θ, binary variable t
    JuMP.@variable(model, x[1:N])
    JuMP.@variable(model, y[1:N])
    JuMP.@variable(model, θ[1:N])
    JuMP.@variable(model, t[1:N], Bin)

    # Objective
    f = vec([[x[j], y[j], θ[j]] for j in 1:N])
    # f_no_theta = vec([[x[j], y[j]] for j in 1:N])
    JuMP.@objective(
        model,
        Min,
        ((f[N] - g)' * Q_g * (f[N] - g)) + sum(q_t * t) + sum((f[j + 1] - f[j])' * Q_r * (f[j + 1] - f[j]) for j in 1:(N-1))
    )

    cover = Set{Pair{Set{Int64}, Set{Int64}}}()
    merged_cover = Set{Pair{Set{Int64}, Set{Int64}}}()
    num_free_face_ineq = 0

    if method != "bigM"
        cover = ClutteredEnvPathOpt.find_biclique_cover(skeleton, free_faces)
        feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces))
        merged_cover = ClutteredEnvPathOpt.biclique_merger(cover, feg)

        (valid_cover, _, _, _) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, cover)
        (valid_cover_merged, _, _, _) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, merged_cover)
    end

    # Footstep location constraints
    if method == "merged" && valid_cover_merged
        J = LightGraphs.nv(skeleton.graph)
        JuMP.@variable(model, λ[1:N, 1:J] >= 0)
        JuMP.@variable(model, z[1:N, 1:length(merged_cover)], Bin)
        for i in 1:N
            for (j,(A,B)) in enumerate(merged_cover)
                JuMP.@constraint(model, sum(λ[i, v] for v in A) <= z[i, j])
                JuMP.@constraint(model, sum(λ[i, v] for v in B) <= 1 - z[i, j])
            end
            JuMP.@constraint(model, sum(λ[i,:]) == 1)
            JuMP.@constraint(model, x[i] == sum(points[j].first * λ[i, j] for j in 1:J))
            JuMP.@constraint(model, y[i] == sum(points[j].second * λ[i, j] for j in 1:J))
        end
    elseif method != "bigM" && valid_cover
        J = LightGraphs.nv(skeleton.graph)
        JuMP.@variable(model, λ[1:N, 1:J] >= 0)
        JuMP.@variable(model, z[1:N, 1:length(cover)], Bin)
        for i in 1:N
            for (j,(A,B)) in enumerate(cover)
                JuMP.@constraint(model, sum(λ[i, v] for v in A) <= z[i, j])
                JuMP.@constraint(model, sum(λ[i, v] for v in B) <= 1 - z[i, j])
            end
            JuMP.@constraint(model, sum(λ[i,:]) == 1)
            JuMP.@constraint(model, x[i] == sum(points[j].first * λ[i, j] for j in 1:J))
            JuMP.@constraint(model, y[i] == sum(points[j].second * λ[i, j] for j in 1:J))
        end
    else
        # Runs if method = "bigM", or if a valid biclique cover is not found
        M, A, b, acc = get_M_A_b(points, free_faces)
        num_free_face_ineq = length(M)
        JuMP.@variable(model, z[1:N, 1:length(free_faces)], Bin)
        for j in 1:N
            for r in 1:length(free_faces)
                ids = (acc[r]+1):(acc[r+1])
                for i in ids
                    # JuMP.@constraint(model, A[r, :]' * [x[j], y[j]] <= b[r] * z[j, r] + M[r] * (1 - z[j, r]))
                    # JuMP.@constraint(model, A[r, :]' * [x[j], y[j]] <= b[r] * z[j, r] + 1 * (1 - z[j, r]))
                    # println("$i")
                    JuMP.@constraint(model, A[i, :]' * [x[j]; y[j]] <= b[i] * z[j, r] + M[i] * (1 - z[j, r]))
                end
            end
            JuMP.@constraint(model, sum(z[j, :]) == 1)
        end
    end

    # Reachability
    # s = [piecewiselinear(model, θ[j], range(0, stop=(2 * pi), length=L), sin) for j in 1:N]
    # c = [piecewiselinear(model, θ[j], range(0, stop=(2 * pi), length=L), cos) for j in 1:N]
    s_break_pts = [0, 5pi/16, 11pi/16, 21pi/16, 27pi/16, 2pi]
    c_break_pts = [0, 3pi/16, 13pi/16, 19pi/16, 29pi/16, 2pi]
    s = [piecewiselinear(model, θ[j], s_break_pts, sin) for j in 1:N]
    c = [piecewiselinear(model, θ[j], c_break_pts, cos) for j in 1:N]

    # For footstep j, the cirles are created from the frame of reference of footstep j-1
    # Use negative if in frame of reference of right foot (j = 1 is left foot)
    # Position is with respect to θ = 0
    for j in 2:N
        if j % 2 == 1
            # In frame of reference of right foot
            JuMP.@constraint(
                model,
                [
                    d1,
                    x[j] - x[j - 1] - c[j-1] * (-p1[1]) + s[j-1] * (-p1[2]),
                    y[j] - y[j - 1] - s[j-1] * (-p1[1]) - c[j-1] * (-p1[2])
                ] in SecondOrderCone()
            )

            JuMP.@constraint(
                model,
                [
                    d2,
                    x[j] - x[j - 1] - c[j-1] * (-p2[1]) + s[j-1] * (-p2[2]),
                    y[j] - y[j - 1] - s[j-1] * (-p2[1]) - c[j-1] * (-p2[2])
                ] in SecondOrderCone()
            )
        else
            # In frame of reference left foot
            JuMP.@constraint(
                model,
                [
                    d1,
                    x[j] - x[j - 1] - c[j-1] * p1[1] + s[j-1] * p1[2],
                    y[j] - y[j - 1] - s[j-1] * p1[1] - c[j-1] * p1[2]
                ] in SecondOrderCone()
            )

            JuMP.@constraint(
                model,
                [
                    d2,
                    x[j] - x[j - 1] - c[j-1] * p2[1] + s[j-1] * p2[2],
                    y[j] - y[j - 1] - s[j-1] * p2[1] - c[j-1] * p2[2]
                ] in SecondOrderCone()
            )
        end
    end

    # Set T to punish extra steps
    for j in 1:N
        if j % 2 == 1
            JuMP.@constraint(model, t[j] => {x[j] == f1[1]}) # use f1 and f2?
            JuMP.@constraint(model, t[j] => {y[j] == f1[2]})
            JuMP.@constraint(model, t[j] => {θ[j] == f1[3]})
        else
            JuMP.@constraint(model, t[j] => {x[j] == f2[1]})
            JuMP.@constraint(model, t[j] => {y[j] == f2[2]})
            JuMP.@constraint(model, t[j] => {θ[j] == f2[3]})
        end
    end

    # Max step distance
    # for j in 3:2:(N-1)
    #     JuMP.@constraint(model, [delta_x_y_max; f_no_theta[j] - f_no_theta[j - 2]] in SecondOrderCone())
    #     JuMP.@constraint(model, [delta_x_y_max; f_no_theta[j + 1] - f_no_theta[j - 1]] in SecondOrderCone())
    # end

    # Max theta difference for individual feet
    # for j in 3:2:(N-1)
    #     JuMP.@constraint(model, -delta_θ_max <= θ[j] - θ[j - 2] <= delta_θ_max)
    #     JuMP.@constraint(model, -delta_θ_max <= θ[j + 1] - θ[j - 1] <= delta_θ_max)
    # end

    # Max θ difference between feet
    for j in 2:N
        JuMP.@constraint(model, -pi/8 <= θ[j] - θ[j-1] <= pi/8)
    end

    # Initial footstep positions
    JuMP.@constraint(model, f[1] .== f1)
    JuMP.@constraint(model, f[2] .== f2)

    # Solve
    JuMP.optimize!(model)

    if method == "merged" && valid_cover_merged
        println("\n\nUsed merged cover.\n\n")
    elseif method != "bigM" && valid_cover
        println("\n\nUsed full cover.\n\n")
        method = "full"
    else
        println("\n\nUsed big-M constraints.\n\n")
        method = "bigM"
    end

    stats = (termination_status(model), solve_time(model), relative_gap(model), simplex_iterations(model),
            node_count(model), LightGraphs.nv(skeleton.graph), length(merged_cover), length(cover), length(free_faces), num_free_face_ineq, method)

    return value.(x), value.(y), value.(θ), value.(t), stats
end

function plot_steps(obstacles, x, y, θ)
    plot()
    ClutteredEnvPathOpt.plot_field(obstacles);
    scatter!(x[1:2:end], y[1:2:end], color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 1:2:length(x)]));
    scatter!(x[2:2:end], y[2:2:end], color="blue", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 2:2:length(x)]));
    quiver!(x, y, quiver=(0.075 * cos.(θ), 0.075 * sin.(θ)))
    display(plot!(title="Footsteps"))
end

function plot_circles(x, y, theta; R1=0.20, R2=0.20, p1=[0, 0.07], p2=[0, -0.27])
    angles = LinRange(0,2*pi, 100)
    #plot()
    for i = 2:length(x)
        plot()
        if i % 2 == 1
            circlex1 = R1*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(-p1[1]) - sin(theta[i-1])*(-p1[2]) )
            circley1 = R1*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(-p1[1]) + cos(theta[i-1])*(-p1[2]) )

            circlex2 = R2*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(-p2[1]) - sin(theta[i-1])*(-p2[2]) )
            circley2 = R2*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(-p2[1]) + cos(theta[i-1])*(-p2[2]) )
            #scatter!([x[i-1]; x[i]], [y[i-1]; y[i]])
            scatter!([x[i-1]], [y[i-1]], color = "blue")
            scatter!([x[i]], [y[i]], color = "red")
            quiver!([x[i-1]; x[i]], [y[i-1]; y[i]],
                quiver=([0.075 * cos(theta[i-1]); 0.075 * cos(theta[i])],
                [0.075 * sin(theta[i-1]); 0.075 * sin(theta[i])])
                )
        else
            circlex1 = R1*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(p1[1]) - sin(theta[i-1])*(p1[2]) )
            circley1 = R1*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(p1[1]) + cos(theta[i-1])*(p1[2]) )

            circlex2 = R2*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(p2[1]) - sin(theta[i-1])*(p2[2]) )
            circley2 = R2*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(p2[1]) + cos(theta[i-1])*(p2[2]) )
            #scatter!([x[i-1]; x[i]], [y[i-1]; y[i]])
            scatter!([x[i-1]], [y[i-1]], color = "red")
            scatter!([x[i]], [y[i]], color = "blue")
            quiver!([x[i-1]; x[i]], [y[i-1]; y[i]],
                quiver=([0.075 * cos(theta[i-1]); 0.075 * cos(theta[i])],
                [0.075 * sin(theta[i-1]); 0.075 * sin(theta[i])])
                )
        end
        plot!(circlex1, circley1, color="dodgerblue")
        plot!(circlex2, circley2, color="maroon")
        display(plot!(legend=false, xlims=(-0.1,1.1), ylims=(-0.1,1.1)))
    end
    # display(plot!(legend=false, xlims=(-0.1,1.1), ylims=(-0.1,1.1)))
end


######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """
"""
function sim_MvNormal(μ, Σ)

    obs_aug = μ + cholesky(Hermitian(Σ), Val(true), check = false).U' * randn(length(μ), 1)

    return obs_aug
end

######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """
"""
function sim_MvNormal_alt(μ,Σ)

    nvars = size(Σ)[1]

    keep = Any[]

    for i in 1:nvars
        if Σ[i,i] != 0.0
            push!(keep, i)
        end
    end

    Σnew = Σ[keep, keep]
    μnew = μ[keep]

    obs = sim_MvNormal(μnew, Σnew)

    obs_aug = zeros(nvars)

    j = 1
    for i in 1:nvars
        if i in keep
            obs_aug[i] = obs[j]
            j = j + 1
        else
            obs_aug[i] = μ[i]
        end
    end

    return obs_aug
end
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """
    SSModel(H, A, F, μ, R, Q, Z)

Description:
A type object containing all parameters necessary to specify a data-generating process in state-space form.
Measurement Equation:
    y_{t} = H β_{t} + A z_{t} + e_{t}
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t}
    e_{t} ~ i.i.d.N(0,R)
    v_{t} ~ i.i.d.N(0,Q)
    z_{t} ~ i.i.d.N(0,Z)
    E(e_t v_s') = 0

Inputs:
- H = measurement equation state vector coefficient matrix.
- A = measurement equation predetermined vector coefficient matrix.
- F = state equation companion matrix.
- μ = state equation intercept vector.
- R = measurement equation error covariance matrix.
- Q = state equation innovation covariance matrix.
- Z = pretermined vector covariance matrix.
"""
@with_kw mutable struct SSModel
    H::Array{Float64,2}   # Measurement equation state vector coefficient matrix
    A::Array{Float64,2}   # Measurement equation predetermined vector coefficient matrix
    F::Array{Float64,2}   # State equation companion matrix
    μ::Array{Float64,1}   # State equation intercept vector
    R::Array{Float64,2}   # Measurement equation error covariance matrix
    Q::Array{Float64,2}   # State equation innovation covariance matrix
    Z::Array{Float64,2}   # Pretermined vector covariance matrix
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    HDFM(nlevels, nvar, nfactors, fassign, flags, varlags, varcoefs, fcoefs, fvars, varvars)

Description:
Creates an object of type `HDFM` that contains all parameters necessary to specify a multi-level linear dynamic factor data-generating process.
This is a convenient alternative to specifying an HDFM directly in state-space form.

Inputs:
- nlevels = number of levels in the multi-level model structure.
- nvar = number of variables.
- nfactors = number of factors for each level (vector of length `nlevels`).
- fassign = determines which factor is assigned to which variable for each level (integer matrix of size `nvar` × `nlevels`).
- flags = number of autoregressive lags for factors of each level (factors of the same level are restricted to having the same number of lags; vector of length `nlevels`).
- varlags = number of observed variable error autoregressive lags (vector of length `nvar`).
- varcoefs = vector of coefficients for each variable in the observation equation (length 1+`nlevels`, where first entry represents the intercept).
- fcoefs = list of `nlevels` number of matrices, for which each row contains vectors of the autoregressive lag coefficients of the corresponding factor.
- fvars = list of `nlevels` number of vectors, where each entry contains the disturbance variance of the corresponding factors.
- varvars = vector of `nvar` number of entries, where each entry contains the innovation variance of the corresponding variable.
"""
@with_kw mutable struct HDFM
    nlevels::Int64                  # number of levels in the multi-level model structure
    nvar::Int64                     # number of variables in the dataset
    nfactors::Array{Int64,1}        # number of factors for each level (vector of length `nlevels`)
    fassign::Array{Int64,2}         # integer matrix of size `nvar` × `nlevels`
    flags::Array{Int64,1}           # number of autoregressive lags for each factor level (vector of length `nlevels`)
    varlags::Array{Int64,1}         # number of obs. variable error autoregressive lags (vector of length `nvar`)
    varcoefs::Array{Float64,2}      # vector of observation equation coefficients for each variables (length 1+`nlevels`)
    varlagcoefs::Array{Float64,2}   # vector of error lag coefficients for each variable (matrix length `nvar`, width max(varlags))
    fcoefs::Array{Any,1}            # list of `nlevels` number of matrices, where each row contains vectors of autoregressive lag coefficients of corresponding factors
    fvars::Array{Any,1}             # list of `nlevels` number of vectors, where each entry contains the disturbance variance of the corresponding factors
    varvars::Array{Float64,1}         # vector of `nvar` number of entries, where each entry contains innovation variances of corresponding variables
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    createSSforHDFM(hdfm)

Description:
Create state-space form coefficient and variance matrices for an HDFM object.

Inputs:
- hdfm = `HDFM` object with all parameters necessary to specify a multi-level linear dynamic factor data-generating process.
"""
function createSSforHDFM(hdfm::HDFM)

    ######################################
    ## Import all HDFM parameters
    @unpack nlevels, nvar, nfactors, fassign, flags, varlags, varcoefs, varlagcoefs, fcoefs, fvars, varvars = hdfm

    ######################################
    ## Specify observation equation coefficient matrix

    # Store number of total lag terms in the state vector
    ntotlags = sum(varlags) + dot(nfactors,flags)           # variable error lags

    # Create empty observation eq. coefficient matrix
    H = zeros(nvar, 1 + ntotlags)               # intercept + total lags

    # Fill out observation eq. coefficient matrix
    for i = 1:nvar
        H[i, 1] = varcoefs[i, 1]
        for j = 1:nlevels
            for k = 1:nfactors[j]
                if k == fassign[i, j]
                    H[i, 1 + sum(nfactors[1:(j-1)]) + k] = varcoefs[i, 1 + j]
                end
            end
        end
    end
    H[:, (1+sum(nfactors)+1):(1+sum(nfactors)+nvar)] = 1.0 .* Matrix(I(nvar))

    ######################################
    ## Specify observation equation error covariance matrix
    R = zeros(nvar, nvar)

    ######################################
    ## Specify state equation companion matrix
    ## and intercept vector

    # Create empty state equation companion matrix and intercept vector
    slength = size(H)[2]                 # length of state vector
    F = zeros(slength, slength)
    μ = zeros(slength)

    # Total number of factors
    ntotfactors = sum(nfactors)

    # Fill out transition eq. companion matrix
    μ[1,1] = 1.0
    for i = 1:nlevels
        for j = 1:nfactors[i]

            rowind = 0
            if i > 1
                rowind = sum(nfactors[1:(i-1)]) + j
            else
                rowind = j
            end

            for k = 1:flags[i]
                F[1+rowind, 1+rowind+(ntotfactors+nvar)*(k-1)] = (fcoefs[i])[j, k] # factor autoregressive lag coefficients
            end
        end
    end
    for i = 1:nvar
        for j = 1:varlags[i]
            F[ntotfactors+i, 1+ntotfactors+(nvar)*(j-1)+i] = varlagcoefs[i, j] # obs. eq. error lag coefficients
        end
    end
    for i = (ntotfactors+nvar+1):(slength)
        for j = 1:(slength-ntotfactors-nvar)
            if i == ntotfactors + nvar + j
                F[i, j] = 1.0
            end
        end
    end

    ######################################
    ## Specify state equation error covariance matrix

    # Create empty state equation error covariance matrix
    Q = zeros(slength, slength)

    # Fill out state equation error covariance matrix
    for i = 1:nlevels
        for j = 1:nfactors[i]

            rowind = 0
            if i > 1
                rowind = sum(nfactors[1:(i-1)]) + j
            else
                rowind = j
            end

            Q[1+rowind, 1+rowind] = (fvars[i])[j]
        end
    end
    for i = 1:nvar
        Q[1+ntotfactors+i, 1+ntotfactors+i] = varvars[i]
    end

    ######################################
    ## Specify all predetermined variable-related parameters
    A = zeros(nvar, nvar)
    Z = zeros(nvar, nvar)

    ######################################
    return H, A, F, μ, R, Q, Z
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """
    convertHDFMtoSS(hdfm)

Description:
Converts an `HDFM` object to an `SSModel` object.

Inputs:
- hdfm = HDFM object containing all HDFM DGP parameters.

Outputs:
- ssmodel = SSModel object containing all state-space model parameters.
"""
function convertHDFMtoSS(hdfm::HDFM)

    H, A, F, μ, R, Q, Z = createSSforHDFM(hdfm::HDFM)

    ssmodel = SSModel(H, A, F, μ, R, Q, Z)

    return ssmodel
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    simulateSSModel(num_obs, ssmodel::SSModel)

Description:
Generate data from a DGP in state space form.
Measurement Equation:
    y_{t} = H β_{t} + A z_{t} + e_{t}
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t}
    e_{t} ~ i.i.d.N(0,R)
    v_{t} ~ i.i.d.N(0,Q)
    z_{t} ~ i.i.d.N(0,Z)
    E(e_t v_s') = 0

Inputs:
- num_obs   = number of observations
- ssmodel   = `SSModel` object containing all parameters necessary to specify a data-generating process in state-space form.
"""
function simulateSSModel(num_obs, ssmodel::SSModel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Convert cov matrices to PSDMat
    # to allow for simulation using MvNormal()
    # (MvNormal() needs a positive definite cov matrix)

    # Create empty data storage matrices
    data_y = zeros(num_obs, size(H)[1])
    data_z = zeros(num_obs, size(Z)[1])
    data_β = zeros(num_obs, size(Q)[1])

    # Initialize β and y
    β0 = inv(I - F) * μ
    y0 = H * β0

    # Initialize z
    if Z == zeros(size(Z)[1], size(Z)[1])
        z0 = zeros(size(Z)[1])
    else
        #z0 = rand(MvNormal(zeros(size(Z)[1]), Z))
        z0 = sim_MvNormal_alt(zeros(size(Z)[1]), Z)
    end

    # Save first observations of y and z
    data_y[1, :] = y0
    data_z[1, :] = z0
    data_β[1, :] = β0

    # Initialize β lag for recursion
    β_lag = β0

    # Recursively generate data
    for t = 2:num_obs
        # Draw transition distrubance
        if Q == zeros(size(Q)[1], size(Q)[1])
            v = zeros(size(Q)[1])
        else
            #v = rand(MvNormal(zeros(size(Q)[1]), Q))
            v = sim_MvNormal_alt(zeros(size(Q)[1]), Q)
        end
        # Record new state observation
        β = μ + F * β_lag + v
        # Draw new z observation
        if Z == zeros(size(Z)[1], size(Z)[1])
            z = zeros(size(Z)[1])
        else
            #z = rand(MvNormal(zeros(size(Z)[1]), Z))
            z = sim_MvNormal_alt(zeros(size(Z)[1]), Z)
        end
        # Draw measurement distrubance
        if R == zeros(size(R)[1], size(R)[1])
            e = zeros(size(R)[1])
        else
            #e = rand(MvNormal(zeros(size(R)[1]), R))
            e = sim_MvNormal_alt(zeros(size(R)[1]), R)
        end
        # Record new measurement observation
        y = H * β + A * z + e
        # Save generated data
        data_y[t, :] = y
        data_z[t, :] = z
        data_β[t, :] = β
        # Update β lag for recursion
        β_lag = β
    end

    # Return data
    return data_y, data_z, data_β
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################


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

using JuliaFEM, Test

# test tet4 stiffness matrix

X = Dict(1 => [2.0, 3.0, 4.0],
         2 => [6.0, 3.0, 2.0],
         3 => [2.0, 5.0, 1.0],
         4 => [4.0, 3.0, 6.0])
element = Element(Tet4, (1, 2, 3, 4))
update!(element, "youngs modulus", 96.0)
update!(element, "poissons ratio", 1/3)
update!(element, "geometry", X)
problem = Problem(Elasticity, "tet4", 3)
add_element!(problem, element)
time = 0.0
assemble!(problem, time)
K_expected = [
     149   108   24   -1    6   12  -54   -48    0  -94   -66  -36
     108   344   54  -24  104   42  -24  -216  -12  -60  -232  -84
      24    54  113    0   30   35    0   -24  -54  -24   -60  -94
      -1   -24    0   29  -18  -12  -18    24    0  -10    18   12
       6   104   30  -18   44   18   12   -72  -12    0   -76  -36
      12    42   35  -12   18   29    0   -24  -18    0   -36  -46
     -54   -24    0  -18   12    0   36     0    0   36    12    0
     -48  -216  -24   24  -72  -24    0   144    0   24   144   48
       0   -12  -54    0  -12  -18    0     0   36    0    24   36
     -94   -60  -24  -10    0    0   36    24    0   68    36   24
     -66  -232  -60   18  -76  -36   12   144   24   36   164   72
     -36   -84  -94   12  -36  -46    0    48   36   24    72  104]
@test isapprox(problem.assembly.K, K_expected)


# Exponential
ϕ = @SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(ExponentialMap(x)),ϕ) ≈ jacobian(ExponentialMap,ϕ)
@test retraction_map(ExponentialMap(ϕ)) == ExponentialMap

ϕ = 1e-6*@SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(ExponentialMap(x)),ϕ) ≈ jacobian(ExponentialMap,ϕ)
@test retraction_map(ExponentialMap(ϕ)) == ExponentialMap

# Vector Part
v = 0.1*@SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(VectorPart(x)),v) ≈
    jacobian(VectorPart, v)
@test retraction_map(VectorPart(v)) == VectorPart

# Gibbs Vectors
g = @SVector rand(3)
@test ForwardDiff.jacobian(x->SVector(CayleyMap(x)),g) ≈ jacobian(CayleyMap, g)
@test retraction_map(CayleyMap(g)) == CayleyMap

# MRPs
p = SVector(rand(MRP))
@test ForwardDiff.jacobian(x->SVector(MRPMap(x)),p) ≈
    jacobian(MRPMap, p)
@test retraction_map(MRPMap(p)) == MRPMap

μ0 = 1/TO.scaling(VectorPart)
jac_eye = [@SMatrix zeros(1,3); μ0*Diagonal(@SVector ones(3))];
@test jacobian(ExponentialMap, p*1e-10) ≈ jac_eye
@test jacobian(MRPMap, p*1e-10) ≈ jac_eye
@test jacobian(CayleyMap, p*1e-10) ≈ jac_eye
@test jacobian(VectorPart, p*1e-10) ≈ jac_eye

############################################################################################
#                                 INVERSE RETRACTION MAPS
############################################################################################

# Exponential Map
Random.seed!(1);
q = rand(UnitQuaternion)
q = UnitQuaternion{ExponentialMap}(q)
qval = SVector(q)
@test ExponentialMap(q) == TO.scaling(ExponentialMap)*logm(q)
@test ExponentialMap(ExponentialMap(q)) ≈ q
@test ExponentialMap(ExponentialMap(ϕ)) ≈ ϕ

function invmap(q)
    μ = TO.scaling(ExponentialMap)
    v = @SVector [q[2], q[3], q[4]]
    s = q[1]
    θ = norm(v)
    M = μ*2atan(θ, s)/θ
    return M*v
end
@test invmap(qval) ≈ TO.scaling(ExponentialMap)*logm(q)

qI = VectorPart(v*1e-5)
@test ForwardDiff.jacobian(invmap, qval) ≈ jacobian(ExponentialMap, q)
@test ForwardDiff.jacobian(invmap, SVector(qI)) ≈ jacobian(ExponentialMap, qI)

# Vector Part
@test VectorPart(q) == TO.scaling(VectorPart)*qval[2:4]
@test VectorPart(VectorPart(q)) ≈ q
@test VectorPart(VectorPart(v)) ≈ v

# Cayley
invmap(q) = 1/q[1] * @SVector [q[2], q[3], q[4]]
@test CayleyMap(q) ≈ TO.scaling(CayleyMap)*invmap(qval)
@test ForwardDiff.jacobian(invmap, qval) ≈ jacobian(CayleyMap, q)
@test CayleyMap(CayleyMap(q)) ≈ q
@test CayleyMap(CayleyMap(g)) ≈ g

# MRP
invmap(q) = TO.scaling(MRPMap)/(1+q[1]) * @SVector [q[2], q[3], q[4]]
@test MRPMap(q) ≈ invmap(qval)
@test ForwardDiff.jacobian(invmap, qval) ≈ jacobian(MRPMap, q)
@test MRPMap(MRPMap(q)) ≈ q
@test MRPMap(MRPMap(p)) ≈ p

# Test near origin
μ0 = TO.scaling(VectorPart)
jacT_eye = [@SMatrix zeros(1,3); μ0*Diagonal(@SVector ones(3))]';
@test isapprox(jacobian(ExponentialMap,qI), jacT_eye, atol=1e-5)
@test isapprox(jacobian(VectorPart,qI), jacT_eye, atol=1e-5)
@test isapprox(jacobian(CayleyMap,qI), jacT_eye, atol=1e-5)
@test isapprox(jacobian(MRPMap,qI), jacT_eye, atol=1e-5)


@testset "Primitives" begin
  @testset "Boxes" begin
    b = Box(P2(0, 0), P2(1, 1))
    @test embeddim(b) == 2
    @test paramdim(b) == 2
    @test coordtype(b) == T
    @test minimum(b) == P2(0, 0)
    @test maximum(b) == P2(1, 1)
    @test extrema(b) == (P2(0, 0), P2(1, 1))

    @test measure(b) == T(1)
    @test P2(1, 1) ∈ b
  end

  @testset "Ball" begin
    b = Ball(P3(1,2,3), T(5))
    @test embeddim(b) == 3
    @test paramdim(b) == 3
    @test coordtype(b) == T
    @test center(b) == P3(1,2,3)
    @test radius(b) == T(5)

    b = Ball(P2(0,0), T(2))
    @test measure(b) ≈ π*(2^2)
    b = Ball(P3(0,0,0), T(2))
    @test measure(b) ≈ (4/3)*π*(2^3)

    b = Ball(P2(0,0), T(2))
    @test P2(1,0) ∈ b
    @test P2(0,1) ∈ b
    @test P2(2,0) ∈ b
    @test P2(0,2) ∈ b
    @test P2(3,5) ∉ b

    b = Ball(P3(0,0,0), T(2))
    @test P3(1,0,0) ∈ b
    @test P3(0,0,1) ∈ b
    @test P3(2,0,0) ∈ b
    @test P3(0,0,2) ∈ b
    @test P3(3,5,2) ∉ b
  end

  @testset "Sphere" begin
    s = Sphere(P3(0,0,0), T(1))
    @test embeddim(s) == 3
    @test paramdim(s) == 2
    @test coordtype(s) == T
    @test center(s) == P3(0, 0, 0)
    @test radius(s) == T(1)

    s = Sphere(P2(0,0), T(2))
    @test measure(s) ≈ 2π*2
    s = Sphere(P3(0,0,0), T(2))
    @test measure(s) ≈ 4π*(2^2)

    s = Sphere(P2(0,0), T(2))
    @test P2(1,0) ∉ s
    @test P2(0,1) ∉ s
    @test P2(2,0) ∈ s
    @test P2(0,2) ∈ s
    @test P2(3,5) ∉ s

    s = Sphere(P3(0,0,0), T(2))
    @test P3(1,0,0) ∉ s
    @test P3(0,0,1) ∉ s
    @test P3(2,0,0) ∈ s
    @test P3(0,0,2) ∈ s
    @test P3(3,5,2) ∉ s
  end

  @testset "Cylinder" begin
    c = Cylinder(P3(1,2,3), P3(4,5,6), T(5))
    @test embeddim(c) == 3
    @test paramdim(c) == 3
    @test coordtype(c) == T
    @test radius(c) == T(5)
    @test height(c) ≈ √27

    @test measure(c) ≈ π*5.0^2*√27
  end
end


module SubModuleRosenBrock

using Test
using F1Method
using LinearAlgebra
using SciMLBase
using ForwardDiff

# Set up:
# - overload `SteadyStateProblem` constructor
# - overload `solve` function
# - define solver algorithm (basic Newton here)
# - define type for that algorithm (here `MyAlg`)

function newton_solve(F, ∇ₓF, x; Ftol=1e-10)
    while norm(F(x)) ≥ Ftol
        x .-= ∇ₓF(x) \ F(x)
    end
    return x
end

# Create a type for the solver's algorithm
struct MyAlg <: SciMLBase.AbstractSteadyStateAlgorithm end

# Overload SciMLBase's solve function
function SciMLBase.solve(prob::SciMLBase.AbstractSteadyStateProblem,
                          alg::MyAlg;
                          Ftol=1e-10)
    # Define the functions according to SciMLBase.SteadyStateProblem type
    p = prob.p
    x0 = copy(prob.u0)
    dx, df = copy(x0), copy(x0)
    F(x) = prob.f(x, p)
    ∇ₓF(x) = prob.f.jac(x, p)
    # Compute `u_steady` and `resid` as per SciMLBase using my algorithm
    x_steady = newton_solve(F, ∇ₓF, x0, Ftol=Ftol)
    resid = F(x_steady)
    # Return the common SciMLBase solution type
    SciMLBase.build_solution(prob, alg, x_steady, resid; retcode=:Success)
end

# This is quasi derived from the Rosenbrock function
# whereby F(x,p) is almost the gradient of the Rosenbrock function
# except the Hessian matroix is made non-symmetric by slightly altering
# F(x,p)

# State function ≈ gradient of Rosenbrock.
# With p = [a, b] the parameters defining the constants
# in Rosenbrocks function, and the slight alteration to make the
# Hessian matrix non-symmetric.
# Specifically, if r(x,y) denotes the rosenbrock function,
#     F₁([x, y], p) = ∂r/∂x
#     F₂([x, y], p) = 0.5 * ∂r/∂y
statefun(x,p) = [
    -2 * (p[1] - x[1]) - 4 * p[2] * (x[2] - x[1]^2) * x[1]
    p[2] * (x[2] - x[1]^2)
]
F = ODEFunction(statefun, jac = (x,p) -> ForwardDiff.jacobian(x -> statefun(x, p), x))

# Define mismatch function f(x,p) and its derivative ∇ₓf(x,p)
# (Note ∇ₓF and ∇ₓf are required by the F1 method)
function state_mismatch(x)
    δ(x) = x .- 1
    return 0.5δ(x)'δ(x)
end
function parameter_mismatch(p)
    δ(p) = log.(p)
    return 0.5δ(p)'δ(p)
end
f(x,p) = state_mismatch(x) + parameter_mismatch(p)
∇ₓf(x,p) = ForwardDiff.jacobian(x -> [f(x,p)], x)

# Initialize the starting state and parameters
x₀ = rand(2)
p₀ = rand(2)
# Initialize the memory cache for storing reusable objects
mem = F1Method.initialize_mem(F, ∇ₓf, x₀, p₀, MyAlg())

# Define the functions via the F1 method
F1_objective(p) = F1Method.objective(f, F, mem, p, MyAlg())
F1_gradient(p) = F1Method.gradient(f, F, ∇ₓf, mem, p, MyAlg())
F1_Hessian(p) = F1Method.hessian(f, F, ∇ₓf, mem, p, MyAlg())

# Define the exact solution and objective for comparison
exact_solution(p) = [p[1], p[1]^2]
exact_objective(p) = f(exact_solution(p), p)

# Run the tests
@testset "objective" begin
    @test exact_objective(p₀) ≈ F1_objective(p₀)
    @test exact_objective(2p₀) ≈ F1_objective(2p₀)
end
@testset "gradient" begin
    @test ForwardDiff.jacobian(p -> [exact_objective(p)], p₀) ≈ F1_gradient(p₀)
    @test ForwardDiff.jacobian(p -> [exact_objective(p)], 2p₀) ≈ F1_gradient(2p₀)
end
@testset "Hessian" begin
    @test ForwardDiff.hessian(p -> exact_objective(p), p₀) ≈ F1_Hessian(p₀)
    @test ForwardDiff.jacobian(p -> [exact_objective(p)], 2p₀) ≈ F1_gradient(2p₀)
end

end # submodule


#--------------------------------------------------------------------------------------------------
"Collection of functions for mapping Number numbers to Hensel code."
#--------------------------------------------------------------------------------------------------
module Hensel
using Mahler
export pIndex, hVector

#--------------------------------------------------------------------------------------------------
# Linear mapping
#--------------------------------------------------------------------------------------------------
"Linear mapping x ∈ [a,b] to y ∈ [0,1]. ."
to01(x::Number, a::Number, b::Number)::Number             = (x-a)/(b-a)
to01(x::Number, (a,b)::Tuple{<:Number, <:Number})::Number = (x-a)/(b-a)

"Linear mapping x ∈ [0,1] to y ∈ [a,b]."
from01(x::Number, a::Number, b::Number)::Number              = x*(b-a)+a
from01(x::Number, (a, b)::Tuple{<:Number, <:Number})::Number = x*(b-a)+a

#--------------------------------------------------------------------------------------------------
# Functions on [a,b]
#--------------------------------------------------------------------------------------------------

"Caclulates Hensel code as integer vector. p is p-adic base, sz - precision.
Asserts: x ∈ [a,b], p > 0, sz > 0."
function hVector(x::Number, a::Number, b::Number, p::Integer, sz::Integer)::Vector{Integer}
    @assert(a <= x <= b)
    return hVector(to01(x, a, b), p, sz)
end

hVector(x::Number, (a, b)::Tuple{<:Real, <:Real}, p::Integer, sz::Integer)::Vector{Integer} =
    hVector(x, a, b, p, sz)

"Calculates integer from real number x∈[a,b]"
iValue(x::Number, ab::Tuple{<:Number, <:Number}, p::Integer, sz::Integer)::Integer =
    iValue(hVector(x, ab, p, sz), p)

"Calculates real number x∈[a,b] from integer"
rValue(n::Integer, p::Integer, sz::Integer, ab::Tuple{<:Number, <:Number})::Number =
    from01(rValue(hVector(n, p, sz), p), ab)

#--------------------------------------------------------------------------------------------------
# Functions on [0,1]
#--------------------------------------------------------------------------------------------------

"Calculates real number x∈[0,1] from Hensel code vector: Σv[i]*p^(-i)/p"
rValue(v::Vector{<:Integer}, p::Integer)::Real =
    foldl((s,i) -> v[i]+s/p, length(v):-1:1, init = p/2)/p

"Calculates real number x∈[0,1] from integer"
rValue(n::Integer, p::Integer, sz::Integer)::Real = rValue(hVector(n, p, sz), p)

"Suppose [0,1] is diveded into p inervals of equal length.
Function retuns index of the interval to which x belongs to.
Indexation starts from 0 (zero). x ∈ [0,1]."
pIndex(x::Number, p::Integer)::Integer = floor(Int, x*p)

"Caclulates Hensel code as integer vector. p is a base, sz - precision.
Asserts: x ∈ [0,1], p > 1, sz > 0."
function hVector(x::Number, p::Integer, sz::Integer)::Vector{Integer}
    @assert(0 <= x <= 1)
    @assert(p > 1)
    @assert(sz > 0)
    if x == 1.0
        return fill(p - 1, sz)
    end
    h = Vector{Integer}(undef, sz)
    for i = 1 : sz
        n = pIndex(x, p)
        x = to01(x, n/p, (n+1)/p)
        h[i] = n
    end
    return h
end

"Calculates Hensel code vector back from its integer representation"
function hVector(n::Integer, p::Integer, sz::Integer)::Vector{Integer}
    h = Vector{Integer}(undef, sz)
    for i = 1 : sz
        h[i] = n % p
        n ÷= p
    end
    return h
end

"Calculates integer from Hensel code vector: Σv[i]*p^i"
iValue(v::Vector{<:Integer}, p::Integer)::Integer =
    foldl((s,i) -> v[i]+s*p, length(v):-1:1, init = 0)

#--------------------------------------------------------------------------------------------------
# Abstruct type
#--------------------------------------------------------------------------------------------------
abstract type AbstractMapping end

#--------------------------------------------------------------------------------------------------
# LinMap
#--------------------------------------------------------------------------------------------------
"Linear mapper integer or hensel code to interval and back"
struct LinMap{T<:Number} <: AbstractMapping
    ab::Tuple{T, T}     # interval
    p::Integer          # p-adic base
    sz::Integer         # p-adic precision
end

"Maps integer to point in [a,b]"
(s::LinMap)(n::Integer) = rValue(n, s.p, s.sz, s.ab)

"Maps hensel code to point in [a,b]"
(s::LinMap)(v::Vector{<:Integer}) = from01(rValue(v, s.p), s.ab)

"Maps x∈[a,b] to integer"
# (s::LinMap)(x::Number)::Integer = iValue(x, s.ab, s.p, s.sz)

"Maps x∈[a,b] to integer, hensel code (integer vector), rational or real number"
function (s::LinMap)(x::Number; return_type::String="integer")
    if return_type == "integer"
        return iValue(x, s.ab, s.p, s.sz)
    elseif return_type == "hensel"
        return hVector(x, s.ab, s.p, s.sz)
    elseif return_type == "rational"
        return Rational(x)
    elseif return_type == "real"
        return x
    end
    throw(DomainError(return_type, "unknown keyword"))
end

#--------------------------------------------------------------------------------------------------
# FunEnv
#--------------------------------------------------------------------------------------------------
"Envelope for calculating functions real -> real as integer -> integer"
struct FunEnv{T1<:AbstractMapping, T2<:AbstractMapping}
    f::Function         # function
    marg::T1            # argument map
    mval::T2            # value map
end
FunEnv(f, arange, ap, asz, vrange, vp, vsz) = FunEnv(f, LinMap(arange, ap, asz), LinMap(vrange, vp, vsz))
FunEnv(f, arange, ap, asz, vrange) = FunEnv(f, arange, ap, asz, vrange, ap, asz)
FunEnv(f, arange, ap, asz) = FunEnv(f, arange, ap, asz, arange)

"Direct call to the function"
(s::FunEnv)(x::Real) = s.f(x)

"Call function real -> real as integer -> [integer | vector | real | rational]"
(s::FunEnv)(x::Integer; return_type::String="integer") =
    s.mval(s.f(s.marg(x)), return_type=return_type)

"Call function real, integer -> real as integer -> [integer | vector | real | rational]"
(s::FunEnv)(x::Integer, i::Integer; return_type::String="integer") =
    s.mval(s.f(s.marg(x), i), return_type=return_type)

(s::FunEnv)((x, i, return_type)::Tuple{Integer, Integer, String}) =
    (s::FunEnv)(x, i, return_type=return_type)

(s::FunEnv)((x, i)::Tuple{Integer, Integer}) = (s::FunEnv)(x, i, return_type="integer")

"For testing how good FunEnv approximates its function"
δ(fe::FunEnv, x::Real) = fe(x) - fe(fe.marg(x), return_type="real")
δ(x::Real, fe::FunEnv) = fe.mval(fe(x)) - fe(fe.marg(x))

#--------------------------------------------------------------------------------------------------
# MahEnv
#--------------------------------------------------------------------------------------------------
"Calculates Mahler expansion of function f of length n"
mexpansion(fe::FunEnv, n::Integer)::Vector = Mahler.mexpansion((x) -> fe(x), n)

"Envelope for Mahler expansion"
struct MahEnv
    fe::FunEnv          # mapping function to integer
    mv::Vector          # vector of Mahler coefficients
end
MahEnv(fe::FunEnv) = MahEnv(fe, mexpansion(fe, fe.marg.sz))
MahEnv(fe::FunEnv, n::Integer) = MahEnv(fe, mexpansion(fe, n))

"Calulation from Mahler expansion"
(s::MahEnv)(n::Integer) = Mahler.mval(n, s.mv)
(s::MahEnv)(n::Integer, m::Integer) = Mahler.mval(n, s.mv, m)
"Optional calculation - can use Mahler, extension of FunEnv call"
function (s::MahEnv)(n::Integer, option::String)
    if option == "Mahler"
        return Mahler.mval(n, s.mv)
    elseif option == "Hensel"
        return (s.fe)(n)
    end
    throw(DomainError(option, "unknown keyword"))
end

(s::MahEnv)(x::Real) = Mahler.mval(s.fe.marg(x), s.mv)
(s::MahEnv)(x::Real, m::Integer) = Mahler.mval(s.fe.marg(x), s.mv, m)

δ(se::MahEnv, n::Integer) = (se)(n) - (se)(n, "Hensel")
δ(se::MahEnv, n::Integer, k::Integer) = (se)(n, k) - (se)(n, "Hensel")

#--------------------------------------------------------------------------------------------------
# Additional functions
#--------------------------------------------------------------------------------------------------

"Logistic map function. x∈[0,1]. For testing.
See Logistic map"
logistic(x::Real, r::Real)::Real = r*x*(1.0 - x)
"n-th iteration of logictic function"
logistic(x::Real, r::Real, n::Integer)::Real = foldl((s,i) -> logistic(s,r), 0:(n-1), init=x)

"Dyadic transformation. x∈[0,1] For testing.
See Dyadic transformation"
dyadic(x::Real)::Real = mod(2*x, 1)
"n-th iteration of dyadic function"
dyadic(x::Real, n::Integer)::Real = foldl((s,i) -> dyadic(s), 0:(n-1), init=x)

"nyadic transformation. x∈[0,1] For testing."
nyadic(x::Real, n::Integer)::Real = mod(n*x, 1)
"k-th iteration of dyadic function"
nyadic(x::Real, n::Integer, k::Integer)::Real = foldl((s,i) -> nyadic(s, n), 0:(k-1), init=x)

"Shift operator for continued fractions. x∈[0,1]
See THE GAUSS-KUZMIN-WIRSING OPERATOR"
cfshift(x::Real)::Real = 1/x - floor(1/x)
"n-th iteration of shift operator for continued fractions"
cfshift(x::Real, n::Integer)::Real = foldl((s,i) -> cfshift(s), 0:(n-1), init=x)

end # module

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

# include("./src/Hensel.jl")
# or
# pkg> activate .
# and then
# using Hensel

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


using ConservationLawsParticles
using RecursiveArrayTools, DifferentialEquations, Plots

# model
V(t, x) = 0.0 #2 + 0.5sin(2x) + 0.3sin(2sqrt(2)*x) + 0.2cos(2sqrt(7)*x)
W(t, x) = 0.0 #x^2/2 - abs(x)
W′(t, x) = 0.0 #x - sign(x)
mob(ρ) = 0.0 # max((ρ-1)^2 * (2ρ+1), 0)
smodel = DiffusiveSampledModel((V,), ((W′,),), (mob,), (Diffusion(1),))
# imodel = DiffusiveIntegratedModel((V,), ((W,),), (mob,))

# initial condition
x0 = ArrayPartition(gaussian_particles(3, 201))

# time span
tspan = (0., 15.)

# ODE system for the particles
dprob = ODEProblem(velocities_diff!, x0, tspan, smodel)

abstol = reltol = 2e-8

@time dsol = solve(dprob, BS5(); abstol=abstol, reltol=reltol);

# plot the particle trajectories
plot(legend=false, title="Trajectories", xlabel="time", ylabel="position")
plot!(dsol, vars=1:5:201; color=:green)
savefig("diffusion1.png")

# plot an animation of the density
anim = @animate for t in range(0.0, 7.0; step=1/24)
    plot(title="Density", xlabel="position", ylabel="density",
        xlims=(-10,10), ylims=(0,1/sqrt(pi)))
    densityplot!(dsol(t); color=:green, label="particles")
    plot!(x -> exp(-x^2/(1+4t))/sqrt(pi*(1+4t)); color=:orange, label="analytic")
end
gif(anim, "diffusion1.gif")


struct MeshPointNM{U,D,C,N,T}
    patch::Simplex{U,D,C,N,T}
    bary::SVector{D,T}
    cart::SVector{U,T}
end

Base.length(m::MeshPointNM) = length(m.cart)
Base.getindex(p::MeshPointNM, i::Int) = p.cart[i]

"""
    cartesian(neighborhood) -> point

Return the cartesian coordinates of the point where `neighborhood` is located.
"""
function cartesian(mp::MeshPointNM)
    mp.cart
end

"""
    parametric(neighborhood) -> point

Return the parameters where `neighborhood` is located.
"""
function parametric(mp::MeshPointNM)
    mp.bary
end

chart(nbd::MeshPointNM) = mp.patch

"Return the barycentric coordinates of `mp`"
barycentric(mp::MeshPointNM) = SVector(mp.bary[1], mp.bary[2], 1-mp.bary[1]-mp.bary[2])

"""
A number defines a neighborhood in euclidian space
"""
jacobian(x::Number) = one(x)

jacobian(mp::MeshPointNM) = volume(mp.patch) * factorial(dimension(mp.patch))

"""
    tangents(neighborhod, i) -> tangent_i

Return the i-th tangent vector at the neighborhood.
"""
function tangents(mp::MeshPointNM, i)
    mp.patch.tangents[i]
end

"""
    normal(neighborhood)

Return the normal at a neighborhood on a surface.
"""
function normal(mp::MeshPointNM)
    mp.patch.normals[1]
end

"""
    neighborhood(chart, params)

Create a neighborhood from a chart and a set of parameter values.
"""
function neighborhood(p::Simplex, bary)
  D = dimension(p)
  T = coordtype(p)
  P = SVector{D,T}
  cart = barytocart(p, bary)
  MeshPointNM(p, P(bary), cart)
end

"""
    center(simplex) -> neighborhood

Create the neighborhood at the center of the simplex, i.e. the point corresponding
to parameter `(1/(D+1), 1/(D+1), ...)` where `D` is the simplex dimension.
"""
@generated function center(p::Simplex{U,D,C,N,T}) where {U,D,C,N,T}
    uv = ones(T,D)/(D+1)
    :(neighborhood(p, $uv))
end


C(x)=((x^2+1.)^0.5+x)^2
Shmin(Sv,Pp,μ)=((Sv-Pp)/C(μ))+Pp
Shmax(Sv,Pp,μ)=C(μ)*(Sv-Pp)+Pp

@userplot zobackogram

@recipe function f(h::zobackogram)

    Sv, Pp, μ = h.args

    Smin=Shmin(Sv,Pp,μ)
    Smax=Shmax(Sv,Pp,μ)

    #45° line. slope=1
    L1=[0,Smax*1.1]

    #Shmin line
    L2=[Smin Smin;Smin Sv]

    #Shmax Horizonal line
    L3=[Sv Smax;Smax Smax]

    #Shmax diagonal Line
    L4=[Smin Sv;Sv Smax]

    #Internal Lines
    L5=[Smin Sv; Sv Sv;Sv Smax]

    xlabel := "Shmin [psi/ft]"
    ylabel := "Shmax [psi/ft]"
    legend := :bottomright
    legendfontsize --> 6

    #45° line. slope=1
    @series begin
        seriestype := :path
        linewidth := 3
        linecolor := :black
        label := "45° Line"
        L1,L1
    end

    @series begin
        seriestype := :path
        linewidth := 2
        linecolor := :darkblue
        label := "Shmin Line"
        L2[:,1],L2[:,2]
    end

    @series begin
        seriestype := :path
        linewidth := 2
        linecolor := :darkred
        label := "Shmax Horizontal Line"
        L3[:,1],L3[:,2]
    end

    @series begin
        seriestype := :path
        linewidth := 2
        linecolor := :darkgreen
        label := "Shmax Strike Slip"
        L4[:,1],L4[:,2]
    end

    @series begin
        seriestype := :path
        linewidth := 1
        linestyle := :dash
        linecolor := :black
        label := "Limits"
        series_annotations := map(x->Plots.text(x ,:left,:top,6),["Shmin=$(round(Smin,digits=1))",
                                "Sv=$(round(Sv,digits=1))","Shmax=$(round(Smax,digits=1))"])
        L5[:,1],L5[:,2]
    end

end


# Removing t = 0, such that Σ is invertible
t = Vector(0.1:0.1:100); p = 2;

# Creating generators U,V that result in a positive-definite matrix Σ
Ut, Vt = spline_kernel(t', p)

K = SymEGRQSMatrix(Ut,Vt,ones(size(Ut,2)))
x = randn(size(K,1))
Kfull = Matrix(K)

# Testing multiplication
@test K*x ≈ Kfull*x
@test K'*x ≈ Kfull'*x

# Testing linear solve
@test K\x ≈ Kfull\x

# Testing (log)determinant
@test logdet(K) ≈ logdet(Kfull)
@test det(K) ≈ det(Kfull)

# Testing show
@test Matrix(K) ≈ tril(K.Ut'*K.Vt) + triu(K.Vt'*K.Ut,1) + Diagonal(K.d)
@test Kfull[3,1] ≈ K[3,1]
@test Kfull[2,2] ≈ K[2,2]
@test Kfull[1,3] ≈ K[1,3]


function compute_int_I(params, α₁, α₂, n)
    (; r, D, dim) = params
    m = length(D)
    R = r[m]

    dr = r[m] * 1e-12
    r₂ = r[1:m] .- dr
    r₁ = [0; r[1:m-1]] .+ dr

    # Check if α₁ is not equal to α₂
    if abs(α₁ - α₂) > 1e-8
        λ₁ = α₁^2
        λ₂ = α₂^2

        # Computation at point r₁
        V₁, dV₁ = compute_v(params, α₁, n, r₁)
        V₂, dV₂ = compute_v(params, α₂, n, r₁)
        I₁ = D .* r₁ .^ (dim - 1) .* (dV₁ .* V₂ - dV₂ .* V₁)

        # Computation at point r₂
        V₁, dV₁ = compute_v(params, α₁, n, r₂)
        V₂, dV₂ = compute_v(params, α₂, n, r₂)
        I₂ = D .* r₂ .^ (dim - 1) .* (dV₁ .* V₂ - dV₂ .* V₁)

        I = (I₂ - I₁) / (λ₁ - λ₂) / R^dim
    else
        α = α₁
        λ = α^2
        nu = compute_ν(n, dim)

        if α > 1e-6
            V, dV = compute_v(params, α, n, r₁)
            I₁ = (
                D .* r₁ .^ dim .* dV .^ 2 +
                (λ * r₁ .^ dim - D .* r₁ .^ (dim - 2) * nu) .* V .^ 2 +
                (dim - 2) * D .* r₁ .^ (dim - 1) .* dV .* V
            )

            V, dV = compute_v(params, α, n, r₂)
            I₂ = (
                D .* r₂ .^ dim .* dV .^ 2 +
                (λ * r₂ .^ dim - D .* r₂ .^ (dim - 2) * nu) .* V .^ 2 +
                (dim - 2) * D .* r₂ .^ (dim - 1) .* dV .* V
            )

            I = (I₂ - I₁) / (2 * λ * R^dim)
        else
            I = (r₂ .^ dim - r₁ .^ dim) / R^dim / dim
        end
    end

    I
end


"""
Functions for the Temperature Size Rule, following Forster et al 2012 & Sentis et al 2017

The function takes 3 arguments :

- dry_mass_293 : the dry mass of the organism at 293.15 K (20 celsius degrees)
- temperature_K : the temperature in Kelvin
- TSR_type : temperature size response slopes

4 different temperature size response slopes are included, in addition to the no-response scenario.
They are called by the following keywords :

1) mean_aquatic : Mean aquatic TS response, the average TS response for aquatic Metazoa
2) mean_terestrial : Mean terrestrial TS response, the average TS response for terrestrial Metazoa
3) maximum : Maximum TS response, the strongest negative TS response
4) reverse : Reverse TS response, the strongest positive TS response

The function is called in the following way :

DM293 = 0.2 # Dry mass at 20 degrees
temperature = 273.15+30 # temperature in K
TSR_t = :mean_aquatic # mean aquatic TS slope for instance

temperature_size_rule(DM20, temperature, TSR_t)

It returns the body mass of the organisms at the given temperature according to the given TS response slope.
"""
function temperature_size_rule(parameters)

    temperature_C = parameters[:T] - 273.15 #convert from kelvins to celsius
    conversion_factor_mass = 6.5

    if length(parameters[:bodymass]) > 1
        wmass = parameters[:bodymass]
        dmass = wmass ./ conversion_factor_mass
        parameters[:dry_mass_293] = dmass

    elseif length(parameters[:dry_mass_293]) > 1
        wmass = conversion_factor_mass .* parameters[:dry_mass_293]
        dmass = parameters[:dry_mass_293]
        if parameters[:TSR_type] == :no_response
            parameters[:bodymass] = wmass
        end

    else
        M = parameters[:Z].^(parameters[:trophic_rank].-1)
        if parameters[:TSR_type] == :no_response
            parameters[:bodymass] = M
        end
        wmass = M
        dmass = wmass ./ conversion_factor_mass
        parameters[:dry_mass_293] = dmass

    end

    if parameters[:TSR_type] == :mean_aquatic
        PCM = -3.90 .- 0.53 .* log10.(parameters[:dry_mass_293])     # PCM = Percentage change in body-mass per degrés C

    elseif parameters[:TSR_type] == :mean_terrestrial
        PCM = -1.72 .+ 0.54 .* log10.(parameters[:dry_mass_293])

    elseif parameters[:TSR_type] == :maximum
        PCM = -8

    elseif parameters[:TSR_type] == :reverse
        PCM = 4

    elseif parameters[:TSR_type] == :no_response
        PCM = 0

    end
        TS_response = log.(PCM/100 .+ 1) # Sign and magnitude of TS response
        parameters[:bodymass] = wmass .* exp.(TS_response .* (temperature_C-20))

end


#### randomized SVD (from Jiahao Chen, based on http://arxiv.org/pdf/0909.4061.pdf)

import Base.LinAlg.SVD

#The simplest possible randomized svd
#Inputs
#   A: input matrix
#   n: Number of singular value/vector pairs to find
#   p: Number of extra vectors to include in computation
function rsvd(A, n, p=0)
    Q = rrange(A, n, p=p)
    rsvd_direct(A, Q)
end

#Algorithm 4.4: randomized subspace iteration
#A must support size(A), multiply and transpose multiply
#p is the oversampling parameter
#q controls the accuracy of the subspace found; it is the "number of power iterations"
#A good heuristic is that when the original scheme produces a basis whose
#approximation error is within a factor C of the optimum, the power scheme produces
#an approximation error within C^(1/(2q+1)) of the optimum.
function rrange(A, l::Integer; p::Integer=5, q::Integer=3)
    p≥0 || error()
    m, n = size(A)
    l <= m || error("Cannot find $l linearly independent vectors of $m x $n matrix")
    Ω = randn(n, l+p)
    Q = q_from_qr(A*Ω)
    for t=1:q
        Q = q_from_qr(A'*Q)
        Q = q_from_qr(A*Q)
    end
    Q = p==0 ? Q : Q[:,1:l]
end

function q_from_qr(Y, l::Integer=-1)
    Q = full(qrfact!(Y)[:Q])
    Q = l<0 ? Q : Q[:,1:l]
end

#Algorithm 5.1: direct SVD
#More accurate
function rsvd_direct(A, Q)
    B=Q'A
    S=svdfact!(B)
    SVD(Q*S[:U], S[:S], S[:Vt])
end


function update_weights!(
    new_weights::Vector,
    weights::Vector,
    training_proportion::Real,
)
    n = round(Int, training_proportion * sum(weights))
    π = weights ./ sum(weights)
    d = Distributions.Multinomial(n, π)
    rand!(d, new_weights)
    return
end


import DielectricUpdateTechique as DUTCH
using Test

nCells = [30,33,61];
res = [1f-3,1f-3,1f-3];

r = DUTCH._create_grid_Greens_function(nCells,res);

@test size(r) == (64,128,128)
@test r[33,65,65] == 0
@test r[34,66,66] == sqrt(res[1]^2 + res[2]^2 + res[3]^2)
@test r[33,66,66] == sqrt(res[2]^2 + res[3]^2)
@test r[1,66,66] == sqrt((32*res[1])^2 + res[2]^2 + res[3]^2)

G = DUTCH._create_Greens_function(r,6.4,res[1]/2f0);

@test size(G) == size(r)
@test maximum(abs.(G)) == abs(G[1,1,1])

Gdv = DUTCH.create_Greens_function(nCells,res,6.4);

@test Gdv[1,1,1] == G[1,1,1]*prod(res)


# ======================================================================
# Description
# ========================================================================
# Calculates MOID between two keplerian orbits in AU
#
# Created: 6/16/20
# Original Author : Anivid Pedros Faura, anivid.pedrosfaura@colorado.edu
# Translated from .m to .jl by: Luke Bury, luke.bury@colorado.edu

# ========================================================================
# Libraries, Functions, and Necessary Paths
# ========================================================================
# # Libraries
# using DifferentialEquations # For using ODEProblem
# using LinearAlgebra # Allows for norm(), cross(), dot(), etc
# using Plots
using Printf # allows use of println(@sprintf())

# ========================================================================
# Dependent functions
# ========================================================================
function RefFrame(A,B)
    ### Reference frame rotation
    # 3-1-3 Tranformation (orbital elements)
    theta1 = A["Omega"]
    theta2 = A["i"]
    theta3 = A["argp"]
    RA = zeros(3,3) #from reference to A frame
    RA[1,1] = cos(theta3)*cos(theta1)-sin(theta3)*sin(theta1)*cos(theta2)
    RA[1,2] = sin(theta1)*cos(theta3)+cos(theta1)*cos(theta2)*sin(theta3)
    RA[1,3] = sin(theta2)*sin(theta3)
    RA[2,1] = -cos(theta1)*sin(theta3) - sin(theta1)*cos(theta2)*cos(theta3)
    RA[2,2] = -sin(theta3)*sin(theta1) + cos(theta3)*cos(theta2)*cos(theta1)
    RA[2,3] = sin(theta2)*cos(theta3)
    RA[3,1] = sin(theta2)*sin(theta1)
    RA[3,2] = -sin(theta2)*cos(theta1)
    RA[3,3] = cos(theta2)

    ### Preparing the orbit
    x = [cos(B["Omega"])*cos(B["argp"]) - sin(B["argp"])*cos(B["i"])*sin(B["Omega"])
         sin(B["Omega"])*cos(B["argp"]) + sin(B["argp"])*cos(B["i"])*cos(B["Omega"])
         sin(B["i"])*sin(B["argp"])]

    y = [-cos(B["Omega"])*sin(B["argp"]) - cos(B["argp"])*cos(B["i"])*sin(B["Omega"])
         -sin(B["Omega"])*sin(B["argp"]) + cos(B["argp"])*cos(B["i"])*cos(B["Omega"])
         sin(B["i"])*cos(B["argp"])];

    z = [sin(B["i"])*sin(B["Omega"])
        -sin(B["i"])*cos(B["Omega"])
        cos(B["i"])]

    xn = RA * x
    yn = RA * y
    zn = RA * z

    incliB = atan(sqrt(zn[1]^2 + zn[2]^2),zn[3])
    omegaB = -atan(zn[1],-zn[2])
    argpB = -atan(xn[3],yn[3])

    return incliB, omegaB, argpB
end

function WaterProcedure(incliB, omegaB, argpB, N, vtrueB, vL, vdis)
    if N<2
        vtrueB = zeros(4,1)
        vL = zeros(4,1)
        vdis = zeros(4,1)
        N = 4
        for j = 1:4
            vtrueB[j] = (0.25+0.5*j)*pi #evenly distributed 0.75 - 1.25 - 1.75 - 2.25
            a = (cos(omegaB)*cos(argpB+vtrueB[j]) - sin(omegaB)*sin(argpB+vtrueB[j])*cos(incliB))
            b = (sin(omegaB)*cos(argpB+vtrueB[j]) + cos(omegaB)*sin(argpB+vtrueB[j])*cos(incliB))
            vL[j] = atan(b,a)
            vdis[j] = 1e6; #set to something large
        end
    end

    return vtrueB, vL, vdis, N
end

# ========================================================================
# Calculating MOID
# ========================================================================
# -------------------------------------------------
# Setting up initial conditions
# -------------------------------------------------
# --------------------------
# Orbital parameters of body A
# --------------------------
A = Dict("sma" => 2.7691652, "e" => 0.0760091, "argp" => 73.59764*π/180, "Omega" => 80.30553*π/180, "i" => 10.59407*π/180)

# --------------------------
# Orbital parameters of body B
# --------------------------
B = Dict("sma" => 2.3655722, "e" => 0.127581, "argp" => 87.42605*π/180, "Omega" => 307.46872*π/180, "i" => 2.09575*π/180)

# -------------------------------------------------
# Calculating MOID
# -------------------------------------------------
# --------------------------
# New orbital elements: A frame
# --------------------------
# incliA = 0, omegaA = 0, argpA = 0
incliB, omegaB, argpB = RefFrame(A,B)

# --------------------------
# Scanning orbits
# --------------------------
# Scanning one full revolution of meridional plane to find local minima
# First guess for MOID is set to be big ~1e6AU

# Angle sweeping steps
cstep = 0.12 #rad - based on Wisnioski paper ~ 0.12rad
stepini = 0.07 #rad - for initial tuning step
stepfin = 1e-5 #rad - for final step of first tuning
stepmin = 1e-14 #rad - threshold step for secondtuning

# Initial Guess
trueB = -2 * cstep
moid = 1e6
dist_o = 1e6
vdis = ones(4,1)*1e6

for i in 1:2
    # First triplet
    rB = B["sma"] * (1-B["e"].^2) / (1+B["e"]*cos(trueB)) #compute the radius for B
    xB = rB * (cos(omegaB)*cos(argpB+trueB) - sin(omegaB)*sin(argpB+trueB)*cos(incliB))
    yB = rB * (sin(omegaB)*cos(argpB+trueB) + cos(omegaB)*sin(argpB+trueB)*cos(incliB))
    zB = rB * sin(argpB+trueB)*sin(incliB)

    rhoB = sqrt(xB^2+yB^2)
    L = atan(yB,xB)

    rA = A["sma"] * (1-A["e"].^2) / (1+A["e"]*cos(L)) #compute the radius for A
    rA2 = A["sma"] * (1-A["e"].^2) / (1-A["e"]*cos(L))

    if abs(rhoB-rA)>abs(rhoB+rA)
        rA = rA2
        L = L - pi
        diff = rhoB + rA2
    else
        diff = rhoB - rA
    end

    global D = zB^2+diff^2 # square of the distance
    # storing
    if i == 1
        global dist_oo = D
    else
        global dist_o = D
        global trueB_o = trueB
        global L_o = L
    end

    global trueB = trueB + cstep
end # for i in 1:2

# Full revolution
N = 0 # number of minima
dmin = D

vtrueB = fill(1,0).*0.0
vL = fill(1,0).*0.0
vdis = fill(1,0).*0.0

while trueB < (2*pi + cstep)
    rB = B["sma"] * (1-B["e"].^2) / (1+B["e"]*cos(trueB)) #compute the radius for B
    xB = rB * (cos(omegaB)*cos(argpB+trueB) - sin(omegaB)*sin(argpB+trueB)*cos(incliB))
    yB = rB * (sin(omegaB)*cos(argpB+trueB) + cos(omegaB)*sin(argpB+trueB)*cos(incliB))
    zB = rB * sin(argpB+trueB)*sin(incliB)

    rhoB = sqrt(xB^2+yB^2)
    L = atan(yB,xB)

    rA = A["sma"] * (1-A["e"].^2) / (1+A["e"]*cos(L)) #compute the radius for A
    rA2 = A["sma"] * (1-A["e"].^2) / (1-A["e"]*cos(L))

    if abs(rhoB-rA)>abs(rhoB+rA)
        rA = rA2
        L = L - pi
        diff = rhoB + rA2
    else
        diff = rhoB - rA
    end

    D = zB^2+diff^2 # square of the distance

    if (dist_o<=D) && (dist_o<=dist_oo)
        # global N = N + 1
        # vtrueB = trueB_o
        # vL = L_o
        # vdis = dist_o

        global vtrueB = append!(vtrueB,trueB_o)
        global vL = append!(vL, L_o)
        global vdis = append!(vdis, dist_o)
    end
    if dmin>D
        global dmin = D
    end

    global dist_oo = dist_o
    global trueB_o = trueB
    global L_o = L
    global dist_o = D
    global trueB = trueB + cstep
end

# -------------------------------------------------
# Water Procedure
# -------------------------------------------------
vtrueB, vL, vdis, N = WaterProcedure(incliB, omegaB, argpB, N, vtrueB, vL, vdis)

# -------------------------------------------------
# PARALLEL TUNING
# -------------------------------------------------
# Move objects separately along their orbits
# Smallest possible distance no longer meridional distance
rBt = fill(1,3).*NaN
rAt = fill(1,3).*NaN
xBt = fill(1,3).*NaN
yBt = fill(1,3).*NaN
zBt = fill(1,3).*NaN
xAt = fill(1,3).*NaN
yAt = fill(1,3).*NaN
k = 1
while k < N+2
    if k<=N
        global moid = vdis[k]
        global trueB_m = vtrueB[k]
        global L_m = vL[k]
        global myStep = stepini
        global threshold = stepfin #maybe here problem
    else
        if N == 2
            # in case of two minima are very close to each other(<1E-4 a.u.)
            # go to "water procedure"
            if (abs(vdis[1]-vdis[2])<1e-4)
                global N = 1
                global vtrueB, vL, vdis, N = WaterProcedure(incliB, omegaB, argpB, N)
                global k = 1
            else
                if (vdis[1]<moid)
                    global moid = vdis[1]
                    global trueB_m = vtrueB[1]
                    global L_m = vL[1]
                end
            end
        else
            # final tuning
            for i in 1:(N-1)
                if vdis[i]<moid
                    global moid = vdis[i]
                    global trueB_m = vtrueB[i]
                    global L_m = vL[i]
                end
            end
        end
        global myStep  = 2*stepini #inital state
        global threshold = stepmin #terminal state
    end

    rBt[2] = B["sma"] * (1-B["e"]^2) / (1+B["e"]*cos(trueB_m)) #compute the radius for B
    xBt[2] = rBt[2] * (cos(omegaB)*cos(argpB+trueB_m) - sin(omegaB)*sin(argpB+trueB_m)*cos(incliB))
    yBt[2] = rBt[2] * (sin(omegaB)*cos(argpB+trueB_m) + cos(omegaB)*sin(argpB+trueB_m)*cos(incliB))
    zBt[2] = rBt[2] * sin(argpB+trueB_m)*sin(incliB)

    rAt[2] = A["sma"] * (1-A["e"].^2) / (1+A["e"]*cos(L_m)) #compute the radius for A
    xAt[2] = rAt[2]*cos(L_m)
    yAt[2] = rAt[2]*sin(L_m)

    aleft = true
    aright = true
    bleft = true
    bright = true

    while (myStep>=threshold)
        lpoints = 0
        k1min   = 1
        k1max   = 3
        i1min   = 1
        i1max   = 3
        calc1   = false
        calc2   = false
        calc3   = false
        calc4   = false

        if bleft
            rBt[1] = B["sma"] * (1-B["e"]^2) / (1+B["e"]*cos(trueB_m-myStep)) #compute the radius for B
            xBt[1] = rBt[1] * (cos(omegaB)*cos(argpB+trueB_m-myStep) - sin(omegaB)*sin(argpB+trueB_m-myStep)*cos(incliB))
            yBt[1] = rBt[1] * (sin(omegaB)*cos(argpB+trueB_m-myStep) + cos(omegaB)*sin(argpB+trueB_m-myStep)*cos(incliB))
            zBt[1] = rBt[1] * sin(argpB+trueB_m-myStep)*sin(incliB)
            lpoints = lpoints + 1;
        end
        if bright
            rBt[3] = B["sma"] * (1-B["e"]^2) / (1+B["e"]*cos(trueB_m+myStep)) #compute the radius for B
            xBt[3] = rBt[3] * (cos(omegaB)*cos(argpB+trueB_m+myStep) - sin(omegaB)*sin(argpB+trueB_m+myStep)*cos(incliB))
            yBt[3] = rBt[3] * (sin(omegaB)*cos(argpB+trueB_m+myStep) + cos(omegaB)*sin(argpB+trueB_m+myStep)*cos(incliB))
            zBt[3] = rBt[3] * sin(argpB+trueB_m+myStep)*sin(incliB)
            lpoints = lpoints + 1;
        end
        if aleft
            rAt[1] = A["sma"] * (1-A["e"]^2) / (1+A["e"]*cos(L_m-myStep)) #compute the radius for A
            xAt[1] = rAt[1]*cos(L_m-myStep)
            yAt[1] = rAt[1]*sin(L_m-myStep)
            lpoints = lpoints + 1
        end
        if aright
            rAt[3] = A["sma"] * (1-A["e"]^2) / (1+A["e"]*cos(L_m+myStep)) #compute the radius for A
            xAt[3] = rAt[3]*cos(L_m+myStep)
            yAt[3] = rAt[3]*sin(L_m+myStep)
            lpoints = lpoints + 1
        end

        k1_t = 2
        i1_t = 2

        if lpoints == 1
            if aleft
                i1max = 1
            end
            if aright
                i1min = 3
            end
            if bright
                k1min = 3
            end
            if bleft
                k1max = 1
            end
        end

        if lpoints == 2
            if (aleft && bright)
                calc1 = true
            end
            if (aleft && bleft)
                calc2 = true
            end
            if (aright && bright)
                calc3 = true
            end
            if (aright && bleft)
                calc4 = true
            end
        end

        for k1 in k1min:k1max
            for i1 in i1min:i1max
                compute = true
                if lpoints == 2
                    if i1 != 1
                        if (k1 != 3 && calc1)||(k1 != 1 && calc2)
                            compute = false
                        end
                    end
                    if i1 != 3
                        if (k1 != 3 && calc3)||(k1 != 1 && calc4)
                            compute = false
                        end
                    end
                end
                if (k1 == 2 && i1 == 2)
                    compute = false
                end
                if (compute)
                    Dx = xBt[k1]-xAt[i1]
                    Dy = yBt[k1]-yAt[i1]
                    Dz = zBt[k1]
                    dist = Dx*Dx + Dy*Dy + Dz*Dz
                    if (dist<moid)
                        moid = dist
                        k1_t = k1
                        i1_t = i1
                    end
                end
            end
        end

        if (k1_t!=2) || (i1_t!=2)
            aleft=false
            aright=false
            bleft=false
            bright=false

            if i1_t!=2
                if i1_t == 1
                    aleft = true
                    L_m = L_m - myStep
                    rAt[3] = rAt[2]
                    xAt[3] = xAt[2]
                    yAt[3] = yAt[2]
                    rAt[2] = rAt[1]
                    xAt[2] = xAt[1]
                    yAt[2] = yAt[1]
                else
                    aright = true
                    L_m = L_m + myStep
                    rAt[1] = rAt[2]
                    xAt[1] = xAt[2]
                    yAt[1] = yAt[2]
                    rAt[2] = rAt[3]
                    xAt[2] = xAt[3]
                    yAt[2] = yAt[3]
                end
            end
            if k1_t!=2
                if k1_t == 1
                   bleft = true
                    trueB_m = trueB_m - myStep
                    rBt[3] = rBt[2]
                    xBt[3] = xBt[2]
                    yBt[3] = yBt[2]
                    zBt[3] = zBt[2]
                    rBt[2] = rBt[1]
                    xBt[2] = xBt[1]
                    yBt[2] = yBt[1]
                    zBt[2] = zBt[1]
                else
                    bright = true
                    trueB_m = trueB_m + myStep
                    rBt[1] = rBt[2]
                    xBt[1] = xBt[2]
                    yBt[1] = yBt[2]
                    zBt[1] = zBt[2]
                    rBt[2] = rBt[3]
                    xBt[2] = xBt[3]
                    yBt[2] = yBt[3]
                    zBt[2] = zBt[3]
                end
            end
        else
            aleft = true
            aright = true
            bleft = true
            bright = true
            myStep = myStep*0.15
        end
    end

    if k <= N
        vdis[k] = moid
        vtrueB[k] = trueB_m
        vL[k] = L_m
    end
    global k = k+1
end

moid = sqrt(moid)

println(@sprintf("\nThe MOID is: %1.16f\n", moid))

# ========================================================================
# Useful Formatting Structures
# ========================================================================
# -------------------------------------------------
#
# -------------------------------------------------

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


using CSV, DataFrames, Statistics

const minimum_similarity = 0.8
const movies = CSV.read("top_10_movies_user_rankings.tsv", delim = '\t')

function user_similarity(target_user)
  similarity = Dict{Symbol,Float64}()

  for user in names(movies[:, 2:end])
    user == target_user && continue

    ratings = movies[:, [user, target_user]]
    common_movies = ratings[(ratings[user] .> 0) .& (ratings[target_user] .> 0), :]
    # common_movies = ratings[(ratings[user] .> 7) .& (ratings[target_user] .> 0), :]

    correlation = try
      cor(common_movies[user], common_movies[target_user])
    catch
      0.0
    end

    similarity[user] = correlation
  end

  similarity
end

function recommendations(target_user)
  recommended = Dict{String,Float64}()

  for (user,similarity) in user_similarity(target_user)
    similarity > minimum_similarity || continue

    ratings = movies[:, [Symbol("Movie title"), user, target_user]]
    recommended_movies = ratings[(ratings[user] .>= 7) .& (ratings[target_user] .== 0), :]

    for movie in eachrow(recommended_movies)
      recommended[movie[Symbol("Movie title")]] = movie[user] * similarity
    end
  end

  recommended
end

for user in names(movies)[2:end]
  println("Recommendations for $user: $(recommendations(user))")
end


#feynmantheory.jl

"""
Implementation of Feynman's original variational technique applied to the Polaron model.

See Feynman 1955:
http://dx.doi.org/10.1103/PhysRev.97.660
"""

using ..PolaronMakie

# Equation 31: The <|X(t) - X(s)|^{-1}> * exp(-|t-w|) effective action.
A_integrand(v, w, τ) = (w^2 * τ + (v^2 - w^2) / v * (1 - exp(-v * τ)))^(-0.5) * exp(-τ)

A(v, w, α) = π^(-0.5) * α * v * QuadGK.quadgk(τ -> A_integrand(v, w, τ), 0, Inf)[1]

# Equation 33: Lowest Free energy E = -B - A where B = -3/(4v)*(v-w)^2.
E(v, w, α) = (3 / (4 * v)) * (v - w)^2 - A(v, w, α)

"""
feynman_variation(α::Float64; v = 7.0, w = 6.0)

    Calculate v and w variational Polaron parameters that minimise the free energy, for the supplied α Frohlich coupling.
    This version uses the original athermal action (Feynman 1955).
    Returns v, w.
"""
function feynman_variation(α; v = 7.0, w = 6.0)
    initial = [v, w]
    lower = [0.0, 0.0]
    upper = [100.0, 100.0]

    f(x) = E(x[1], x[2], α)
    od = Optim.OnceDifferentiable(f, initial; autodiff = :forward)
    solution = Optim.optimize(od, lower, upper, initial, Fminbox(BFGS()))
    v, w = Optim.minimizer(solution)

    return v, w
end


using Optim

export SES, loss, fit, forecast

struct SES
    α::Float64
    l0::Float64

    SES() = new(0, 0)
    SES(α::Number, l0::Number) = new(Float64(α), Float64(l0))
end

function loss(model::SES, time_series)
    α = model.α
    l0 = model.l0

    loss = 0
    N = length(time_series)
    y_pred = 0

    for i in 1:(N)
        if i == 1
            y_pred = l0
        else
            y_pred = time_series[i - 1] * α + y_pred * (1 - α)
        end

        loss += (time_series[i] - y_pred)^2
    end

    return loss
end

function fit(model::SES, y)
    lower = [0.0, -Inf]
    upper = [1.0, Inf]
    initial = [model.α, model.l0]

    function loss_(parameters::Array{Float64})
        α = parameters[1]
        l0 = parameters[2]
        return loss(SES(α, l0), y)
    end

    res = Optim.optimize(loss_, lower, upper, initial)
    optimal = Optim.minimizer(res)
    return SES(optimal[1], optimal[2])
end

function forecast(model, time_series, forecast_length)
    α = model.α
    l0 = model.l0

    N = length(time_series)
    y_pred = 0
    pred = []

    for i in 1:(N)
        if i == 1
            y_pred = l0
        else
            y_pred = time_series[i - 1] * α + y_pred * (1 - α)
        end
    end

    y_pred = time_series[N] * α + y_pred * (1 - α)

    for j in 1:forecast_length
        push!(pred, y_pred)
    end

    return pred
end


using WAV

FILENAME = "british_speaking.wav"

require("k_svd")

wav, FS = wavread(FILENAME)

midpoint = int(size(wav)[1]/2)

# We want to build a dictionary off of the first half of the file
# to see how well we can emulate the sound in the second half of the file

function training_data(num_samples, sample_dimension)
    # Make sure to only get samples from the first half.
    # Second half is what we will test on.
    max_index = midpoint - sample_dimension
    starting_values = rand(1:max_index, num_samples)
    data = [wav[i:(i+sample_dimension-1)] for i in starting_values]

    hcat(data...)
end

function encode_half_2(D, sparsity, seconds=8)
    sample_length = size(D)[1]
    output = zeros(seconds * FS)
    num_frames = int(seconds * FS / sample_length)
    for i = 0:(num_frames-1)
        start_offset = i*sample_length
        end_offset = start_offset + sample_length
        signal_piece = wav[midpoint + start_offset:midpoint + end_offset-1]
        sparse_encoding = orthogonal_matching_persuit(signal_piece, D, sparsity)
        sparsely_encoded = D*sparse_encoding
        output[start_offset+1:end_offset] = sparsely_encoded
    end
    output /= maximum(abs(output[:,1]))
    wavwrite(wav[midpoint:end], "original_second_half.wav", Fs=FS)
    wavwrite(output, "encoded_second_half.wav", Fs=FS)
end

function demo()
    S = training_data(1000, 200);
    sparsity = 40
    # Slow step...
    D = k_svd(S, 30, 500, sparsity);
    encode_half_2(D, sparsity)
end


#Function for getting solution to radial Schrodinger equation using numerov algorithm
#23/07/2017

function numerovfunc(atom,nn,ll,jj)

#Get parameters
ZZ, spin, eneigval, params = GetAtomParams(atom,nn,ll,jj)
alpha_c = params[1]
a_1 = params[2]
a_2 = params[3]
a_3 = params[4]
a_4 = params[5]
r_c = params[6]

#Create vector to integrate over
hh = 0.01 #Choose this to be small so that O(hh^6) may be ignored
h2 = hh^2 #For efficiency
r_i = hh/2 #Define starting point
r_o = 2*nn*(nn+15)  #Define end point
x_i = log(r_i) #Change of variables from Callagher, 2005
x_o = log(r_o)
xx = collect(x_i:hh:x_o)
rr = exp(xx) #For efficiency

#Set up Schrodinger equation from Gallagher (2005), note Pritchard (2012)

LdotS = (jj*(jj+1)-ll*(ll+1)-spin*(spin+1))/2 #Spin-orbit coupling

spinorbitpotential = finestrucconst^2*LdotS./(2*rr.^3) #Fine structure splitting

radialcharge = 1 + (ZZ-1)*exp(-a_1*rr) - rr.*(a_3+a_4*rr).*exp(-a_2*rr) #Effective nuclear charge

coulombpotential = -radialcharge./rr - (alpha_c./(2*(rr.^4))).*(1-exp(-((rr/r_c).^6))) #Corrected Coulomb potential

totalpotential = spinorbitpotential+coulombpotential #total potential

cenfugterm = (ll+1/2)^2 #centrifugal term from Schrodinger equation

#Apply Numerov algorithm
G_x = 2*exp(2*xx).*(totalpotential-eneigval) + cenfugterm #Coeficient in differential equation

T_x = h2*G_x/12 #For efficiency

Ysoln_vec = zeros(size(xx)) #Generate a placeholder vector for solutions

Ysoln_vec[end] = -1e-10 #The function must go to zero at infinity

Ysoln_vec[end-1] = -2e-10

#To perform the iteration with convenient indexing the vectors xx, G_x and soln_vec are flipped

fYsoln = flipdim(Ysoln_vec,1)
fT_x = flipdim(T_x,1)

for ii = 3:length(xx)
        fYsoln[ii] = ((2+10*fT_x[ii-1])*fYsoln[ii-1] -  (1-fT_x[ii-2])*fYsoln[ii-2])/(1-fT_x[ii])
end

Ysoln_vec = flipdim(fYsoln,1) #Return solutions to proper ordering

#Normalise (method adapted from Gallagher, 2005)
normconstvec = zeros(length(rr)-1)
for mm = 2:length(rr)
    deltar = rr[mm]-rr[mm-1]
    normconstvec[mm-1] = (Ysoln_vec[mm]^2)*rr[mm]*deltar
end

normconst = sqrt(sum(normconstvec))
normY_sol = Ysoln_vec/normconst

#Saving and loading seems to hurt performance
#filename = filenamemake(atom,nn,ll,jj)

#save(filename,"rr",rr,"normY_sol",normY_sol)

return normY_sol, rr
end


using LinearAlgebra, PolynomialRoots
include("findlocalmaxima.jl")

function rootMUSIC(Y,K,SNR_db=20,m=0)
    # Estimation of signal parameters via rotational invariance techniques
    # (ESPRIT) algorithm for line spectral estimation from measurements.
    # Can handle multiple snapshots by constructing a low rank Hankel matrix
    # from the data.
    #
    # Inputs:
    # Y: Data matrix
    # K: number of sinusoids
    # SNR_db: SNR in each dimension of the data in dB
    # m: Hankel parameter of the Hankel matrix of the data
    #
    # Parameters:
    # N = size(Y,1): number of measurements per snapshot
    # L = size(Y,2): number of snapshots
    #
    # Output:
    # f_hat: estimate of the frequencies of the K sinusoids in [rad]

(N,L) = size(Y)

if m == 0
m = ceil(Int,N/2)
end

R = zeros(m,m)
for j = 1:L, i = m:N
    R = R + Y[i-m+1:i,j]*Y[i-m+1:i,j]'/(L*N)
end

    evdR = eigen(R)
    p = sortperm(real(evdR.values))[1:m-K]
    Un = evdR.vectors[:,p]

    M = Un*Un'

    # Construct polynomial
    polcoeffs = Array{Complex{Float64},1}(undef,2*m-1)
    for k = 1:2*m-1
        polcoeffs[k] = sum(diag(M,k-m))
    end

    # Find roots of polynomial
    polyroots = roots(polcoeffs)

    # Remove roots with multiplicity > 1
    idcs = []
    for i = 1:length(polyroots)-1, j = i+1:length(polyroots)
        if abs(polyroots[i]-polyroots[j]) < 0.5/SNR_db + 0.001
            push!(idcs,i)
        end
    end

    polyroots2 = polyroots[setdiff(1:length(polyroots),idcs)]

    # Find K roots closest to unit circle
    dist2UnitCircle = abs.(abs.(polyroots2) .- 1)
    p = sortperm(dist2UnitCircle)[1:K]
    closestRoots = polyroots2[p]

    # Calculate frequencies from roots
    f_hat = (angle.(closestRoots)/(2*pi) .+ 1).%1
end


export propagate
export fresnel_kernel, rs_kernel
export point_source_propagate
export lens_propagate, four_f_propagate, mla_propagate

"""
    fresnel_kernel(fx, fy, z, λ, n=1.0)

Calculates the Fresnel propagation kernel in Fourier space.
`fx`, `f_y` are the spatial frequency. `z` is the propagation distance.
`λ` is the wavelength and `n` the refractive index.
"""
function fresnel_kernel(fx, fy, z, λ, n=1)
    κ = calc_κ(λ, n)
    return @. exp(1im * 2 * π * κ * z) * exp(-1im * π * z / κ * (fx^2 + fy^2))
end

"""
    rs_kernel(fx, fy, z, λ, n=1.0)

Calculates the Rayleigh-Sommerfeld propagation kernel in Fourier space.
`fx`, `f_y` are the spatial frequencies. `z` is the propagation distance.
`λ` is the wavelength and `n` the refractive index.
"""
function rs_kernel(fx, fy, z, λ=λ0, n=1)
    κ = calc_κ(λ, n)
    c = complex(κ^2 - (fx^2 + fy^2))
    return exp(1im * z * 2 * π * sqrt(c))
end

"""
    propagate!(arr, L, z; kernel=rs_kernel, λ=550e-9, n=1)

Propagate an electric field `arr` with a array size (in physical dimensions like meter etc) of `L` over the distance `z`.
Per default `kernel=rs_kernel` (Rayleigh-Sommerfeld) is the propagation kernel.
`λ` is the wavelength and `n` the refractive index.
"""
function propagate!(arr, L, z; kernel=rs_kernel, λ=λ0, n=1, fft_plan=plan_fft!(arr))
    # trivial propagation
    if iszero(z)
        # return a copy instead of arr
        # therefore we can rely on, that we don't modify the initial array
        return copy(arr)
    end

    # array with the frequencies in Fourier space
    freq_x = to_gpu_or_cpu(arr, Zygote.@ignore eltype(arr).(fftfreq(size(arr)[2], size(arr)[2] / L)'))
    freq_y = to_gpu_or_cpu(arr, Zygote.@ignore eltype(arr).(fftfreq(size(arr)[1], size(arr)[1] / L)))
    arr_ft = fft_plan * arr

    κ = eltype(arr)(calc_κ(λ, n))
    c_sqrt = Zygote.@ignore (κ^2 .- (freq_x.^2 .+ freq_y.^2))
    c = Zygote.@ignore 1im * eltype(arr)(z * 2π)
    c_exp = Zygote.@ignore (exp.(c .* sqrt.(c_sqrt)))
    out_ft = arr_ft .* c_exp
    #out_ft = arr_ft .* kernel.(freq_x, freq_y, Ref(z), Ref(λ), Ref(n))
    p_inv = Zygote.@ignore inv(fft_plan)
    out = p_inv.scale .* (p_inv.p * out_ft)
    return out
end

function propagate(arr, L, z; kernel=rs_kernel, λ=λ0, n=1)
    propagate!(complex.(arr), L, z, kernel=kernel, λ=λ, n=n)
end

function propagate_tilted_plane(arr, L, z, tilting_ϕ, dim=1; kernel=rs_kernel, λ=λ0, n=1, fft_plan=plan_fft!(arr))
    Zygote.@ignore GC.gc(true)
    Δzs = z .+ fftpos(L, size(arr, 1)) .* tan(deg2rad(tilting_ϕ))

    arr2 = propagate!(arr, L, Δzs[1], fft_plan=fft_plan)
    if dim == 1
        out = arr2[1:1, :]
    else
        out = arr2[:, 1:1]
    end
    for (i, Δz) in enumerate(Δzs)
        Δz = Δzs.step
        if isone(i)
            continue
        end
        arr2 = propagate!(arr2, L, Δz, kernel=rs_kernel, λ=λ, n=n, fft_plan=fft_plan)
        if dim == 1
            out = vcat(out, arr2[i:i, :])
        else
            out = hcat(out, arr2[:, i:i])
        end
    end
    return out
end

"""
    point_source_propagate(L, size, point; λ=550e-9, n=1, dtype=ComplexF64)

Propagate a point source in a field of width `L` with array size `size`
over a distance `z` in a medium with refractive index `n`.
Point source is located at position `point(x, y, z)` and then propagated to
a location z=0. `z<0` in `point` means forward propagation.

This is based on the analytical solution in real space and not on Fourier
space propagation. The latter one suffers from artifacts while microscopic large `L`.
This function should be always preferred for point sources.
"""
function point_source_propagate(L, size, point::Point; λ=λ0, n=1, dtype=ComplexF64)
    x0, y0, z = point.x, point.y, point.z
    out = zeros(dtype, size)
    k = calc_k(λ, n)
    # if we are at the plane of the point source
    # return numerical delta
    if iszero(z)
        j = argmin(map(x -> abs(x - x0), fftpos(L, size[2])))
        i = argmin(map(y -> abs(y - y0), fftpos(L, size[1])))
        out[i, j] = 1
        return out
    end
    # calculate values
    @inbounds for (j, x) in enumerate(fftpos(L, size[2]))
        for (i, y) in enumerate(fftpos(L, size[1]))
            # z>0 means that we back propagate the PS to z=0
            r = dtype(sign(-z) * sqrt((x-x0)^2 + (y-y0)^2 + z^2))
            out[i, j] = 1/r .* exp(1im * k * r);
        end
    end
    out = out ./ out[argmax(abs2.(out))]
    return out
end

"""
    lens_propagate(arr, L, f; λ=550e-9, n=1, d=nothing)

Propagate an electrical field of width `L` at a distance `d` in
front of a lens with focal length `f` in medium with refractive index `n`
to the back focal plane of the lens.
Per default `d=nothing` meaning that we set `d=f`.
Assume same sampling and size in both dimensions.
As return parameter we get the resulting field and the new field size.
Essentially the lens (assuminig infinite large of that lens) performs
a scaled Fourier transform. Therefore we get a new field size.

 # Reference
Based on:
* "Computational Fourier Optics. A MATLAB Tutorial", D. Voelz, (2011).
* Goodman, Joseph W. Introduction to Fourier optics
"""
function lens_propagate(arr, L, f; λ=λ0, n=1, d=nothing, fft_plan=plan_fft!(complex.(arr)))
    if isnothing(d)
        d = f
    end
    κ = calc_κ(λ, n)
    # shift to center again
    if typeof(arr) <: Array{<:Real}
        arr = complex.(arr)
    end
    out_f = fftshift(fft_plan * ifftshift(arr))
    dx = L / size(arr)[2]
    dy = L / size(arr)[1]

    # lens performs scaled fourier transform. Therefore new field size
    L_new = λ * f / dx
    aux_x =  Zygote.@ignore fftpos(L_new, size(arr)[2])'
    aux_y = Zygote.@ignore fftpos(L_new, size(arr)[1])
    x = Zygote.@ignore to_gpu_or_cpu(arr, aux_x)
    y = Zygote.@ignore to_gpu_or_cpu(arr, aux_y)
    c_exp = eltype(arr)(π * κ / f * (1 - d / f))
    c = Zygote.@ignore  .- 1im ./ (λ .* f) .* cis.(c_exp .* (x.^2 .+ y.^2))
    out = out_f .* c .* eltype(out_f)(dx * dy)

    return out, L_new
end

"""
    four_f_propagate(arr, L, f1, f2, NA; fft_plan=fft_plan, d=nothing)

Propagate the electrical field `arr` (field size `L`) from front
focal plane to the back focal plane of a 4f system.
The focal length of the first lens is `f1` and the second lens `f`.
Magnification of the system is then `f2/f1`.
`d` is the distance from lens to the point of the `arr` in front
of the lens. If `d == nothing` then we propagate from object plane to
image plane.

"""
function four_f_propagate(arr, L, f1, f2, NA; d=nothing, fft_plan=plan_fft!(complex.(arr)))
	radius = NA * 2 * f1 / 2

    out = lens_propagate(arr, L, f1, fft_plan=fft_plan, d=d)
    E1 = out[1]
    L1 = out[2]
    # check that the field size is large enough that the pupil fits
    @assert radius < L1 / 2
	E1_ = circ(E1, L1, radius)

    # tuple of E2 and L2
    out2 = lens_propagate(E1_, L1, f2, fft_plan=fft_plan)
    return out2
end

"""
    mla_propagate(E, L, fmla, p_mla, z1, z2)

Propagate an electrical field `E` placed a distance `z1` in front
of a microlns array, to a distance `z2` behind the microlens array.
`L` is the size of the field and `fmla` is the focal length of the lenslets.
`p_mla` is microlenses pitch.
"""
function mla_propagate(E, L, fmla, p_mla, z1, z2)

    E_front = propagate(E, L, z1)
	mla = micro_lens_array(fmla, L, p_mla, size(E))
	E_mla = E_front .* mla
	E_behind = propagate(E_mla, L, z2);
	return E_behind, real(mla)
end


s = open("input.txt") do file
    readlines(file);
end

function EvalExpression(expr)
    parts = split(expr, " ");
    curVal = parse(Int, parts[1])

    i = 2;
    while i <= length(parts)
        op = parts[i];
        num = parse(Int, parts[i+1]);

        if op == "+"
            curVal += num;
        else
            curVal *= num;
        end

        i += 2;
    end

    return curVal
end

function DoTheMath(expression)
    expressionStack = [];
    currentExpr = "";
    for i = 1:length(expression)
        if expression[i] == '('
            push!(expressionStack, currentExpr);
            currentExpr = "";
        elseif expression[i] == ')'
            val = EvalExpression(currentExpr);
            currentExpr = pop!(expressionStack) * "$(val)";
        else
            currentExpr *= string(expression[i])
        end
    end

    val = EvalExpression(currentExpr);
    return val;
end

answer = sum(map(DoTheMath, s))


using Deconvolution
using Test, Random, FFTW

Random.seed!(42) # Fixed random seed

##### Wiener deconvolution

x = range(0, stop = 10, length = 15)
s = sinpi.(x) .- 1.5cos.(x)
n = rand(length(s))
blurring_ft  = exp.(-0.001 * (range(-div(length(x), 2), length = length(x)) .^ 2) .^ (5 // 6))
blurred_s_ft = fftshift(blurring_ft) .* fft(s) .+ fft(n)
blurred_s    = real(ifft(blurred_s_ft))
blurring     = ifft(fftshift(blurring_ft))
@testset "No blurring with noise scalar parameter" begin
    @test wiener(s + n, s, 0.5) ≈
        [-1.1111473161507488, -0.17041312679112874, -1.3469825775608049, 1.2447585407259756,
          2.578860561637874, 1.09191362235022, 1.5339504088753226, -0.34582554585559355,
          -1.316232389748769, -0.21271451029702804, -0.6157543266519075, -0.3338395657941805,
          2.1195369797597614, 1.060270051101612, 1.4446941593151053]
end

@testset "Blurring with noise scalar parameter" begin
    @test wiener(blurred_s, s, 0.5, blurring) ≈
        [-1.1086114591808454, -0.17138676636801406, -1.3482585663111306, 1.2427663740234156,
         2.5813314219359276, 1.0940885915582341, 1.532411498782748, -0.34834962964626176,
         -1.311080379275804, -0.21869338550577228, -0.6084828281998079, -0.3366932366271657,
         2.117735346098039, 1.0624233782615877, 1.4429704301669912]
end

@testset "No blurring with noise array" begin
    @test wiener(blurred_s, s, ifftshift(n)) ≈ [-1.4080797956919306, -0.49286583995619976,
        -1.6205631520484833, 0.7973315264441291, 2.041487785980698, 0.7068033074341512,
        1.303796832927927, -0.5691518774600562, -1.7799047917607576, -0.7841432287697161,
        -0.9149839099122377, -0.4555342804458623, 1.761794603633166, 0.5975996848388563, 0.9542148100114622]
end

@testset "Blurring with noise array" begin
    @test wiener(blurred_s, s, ifftshift(n), blurring) ≈ [-1.4250670444336333, -0.48221282618655026,
        -1.636981677852502, 0.8020690789792746, 2.053748123047759, 0.6974830808841243, 1.3177198257314031,
        -0.5714565338646487, -1.7922287374907637, -0.7809316185676288, -0.9185450961060387,
        -0.46118637814202684, 1.7785714975739182, 0.5887392863050276, 0.967947374977039]
end

##### Richardson-Lucy deconvolution

@testset "Richardson-Lucy deconvolution tests" begin
    k=-0.001
    blurring_ft  = exp.(-k * (range(-div(length(x), 2), length = length(x)) .^ 2) .^ (5 // 6))
    blurred_s_ft = fftshift(blurring_ft) .* fft(s)
    blurred_s    = real(ifft(blurred_s_ft))
    blurring     = ifft(fftshift(blurring_ft))

    estimated = lucy(blurred_s, blurring, iterations=1)

    @test sum(abs.(s .- estimated)) < sum(abs.(s .- blurred_s))

    k =-0.0015 # should improve also when blur params are not perfectly estimated
    blurring_ft  = exp.(-k * (range(-div(length(x), 2), length = length(x)) .^ 2) .^ (5 // 6))
    blurring     = ifft(fftshift(blurring_ft))

    estimated = lucy(blurred_s, blurring, iterations=1)

    @test sum(abs.(s .- estimated)) < sum(abs.(s .- blurred_s))
end