using Test, GeometryBasics

@testset "algorithms.jl" begin
    cube = Rect(Vec3f0(-0.5), Vec3f0(1))
    cube_faces = decompose(TriangleFace{Int}, faces(cube))
    cube_vertices = decompose(Point{3,Float32}, cube)
    @test area(cube_vertices, cube_faces) == 6
    mesh = Mesh(cube_vertices, cube_faces)
    @test GeometryBasics.volume(mesh) ≈ 1
end

@testset "Cylinder" begin
    @testset "constructors" begin
        o, extr, r = Point2f0(1, 2), Point2f0(3, 4), 5.0f0
        s = Cylinder(o, extr, r)
        @test typeof(s) == Cylinder{2,Float32}
        @test typeof(s) == Cylinder2{Float32}
        @test origin(s) == o
        @test extremity(s) == extr
        @test radius(s) == r
        #@test abs(height(s)- norm([1,2]-[3,4]))<1e-5
        h = norm(o - extr)
        @test isapprox(height(s), h)
        #@test norm(direction(s) - Point{2,Float32}([2,2]./norm([1,2]-[3,4])))<1e-5
        @test isapprox(direction(s), Point2f0(2, 2) ./ h)
        v1 = rand(Point{3,Float64})
        v2 = rand(Point{3,Float64})
        R = rand()
        s = Cylinder(v1, v2, R)
        @test typeof(s) == Cylinder{3,Float64}
        @test typeof(s) == Cylinder3{Float64}
        @test origin(s) == v1
        @test extremity(s) == v2
        @test radius(s) == R
        @test height(s) == norm(v2 - v1)
        #@test norm(direction(s) - Point{3,Float64}((v2-v1)./norm(v2-v1)))<1e-10
        @test isapprox(direction(s), (v2 - v1) ./ norm(v2 .- v1))
    end

    @testset "decompose" begin

        o, extr, r = Point2f0(1, 2), Point2f0(3, 4), 5.0f0
        s = Cylinder(o, extr, r)
        positions = Point{3,Float32}[(-0.7677671, 3.767767, 0.0),
                                     (2.767767, 0.23223293, 0.0),
                                     (0.23223293, 4.767767, 0.0),
                                     (3.767767, 1.2322329, 0.0), (1.2322329, 5.767767, 0.0),
                                     (4.767767, 2.232233, 0.0)]
        @test decompose(Point3f0, Tesselation(s, (2, 3))) ≈ positions

        FT = TriangleFace{Int}
        faces = FT[(1, 2, 4), (1, 4, 3), (3, 4, 6), (3, 6, 5)]
        @test faces == decompose(FT, Tesselation(s, (2, 3)))

        v1 = Point{3,Float64}(1, 2, 3)
        v2 = Point{3,Float64}(4, 5, 6)
        R = 5.0
        s = Cylinder(v1, v2, R)
        positions = Point{3,Float64}[(4.535533905932738, -1.5355339059327373, 3.0),
                                     (7.535533905932738, 1.4644660940672627, 6.0),
                                     (3.0412414523193148, 4.041241452319315,
                                      -1.0824829046386295),
                                     (6.041241452319315, 7.041241452319315,
                                      1.9175170953613705),
                                     (-2.535533905932737, 5.535533905932738,
                                      2.9999999999999996),
                                     (0.46446609406726314, 8.535533905932738, 6.0),
                                     (-1.0412414523193152, -0.04124145231931431,
                                      7.0824829046386295),
                                     (1.9587585476806848, 2.9587585476806857,
                                      10.08248290463863), (1, 2, 3), (4, 5, 6)]

        @test decompose(Point3{Float64}, Tesselation(s, 8)) ≈ positions

        faces = TriangleFace{Int}[(3, 2, 1), (4, 2, 3), (5, 4, 3), (6, 4, 5), (7, 6, 5),
                                  (8, 6, 7), (1, 8, 7), (2, 8, 1), (3, 1, 9), (2, 4, 10),
                                  (5, 3, 9), (4, 6, 10), (7, 5, 9), (6, 8, 10), (1, 7, 9),
                                  (8, 2, 10)]
        @test faces == decompose(TriangleFace{Int}, Tesselation(s, 8))

        m = triangle_mesh(Tesselation(s, 8))

        @test GeometryBasics.faces(m) == faces
        @test GeometryBasics.coordinates(m) ≈ positions
        m = normal_mesh(s)# just test that it works without explicit resolution parameter
        @test m isa GLNormalMesh

        muv = uv_mesh(s)
        @test Rect(Point.(texturecoordinates(muv))) == FRect2D(Vec2f0(0), Vec2f0(1.0))
    end
end

@testset "HyperRectangles" begin
    a = Rect(Vec(0, 0), Vec(1, 1))
    pt_expa = Point{2,Int}[(0, 0), (1, 0), (0, 1), (1, 1)]
    @test decompose(Point{2,Int}, a) == pt_expa
    mesh = normal_mesh(a)
    @test decompose(Point2f0, mesh) == pt_expa

    b = Rect(Vec(1, 1, 1), Vec(1, 1, 1))
    pt_expb = Point{3,Int64}[[1, 1, 1], [1, 1, 2], [1, 2, 2], [1, 2, 1], [1, 1, 1],
                             [2, 1, 1], [2, 1, 2], [1, 1, 2], [1, 1, 1], [1, 2, 1],
                             [2, 2, 1], [2, 1, 1], [2, 2, 2], [1, 2, 2], [1, 1, 2],
                             [2, 1, 2], [2, 2, 2], [2, 1, 2], [2, 1, 1], [2, 2, 1],
                             [2, 2, 2], [2, 2, 1], [1, 2, 1], [1, 2, 2]]
    @test decompose(Point{3,Int}, b) == pt_expb
    mesh = normal_mesh(b)
end

NFace = NgonFace

@testset "Faces" begin
    @test convert_simplex(GLTriangleFace, QuadFace{Int}(1, 2, 3, 4)) ==
          (GLTriangleFace(1, 2, 3), GLTriangleFace(1, 3, 4))
    @test convert_simplex(NFace{3,ZeroIndex{Int}}, QuadFace{ZeroIndex{Int}}(1, 2, 3, 4)) ==
          (NFace{3,ZeroIndex{Int}}(1, 2, 3), NFace{3,ZeroIndex{Int}}(1, 3, 4))
    @test convert_simplex(NFace{3,OffsetInteger{3,Int}},
                          NFace{4,OffsetInteger{2,Int}}(1, 2, 3, 4)) ==
          (NFace{3,OffsetInteger{3,Int}}(1, 2, 3), NFace{3,OffsetInteger{3,Int}}(1, 3, 4))
    @test convert_simplex(LineFace{Int}, QuadFace{Int}(1, 2, 3, 4)) ==
          (LineFace{Int}(1, 2), LineFace{Int}(2, 3), LineFace{Int}(3, 4),
           LineFace{Int}(4, 1))

    @testset "NgonFace ambiguity" begin
        face = NgonFace((1, 2))
        @test convert_simplex(NgonFace{2,UInt32}, face) === (NgonFace{2,UInt32}((1, 2)),)
        @test convert_simplex(typeof(face), face) === (face,)
        face = NgonFace((1,))
        @test convert_simplex(NgonFace{1,UInt32}, face) === (NgonFace{1,UInt32}((1,)),)
        @test convert_simplex(typeof(face), face) === (face,)
    end
end

@testset "Normals" begin
    n64 = Vec{3,Float64}[(0.0, 0.0, -1.0), (0.0, 0.0, -1.0), (0.0, 0.0, -1.0),
                         (0.0, 0.0, -1.0), (0.0, 0.0, 1.0), (0.0, 0.0, 1.0),
                         (0.0, 0.0, 1.0), (0.0, 0.0, 1.0), (-1.0, 0.0, 0.0),
                         (-1.0, 0.0, 0.0), (-1.0, 0.0, 0.0), (-1.0, 0.0, 0.0),
                         (1.0, 0.0, 0.0), (1.0, 0.0, 0.0), (1.0, 0.0, 0.0), (1.0, 0.0, 0.0),
                         (0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0),
                         (0.0, -1.0, 0.0), (0.0, -1.0, 0.0), (0.0, -1.0, 0.0),
                         (0.0, -1.0, 0.0),]
    n32 = map(Vec{3,Float32}, n64)
    r = triangle_mesh(centered(Rect3D))
    # @test normals(coordinates(r), GeometryBasics.faces(r)) == n32
    # @test normals(coordinates(r), GeometryBasics.faces(r)) == n64
end

@testset "HyperSphere" begin
    sphere = Sphere{Float32}(Point3f0(0), 1.0f0)

    points = decompose(Point, Tesselation(sphere, 3))
    point_target = Point{3,Float32}[[0.0, 0.0, 1.0], [1.0, 0.0, 6.12323e-17],
                                    [1.22465e-16, 0.0, -1.0], [-0.0, 0.0, 1.0],
                                    [-1.0, 1.22465e-16, 6.12323e-17],
                                    [-1.22465e-16, 1.49976e-32, -1.0], [0.0, -0.0, 1.0],
                                    [1.0, -2.44929e-16, 6.12323e-17],
                                    [1.22465e-16, -2.99952e-32, -1.0]]
    @test points ≈ point_target

    f = decompose(TriangleFace{Int}, Tesselation(sphere, 3))
    face_target = TriangleFace{Int}[[1, 2, 5], [1, 5, 4], [2, 3, 6], [2, 6, 5], [4, 5, 8],
                                    [4, 8, 7], [5, 6, 9], [5, 9, 8]]
    @test f == face_target
    circle = Circle(Point2f0(0), 1.0f0)
    points = decompose(Point2f0, Tesselation(circle, 20))
    @test length(points) == 20
    tess_circle = Tesselation(circle, 32)
    mesh = triangle_mesh(tess_circle)
    @test decompose(Point2f0, mesh) ≈ decompose(Point2f0, tess_circle)
end

@testset "Rectangles" begin
    rect = FRect2D(0, 7, 20, 3)
    @test (rect + 4) == FRect2D(4, 11, 20, 3)
    @test (rect + Vec(2, -2)) == FRect2D(2, 5, 20, 3)

    @test (rect - 4) == FRect2D(-4, 3, 20, 3)
    @test (rect - Vec(2, -2)) == FRect2D(-2, 9, 20, 3)

    base = Vec3f0(1, 2, 3)
    wxyz = Vec3f0(-2, 4, 2)
    rect = FRect3D(base, wxyz)
    @test (rect + 4) == FRect3D(base .+ 4, wxyz)
    @test (rect + Vec(2, -2, 3)) == FRect3D(base .+ Vec(2, -2, 3), wxyz)

    @test (rect - 4) == FRect3D(base .- 4, wxyz)
    @test (rect - Vec(2, -2, 7)) == FRect3D(base .- Vec(2, -2, 7), wxyz)

    rect = FRect2D(0, 7, 20, 3)
    @test (rect * 4) == FRect2D(0, 7 * 4, 20 * 4, 3 * 4)
    @test (rect * Vec(2, -2)) == FRect2D(0, -7 * 2, 20 * 2, -3 * 2)

    base = Vec3f0(1, 2, 3)
    wxyz = Vec3f0(-2, 4, 2)
    rect = FRect3D(base, wxyz)
    @test (rect * 4) == FRect3D(base .* 4, wxyz .* 4)
    @test (rect * Vec(2, -2, 3)) == FRect3D(base .* Vec(2, -2, 3), wxyz .* Vec(2, -2, 3))

    rect1 = Rect(Vec(0.0, 0.0), Vec(1.0, 2.0))
    rect2 = Rect(0.0, 0.0, 1.0, 2.0)
    @test rect1 isa GeometryBasics.HyperRectangle{2,Float64}
    @test rect1 == rect2

    split1, split2 = GeometryBasics.split(rect1, 2, 1)
    @test widths(split1) == widths(split2)
    @test origin(split1) == Vec(0, 0)
    @test origin(split2) == Vec(0, 1)
    @test in(split1, rect1)
    @test !in(rect1, split1)

    prim = Rect(0.0, 0.0, 1.0, 1.0)
    @test length(prim) == 2

    @test width(prim) == 1.0
    @test height(prim) == 1.0

    b1 = Rect2D(0.0, 0.0, 2.0, 2.0)
    b2 = Rect2D(0, 0, 2, 2)
    @test isequal(b1, b2)

    pt = Point(1.0, 1.0)
    b1 = Rect(0.0, 0.0, 1.0, 1.0)
    @test in(pt, b1)

    rect = Rect(0.0, 0.0, 1.0, 1.0)
    @test GeometryBasics.positive_widths(rect) isa GeometryBasics.HyperRectangle{2,Float64}

    h1 = Rect(0.0, 0.0, 1.0, 1.0)
    h2 = Rect(1.0, 1.0, 2.0, 2.0)
    @test union(h1, h2) isa GeometryBasics.HyperRectangle{2,Float64}
    @test GeometryBasics.diff(h1, h2) == h1
    @test GeometryBasics.intersect(h1, h2) isa GeometryBasics.HyperRectangle{2,Float64}

    b = Rect(0.0, 0.0, 1.0, 1.0)
    v = Vec(1, 2)
    @test update(b, v) isa GeometryBasics.HyperRectangle{2,Float64}
    v = Vec(1.0, 2.0)
    @test update(b, v) isa GeometryBasics.HyperRectangle{2,Float64}

    p = Vec(5.0, 4.0)
    rect = Rect(0.0, 0.0, 1.0, 1.0)
    @test min_dist_dim(rect, p, 1) == 4.0
    @test min_dist_dim(rect, p, 2) == 3.0
    @test max_dist_dim(rect, p, 1) == 5.0
    @test max_dist_dim(rect, p, 2) == 4.0

    rect1 = Rect(0.0, 0.0, 1.0, 1.0)
    rect2 = Rect(3.0, 1.0, 4.0, 2.0)
    @test min_dist_dim(rect1, rect2, 1) == 2.0
    @test min_dist_dim(rect1, rect2, 2) == 0.0
    @test max_dist_dim(rect1, rect2, 1) == 7.0
    @test max_dist_dim(rect1, rect2, 2) == 3.0

    @test !before(rect1, rect2)
    rect1 = Rect(0.0, 0.0, 1.0, 1.0)
    rect2 = Rect(3.0, 2.0, 4.0, 2.0)
    @test before(rect1, rect2)

    @test !meets(rect1, rect2)
    rect2 = Rect(1.0, 1.0, 4.0, 2.0)
    @test meets(rect1, rect2)

    rect1 = Rect(1.0, 1.0, 2.0, 2.0)
    rect2 = Rect(0.0, 0.0, 2.0, 1.0)
    @test !overlaps(rect1, rect2)
    rect1 = Rect(1.0, 1.0, 2.0, 2.0)
    rect2 = Rect(1.5, 1.5, 2.0, 2.0)
    @test overlaps(rect1, rect2)

    rect1 = Rect(1.0, 1.0, 2.0, 2.0)
    rect2 = Rect(0.0, 0.0, 2.0, 1.0)
    @test !GeometryBasics.starts(rect1, rect2)
    rect2 = Rect(1.0, 1.0, 1.5, 1.5)
    @test !GeometryBasics.starts(rect1, rect2)
    rect2 = Rect(1.0, 1.0, 3.0, 3.0)
    @test GeometryBasics.starts(rect1, rect2)

    rect1 = Rect(1.0, 1.0, 2.0, 2.0)
    rect2 = Rect(0.0, 0.0, 4.0, 4.0)
    @test during(rect1, rect2)
    rect1 = Rect(0.0, 0.0, 2.0, 3.0)
    rect2 = Rect(1.0, 1.0, 4.0, 2.0)
    @test !during(rect1, rect2)

    rect1 = Rect(1.0, 1.0, 2.0, 2.0)
    rect2 = Rect(0.0, 0.0, 4.0, 4.0)
    @test !finishes(rect1, rect2)
    rect1 = Rect(1.0, 0.0, 1.0, 1.0)
    rect2 = Rect(0.0, 0.0, 2.0, 1.0)
    @test !finishes(rect1, rect2)
    rect1 = Rect(1.0, 1.0, 1.0, 2.0)
    rect2 = Rect(0.0, 0.0, 2.0, 3.0)
    @test finishes(rect1, rect2)

end


export bode

"""
### function bode(sys::StateSpace, w::Array=[]; plot=true)

Compute the bode diagram.

##### Args

* sys: System in state space representation.
* w: Array of frequencies in which the bode diagram will be computed. If empty,
  the frequencies will be chosen according to the system poles.

##### Parameters

* plot: Plot the bode diagram (default = `true`).

##### Returns

Gain (dB), Phase (deg), w (freq. array) for each input/output pair.

"""

function bode(sys::StateSpace, w::Array=[]; plot=true)
    A, B, C, D = sys.A, sys.B, sys.C, sys.D

    # Number of inputs.
    num_u = size(B,2)

    # Number of outputs.
    num_y = size(C,1)

    # Number of states.
    num_states = size(A,1)

    # Identity matrix.
    eye_s = eye(num_states)

    # Check if the vector of frequencies was provided.
    if (isempty(w))
        # If not, compute the bode diagram from a frequency two decades lower
        # than the minimum frequency and two decades higher than the maximum
        # frequency.
        poles = pole(sys)

        # Get the minimum and maximum frequency of the system removing the poles
        # at origin.
        highest_freq  = maximum(abs.(real(poles[poles .!= 0])))
        smallest_freq = minimum(abs.(real(poles[poles .!= 0])))

        w_initial = (smallest_freq != 0) ? smallest_freq/100.0 : 1e-2
        w_final   = (highest_freq != 0)  ?  highest_freq*100.0 : 1e+2

        # Compute 500 logarithmically spaced points.
        w = logspace(log10(w_initial), log10(w_final), 500)
    end

    # Frequency response for each pair input/output.
    gain  = Vector{Vector{Float64}}(length(w))
    phase = Vector{Vector{Float64}}(length(w))

    for i = 1:length(w)
        s = Complex(0,w[i])

        # `fr` is always a vector. Hence, we can squeeze one dimension here.
        fr = squeeze(C*inv(s*eye_s-A)*B+D,1)

        gain[i]  = 20*log10.([ norm(f) for f in fr ])
        phase[i] = [ angle(f)*180.0/pi for f in fr ]

        if i > 1
            # Make sure that the phase plot is continuous. This can be done by
            # analysing the difference between the current phase and the latest
            # one and checking if it is higher than +180.0 or lower than -180.0.
            Δ = phase[i] - phase[i-1]
            phase[i] += (Δ .> 180.0) * (-360.0) + (Δ .< -180.0) * (+360.0)
        end
    end

    ( plot ) && plot_bode(gain, phase, w, num_u, num_y)

    gain, phase, w
end

"""
### function bode(G::TransferFunction, w::Array=[]; plot=true)

Compute the bode diagram.

##### Args

* G: Transfer function of the dynamic system.
* w: Array of frequencies in which the bode diagram will be computed. If empty,
  the frequencies will be chosen according to the system poles.

##### Parameters

* plot: Plot the bode diagram (default = `true`).

##### Returns

Gain (dB), Phase (deg), w (freq. array) for each input/output pair.

"""

function bode(G::TransferFunction, w::Array=[]; plot=true)
    # Convert the transfer function to state space and compute the Bode diagram.
    bode(tf2ss(G), w; plot=plot)
end


# ---
# title: 724. Find Pivot Index
# id: problem724
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/find-pivot-index/description/>
# hidden: true
# ---
#
# Given an array of integers `nums`, write a method that returns the "pivot"
# index of this array.
#
# We define the pivot index as the index where the sum of all the numbers to the
# left of the index is equal to the sum of all the numbers to the right of the
# index.
#
# If no such index exists, we should return -1. If there are multiple pivot
# indexes, you should return the left-most pivot index.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,7,3,6,5,6]
#     Output: 3
#     Explanation:
#     The sum of the numbers to the left of index 3 (nums[3] = 6) is equal to the sum of numbers to the right of index 3.
#     Also, 3 is the first index where this occurs.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,2,3]
#     Output: -1
#     Explanation:
#     There is no index that satisfies the conditions in the problem statement.
#
#
#
#
# **Constraints:**
#
#   * The length of `nums` will be in the range `[0, 10000]`.
#   * Each element `nums[i]` will be an integer in the range `[-1000, 1000]`.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# n_x = 5
# n_h = 4
# n_y = 1

# W1 = randn(n_h, n_x)
# println(W1)

import Random
Random.seed!(210)

# function initialize_parameters(layer_dims)
#     parameters = Dict{String, Array{Float64}}()
#     for i = 2:length(layer_dims)
#         parameters[string("W", i-1)] = randn(layer_dims[i], layer_dims[i-1]) * 0.01
#         parameters[string("b", i-1)] = zeros(layer_dims[i], 1)
#     end
#     return parameters
# end

include("NeuralNetwork.jl")

parameters = initialize_parameters([5 10 1])

function printdict(dict)
    for (key, value) in dict
        println(key, " ", size(value))
    end
end
function printdict_full(dict)
    for (key, value) in dict
        println(key, " ", (value))
    end
end
printdict(parameters)

# using LinearAlgebra

# function sigmoid(X)
#     return 1 ./ (1 .+ exp.(-X))
# end
# function sigmoid(x)
#     return 1 / (1 + exp(-x))
# end

# function relu(x)
#     if x>0
#         return x
#     else
#         return 0
#     end
# end

# function forward_prop(X::Array{Float64}, parameters::Dict{String, Array{Float64}}, layer_dims::Array{Int}, activations=Nothing)
#     num_layers = length(layer_dims)

#     if activations === Nothing
#         activations = Dict()
#         for i = 1:num_layers-1
#             # println("relu")
#             activations[i] = relu
#         end
#         # println(sigmoid)
#         activations[num_layers] = sigmoid
#     end
#     # println(activations)
#     caches = Dict()

#     Z1 = parameters[string("W", 1)] * X .+ parameters[string("b", 1)]
#     A1 = activations[1].(Z1)
#     caches[string("Z", 1)] = Z1
#     caches[string("A", 1)] = A1

#     Ai = Nothing
#     for i = 2:num_layers-1
#         Zi = parameters[string("W", i)] * caches[string("A", i-1)] .+ parameters[string("b", i)]
#         Ai = activations[i].(Zi)
#         caches[string("Z", i)] = Zi
#         caches[string("A", i)] = Ai
#     end

#     return Ai, caches
# end

X = randn(5, 100)
a2, caches = forward_prop(X, parameters, [5 10 1], (relu, sigmoid))
# a2, caches = forward_prop(X, parameters, [5 10 1])
println(a2)
# println(caches)

# function cost_binary(Y, Ŷ)
#     @assert length(Y) == length(Ŷ)
#     m = length(Y)

#     cost = - sum(Y .* log.(Ŷ) .+ (1 .- Y) .* log.(1 .- Ŷ)) / m
# end

Y = [0.777085  0.248755  0.660713  0.37982  0.978839  0.182363  0.07204  0.855448  0.37393  0.40736]
@assert cost_binary(Y, Y) == 0.4982477807320199

# function sigmoid_back(x)
#     sig = sigmoid(x)
#     return sig*(1-sig)
# end

# function relu_back(x)
#     if x>0
#         return 1
#     else
#         return 0
#     end
# end

# function backward_prop(Y::Array{Float64}, Ŷ::Array{Float64}, parameters::Dict{String, Array{Float64}}, caches::Dict{String, Array{Float64}}, layer_dims::Array{Int}, activations=Nothing)

#     num_layers = length(layer_dims)
#     @assert length(Y) == length(Ŷ)
#     m = size(Y)[2]

#     if activations === Nothing
#         activations = Dict()
#         for i = 1:num_layers-1
#             activations[i] = relu_back
#         end
#         activations[num_layers] = sigmoid_back
#     else
#         activations_orig = activations
#         activations = []
#         for activation in activations_orig
#             push!(activations, @eval ($(Symbol("$activation", "_back"))))
#         end
#     end
#     # println(activations)

#     dA = sum(.- Y ./ Ŷ .+ (1 .- Y) ./ (1 .- Ŷ))

#     grads = Dict{String, Array{Float64}}()

#     for l in num_layers-1:-1:1
#         dZ = dA .* activations[l].(caches[string("Z", l)])
#         grads[string("dw", l)] = 1/m .* (dZ * transpose(caches[string("A", l-1)]))
#         grads[string("db", l)] = 1/m .* sum(dZ, dims=2)
#         dA = transpose(parameters[string("W", l)]) * dZ
#     end

#     return grads
# end

# grads = backward_prop(randn(1, 100), a2, parameters, caches, [5 10 1])
# printdict(grads)
# @assert size(grads["dw1"]) == size(parameters["W1"])
grads = backward_prop(randn(1, 100), a2, parameters, caches, [5 10 1], [relu, sigmoid])
printdict(grads)
@assert size(grads["dw1"]) == size(parameters["W1"])

# function update_parameters(parameters::Dict{String, Array{Float64}}, grads::Dict{String, Array{Float64}}, layer_dims::Array{Int}, learning_rate::Float64)
#     num_layers = length(layer_dims)
#     for l = 1:num_layers-1
#         parameters[string("W", l)] += learning_rate .* grads[string("dw", l)]
#         parameters[string("b", l)] += learning_rate .* grads[string("db", l)]
#     end
#     return parameters
# end

printdict_full(parameters)
update_parameters(parameters, grads, [5 10 1], 0.01)
printdict_full(parameters)

# function neural_network_dense(X, Y, layer_dims::Array{Int}, num_iterations::Int, learning_rate::Float64, activations=Nothing)
#     num_layers = length(layer_dims) # calculate number of layers
#
#     # if activations are not given, assume that all hidden layers have relu and output layer has sigmoid
#     if activations === Nothing
#         activations = Array{Function}(undef, num_layers-1)
#         for i = 1:num_layers-2
#             activations[i] = relu
#         end
#         activations[num_layers-1] = sigmoid
#     end
#     activations = Tuple(activations)
#     println(activations)
#     activations_back = []
#     for activation in activations
#         push!(activations_back, @eval ($(Symbol("$activation", "_back"))))
#     end
#     activations_back = Tuple(activations_back)
#     println(activations_back)
#
#     parameters = initialize_parameters(layer_dims)
#
#     for iteration = 1:num_iterations
#         Ŷ, caches = forward_prop(X, parameters, layer_dims, activations)
#         grads = backward_prop(Y, Ŷ, parameters, caches, layer_dims, activations_back)
#         parameters = update_parameters(parameters, grads, layer_dims, learning_rate)
#
#         if iteration % 1 == 0
#             cost = cost_binary(Y, Ŷ)
#             println("Cost at iteration $iteration is $cost")
#         end
#     end
#
#     return parameters
# end

X = randn(10, 100)
Y = randn(1, 100)
layer_dims = [10 10 1]
neural_network_dense(X, Y, layer_dims, 100, 0.05)


struct YeeGrid{T,nD}
	x::StructuredGrid{nD}
	y::StructuredGrid{nD}
	z::StructuredGrid{nD}
end

function YeeGrid{T}(refgrid::StructuredGrid{nD}) where {nD,T}
	YeeGrid{T,nD}(
		nodes(T,XComponent,refgrid),
		nodes(T,YComponent,refgrid),
		nodes(T,ZComponent,refgrid))
end
function YeeGrid{T}(nodes...) where T
	refgrid = StructuredGrid(nodes...)
	YeeGrid{T}(refgrid)
end

Base.size(yeegrid::YeeGrid) = (size(yeegrid.x),size(yeegrid.y),size(yeegrid.z))


# Try to approximate squared-exponential kernel:
using PyPlot
using Optim

function gaussian_approx5()

ntime = 1000
time = collect(linspace(0,10,ntime))
gaussian=exp(-time.^2/2.)

function compute_chi(x)
model = zeros(ntime)
c1 = 1.0
c2 = 0.0
for k=1:nexp
  model += x[1+(k-1)*2].*exp(-x[2+(k-1)*2].*time)
  c1 -= x[1+(k-1)*2]
  c2 -= x[1+(k-1)*2]*x[2+(k-1)*2]
end
c2 = c2/c1
model += c1.*exp(-c2.*time)
return sum((gaussian-model).^2)
end

nexp = 6
#nexp = 2
#nexp = 4
npar = 2*nexp
abest=zeros(npar)
abest=[2.7518074963750734,1.5312499407621147,-0.3546256365636662,4.753504292338284,-1.3968344706386124,1.7967431129812566,1.0,0.5,0.1,0.1,-0.2,0.5]

atrial = zeros(npar)
for k=1:npar
  atrial[k] = abest[k]
end
model = zeros(ntime)
c1 = 1.0
c2 = 0.0
for k=1:nexp-1
  model += atrial[1+(k-1)*2].*exp(-atrial[2+(k-1)*2].*time)
  c1 -= atrial[1+(k-1)*2]
  c2 -= atrial[1+(k-1)*2]*atrial[2+(k-1)*2]
end
c2 = c2/c1
model += c1.*exp(-c2.*time)
clf()
plot(time,gaussian)
plot(time,model)
#chibest=sum((gaussian-model).^2)
chibest=compute_chi(atrial)

println("Initial chi-square: ",chibest)
read(STDIN,Char)

result = optimize(compute_chi, atrial, BFGS(), OptimizationOptions(autodiff = true))

println(Optim.minimizer(result),Optim.minimum(result))
abest= Optim.minimizer(result)
clf()
plot(time,gaussian)
plot(time,model)
model = zeros(ntime)
c1 = 1.0
c2 = 0.0
for k=1:nexp
  model += abest[1+(k-1)*2].*exp(-abest[2+(k-1)*2].*time)
  c1 -= abest[1+(k-1)*2]
  c2 -= abest[1+(k-1)*2]*abest[2+(k-1)*2]
end
c2 = c2/c1
model += c1.*exp(-c2.*time)
plot(time,(gaussian-model))
chi=sum((gaussian-model).^2)
println("New minimum: ",chi," ",abest," ",c1," ",c2," ",std(gaussian-model)," ",maximum(abs(gaussian-model)))
read(STDIN,Char)

return
end


module TestEnsembleGibbs

using Test: @test, @testset
using Distributions
using Ensemble
using Ensemble.Stats: logsumexp
using Statistics

const snr_half = 3.0
const mu = log(3.0)
const sigma = 1.0/3.0
const lambda = 100.0

const ptrue = [log(lambda), mu, log(sigma)]

log_odds(x) = 0.5*(x*x - snr_half*snr_half)

function logpselect(x)
    lo = log_odds(x)

    return lo - logsumexp(0.0, lo)
end

function logpnselect(x)
    lo = log_odds(x)

    return -logsumexp(0.0, lo)
end

function draw(lambda=lambda, mu=mu, sigma=sigma)
    n = rand(Poisson(lambda))
    xs = exp.(mu .+ sigma.*randn(n))
    psel = exp.(logpselect.(xs))

    sel = rand(n) .< psel

    xs[sel], xs[.~(sel)]
end

function lognorm_logpdf(xs, mu, sigma)
    out = zeros(size(xs)...)

    log_inv_sqrt_2pi = -0.5*log(2.0*pi)
    log_sigma = log(sigma)
    sigma2 = sigma*sigma

    for i in eachindex(xs)
        log_x = log(xs[i])
        y = log_x - mu
        out[i] = log_inv_sqrt_2pi - log_sigma - log_x - 0.5*y*y/sigma2
    end

    out
end

function make_lnprob(xsdet)
    ndet = length(xsdet)
    function lnprob(params, xsndet)
        nndet = length(xsndet)
        log_lambda, mu, log_sigma = params
        lambda = exp(log_lambda)
        sigma = exp(log_sigma)

        sum(logpselect.(xsdet)) + sum(logpnselect.(xsndet)) + (ndet + nndet)*log_lambda - lambda + sum(lognorm_logpdf(xsdet, mu, sigma)) + sum(lognorm_logpdf(xsndet, mu, sigma))
    end
end

function gibbsupdate(params, xsndet)
    log_lambda, mu, log_sigma = params
    lambda = exp(log_lambda)
    sigma = exp(log_sigma)

    xsdet, xsndet = draw(lambda, mu, sigma)
    xsndet
end

function testall()
    xsdet, xsndet = draw()

    lnprob = make_lnprob(xsdet)
    ps = zeros(3, 100)
    for i in 1:100
        for j in 1:3
            ps[j,i] = ptrue[j] + 1e-3*randn()
        end
    end
    gs = EnsembleGibbs.gibbses(ps, [nothing for i in 1:100], gibbsupdate)
    lnprobs = EnsembleGibbs.lnprobs(ps, gs, lnprob)

    for i in 1:1000
        ps, gs, lnprobs = EnsembleGibbs.update(ps, gs, lnprobs, lnprob, gibbsupdate)
    end

    pts = zeros(3, 100, 100)
    for i in 1:100
        for j in 1:10
            ps, gs, lnprobs = EnsembleGibbs.update(ps, gs, lnprobs, lnprob, gibbsupdate)
        end
        pts[:,:,i] = ps
    end

    @testset "EnsembleGibbs tests" begin
        @test((mean(pts[1,:,:] .- ptrue[1])/std(pts[1,:,:]) < 3))
        @test((mean(pts[2,:,:] .- ptrue[2])/std(pts[2,:,:]) < 3))
        @test((mean(pts[3,:,:] .- ptrue[3])/std(pts[3,:,:]) < 3))
    end
end

end


@using LinearAlgebraicRepresentation as LAR
@using LinearAlgebra as LA
@using ViewerGL as GL
include("../../src/Lar.jl")
LarA = Lar.Arrangement

expl = 2.0

store = [];
scaling = 1.20;
V,(VV,EV,FV,CV) = LAR.cuboid([0.25,0.25,0.25],true,[-0.25,-0.25,-0.25]);
mybox = (V,CV,FV,EV);
for k=1:20
	size = rand()*scaling
	scale = LAR.s(size,size,size)
	transl = LAR.t(rand(3)...)
	alpha = 2*pi*rand()
	rx = LAR.r(alpha,0,0); ry = LAR.r(0,alpha,0); rz = LAR.r(0,0,alpha)
	rot = rx * ry * rz
	str = LAR.Struct([ transl, scale, rot, mybox ])
	obj = LAR.struct2lar(str)
	vs = obj[1]
	diag = LinearAlgebra.norm(vs[:,8]-vs[:,1])
	if diag > 1/5
		push!(store, obj)
	end
end

str = LAR.Struct(store);
V,CV,FV,EV = LAR.struct2lar(str);

GL.VIEW([ GL.GLPol(V,CV, GL.COLORS[1]) ]);

cop_EV = LAR.coboundary_0(EV::Lar.Cells);
cop_FE = LAR.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);

model = Lar.Model(V);
Lar.addModelCells!(model, 1, cop_EV)
Lar.addModelCells!(model, 2, cop_FE)

sp_idx = LarA.spaceIndex(model, 2)

de_models = [
	LarA.face_decomposition(model, face_idx, sp_idx[face_idx])
	for face_idx = 1 : size(model, 2, 1)
]

m = deepcopy(de_models[1])
for i = 2 : length(de_models)
	Lar.uniteModels!(m, de_models[i])
end

triangulated_faces = LAR.triangulate2D(convert(Lar.Points, m.G'), m.T[1:2])
FVs = convert(Array{Lar.Cells}, triangulated_faces);
GL.VIEW( GL.GLExplode(m.G,FVs,expl,expl,expl,99,1) );

#-------------------------------------------------------------------------------
#
# V, copEV, copFE, copCF = Lar.Arrangement.spatial_arrangement(
# 	 W::Lar.Points, cop_EV::Lar.ChainOp, cop_FE::Lar.ChainOp);
#
#
# triangulated_faces = Lar.triangulate2D(V, [copEV, copFE]);
# FVs = convert(Array{Lar.Cells}, triangulated_faces);
# V = convert(Lar.Points, V');
# GL.VIEW( GL.GLExplode(V,FVs,1.5,1.5,1.5,99,1) );
#
#
# EVs = Lar.FV2EVs(copEV, copFE); # polygonal face fragments
# GL.VIEW( GL.GLExplode(V,EVs,1.5,1.5,1.5,99,1) );


using PictureSegmentation, Images, TestImages
using Test

function generate_3d_img()
    img=Array{RGB{Float64},3}(undef,100,100,100)
    for i in CartesianIndices(img)
        if (i[1]-50)^2 + (i[2]-50)^2 < (i[3]/4)^2
            img[i]=RGB(i[3]/100,0,1-(i[3]/100))
        else
            img[i]=RGB(0,0,0)
        end
    end
    for i = 10:40
        for j = 10:40
            for k = 10:40
                img[i,j,k] = RGB(0.25,0.66,0.02)
            end
        end
    end
    return img
end

@testset "PictureSegmentation.jl" begin
    include("growcut.jl")
    include("growcut3d.jl")
    include("seedpixel.jl")
end


using Distributions
using SMC
using KernelDensity
using PyPlot
using LaTeXStrings

include("../math.jl")
include("../aide.jl")

# where to write plots to
const PLOT_DIR = "plots"
if !Base.Filesystem.isdir(PLOT_DIR)
    Base.Filesystem.mkdir(PLOT_DIR)
end

"""
An unnormalized target density defined by a Gaussian prior and a mixture-of-Gaussian likelihood.
"""
struct Problem
    prior_mean::Float64
    prior_std::Float64
    proportions::Array{Float64,1}
    likelihood_means::Array{Float64,1}
    likelihood_stds::Array{Float64,1}
end

function log_joint_density(problem::Problem, x::Float64)
    prior = Normal(problem.prior_mean, problem.prior_std)
    ld = logpdf(prior, x)
    log_likelihoods = Array{Float64,1}(length(problem.likelihood_means))
    for i=1:length(problem.likelihood_means)
        likelihood = Normal(problem.likelihood_means[i],
                            problem.likelihood_stds[i])
        log_likelihoods[i] = log(problem.proportions[i]) + logpdf(likelihood, x)
    end
    ld + logsumexp(log_likelihoods)
end

"""
Exact sampler for the given inference problem.
"""
struct ExactSampler <: Sampler
    problem::Problem
end

function simulate(s::ExactSampler)
    error("not implemented!")
end

function regenerate(s::ExactSampler, x)
    log_joint_density(s.problem, x)
end

"""
Sampling importance resampling scheme.

The distribution is a Gaussian distribution with mean `mean` and
standard-deviation `std`.  `SMC.forward` Samples from the proposal and returns
the importance weight with respect to the given inference problem `p`.
`SMC.backward` returns the importance weight given a latent sample.
"""
struct GaussianInitializer
    p::Problem
    mean::Float64
    std::Float64
end

function SMC.forward(init::GaussianInitializer)
    dist = Normal(init.mean, init.std)
    x = rand(dist)
    (x, log_joint_density(init.p, x) - logpdf(dist, x))
end

function SMC.backward(init::GaussianInitializer, x::Float64)
    dist = Normal(init.mean, init.std)
    log_joint_density(init.p, x) - logpdf(dist, x)
end

function make_sir_scheme(problem::Problem,
                   init_mean::Float64,
                   init_std::Float64,
                   num_particles::Int)
    initializer = GaussianInitializer(problem, init_mean, init_std)
    SMCScheme(initializer, [], num_particles)
end

function make_offset_proposal_scheme(problem::Problem, num_particles::Int)
    make_sir_scheme(problem, 2.0, 0.5, num_particles)
end

function make_broad_proposal_scheme(problem::Problem, num_particles::Int)
    make_sir_scheme(problem, 0.0, 1.0, num_particles)
end

function make_sampler(sir_scheme::SMCScheme, problem::Problem)
    SMCSampler(sir_scheme, (x) -> log_joint_density(problem, x))
end

const OFFSET_PROPOSAL = "offset"
const BROAD_PROPOSAL = "broad"

"""
Returns a dictionary mapping key to a vector of tuples (sample, log marginal
likelihood estimate)
"""
function run_sir(problem::Problem, num_particles_list::Vector{Int}, num_replicates::Int)
    sir_outputs = Dict()
    for n in num_particles_list
        key = (OFFSET_PROPOSAL, n)
        scheme = make_offset_proposal_scheme(problem, n)
        sir_outputs[key] = map((i) -> smc(scheme), 1:num_replicates)
        key = (BROAD_PROPOSAL, n)
        scheme = make_broad_proposal_scheme(problem, n)
        sir_outputs[key] = map((i) -> smc(scheme), 1:num_replicates)
    end
    return sir_outputs
end

function get_log_marginal_likelihood_estimates(sir_outputs::Dict, proposal_name::String, num_particles::Int)
    map((output) -> output[2],  sir_outputs[(proposal_name, num_particles)])
end

function get_samples(sir_outputs::Dict, proposal_name::String, num_particles::Int)
    map((output) -> output[1],  sir_outputs[(proposal_name, num_particles)])
end

"""
Returns a dictionary mapping key to kernel density estimators
"""
function kernel_density_estimators(sir_outputs::Dict)
    kd_estimators = Dict()
    for (key, output) in sir_outputs
        samples = map((o) -> o[1], output)
        kd_estimators[key] = InterpKDE(kde(samples))
    end
    return kd_estimators
end

function plot_sir_sampling_densities(problem::Problem, kd_estimators::Dict, fname::String)

    width = 2.5
    height = 2.5
    plt[:figure](figsize=(width * 1.3, height))
    ax = plt[:gca]()

    # number of particles to plot sampling densities for
    all_num_particles_to_plot_samples = [1, 10, 100, 1000]

    # location of sampling densities on the y-axis
    yticks = [4, 3, 2, 1]

    # location of floor of plot on the y-axis
    density_y_base = 0

    # height of the target density on the y-axis
    density_actual_span = 0.7

    # x-axis limits
    data_lim = (-3, 4)

    # x-coordinates at which to evaluate densities
    density_xs = linspace(-3, 4, 100)

    # plot the unnormalized target density
    density_ys = map((x) -> exp(log_joint_density(problem, x)), density_xs)
    density_ys_span = maximum(density_ys) - minimum(density_ys)
    density_ys *= density_actual_span / density_ys_span
    plt[:plot](density_xs, density_y_base + density_ys, color="black")

    # plot the sampling densities (kernel density estimates)
    for (i, num_particles) in enumerate(all_num_particles_to_plot_samples)
        ys = pdf(kd_estimators[(OFFSET_PROPOSAL, num_particles)], density_xs)
        ys = (ys / maximum(ys)) * density_actual_span + yticks[i]
        plt[:plot](density_xs, ys, color="magenta", zorder=100, alpha=0.7)
        ys = pdf(kd_estimators[(BROAD_PROPOSAL, num_particles)], density_xs)
        ys = (ys / maximum(ys)) * density_actual_span + yticks[i]
        plt[:plot](density_xs, ys, color="blue", alpha=0.7)
    end

    # labels
    plt[:text](1.5, density_y_base + density_actual_span/9.0, "posterior\ndensity", fontsize=10)
    plt[:text](-1.7, density_y_base + density_actual_span/6.0, "L", fontsize=12)
    plt[:text](0.3, density_y_base + density_actual_span/6.0, "R", fontsize=12)
    ax[:set_xlim](data_lim)
    exponents = [0, 1, 2, 3]
    exponent_strings = map((i) -> @sprintf("%d", i), exponents)
    ax[:set_yticks](yticks)
    ax[:set_yticklabels](map((exponent_string) -> latexstring("\$10^{$exponent_string}\$"),
                             exponent_strings))
    ax[:tick_params](axis="both", which="major", labelsize=9)
    plt[:ylabel]("Number of particles\n")

    # arrow for the number of particles on y-axis
    ax[:annotate]("", xy=(-0.2, 0.0), xycoords="axes fraction", xytext=(-0.2, 1.0),
                arrowprops=Dict([("arrowstyle", "->"), ("color", "black")]))

    plt[:tight_layout]()
    plt[:savefig](fname)
end

function draw_arrow(x, y, xlen, ylen, ax)
    ax[:arrow](x, y, xlen, ylen, fc="black", ec="black",
	    transform=ax[:transAxes], length_includes_head=true,
	    head_width=0.05, head_length=0.05)
end

function plot_log_marginal_likelihoods(sir_outputs::Dict, all_num_particles::Vector{Int}, fname::String)

    function plot_single_series(proposal_name::String, label, color)
        # a vector of a vector of IID estimates
        # there is one vector of IID estimates for each number of particles
        individual_estimates = map((num_particles) -> get_log_marginal_likelihood_estimates(sir_outputs, proposal_name, num_particles), all_num_particles)
        estimates = map(mean, individual_estimates)
        stderrs = map((data::Vector{Float64}) -> std(data)/sqrt(length(data)), individual_estimates)
        plt[:plot](all_num_particles, estimates, color=color, label=label)
        lower = estimates - stderrs
        upper = estimates + stderrs
        plt[:fill_between](all_num_particles, lower, estimates, color=color, alpha=0.3)
        plt[:fill_between](all_num_particles, estimates, upper, color=color, alpha=0.3)
    end

    width = 2.5
    height = 2.5
    plt[:figure](figsize=(width, height))
    ax = plt[:gca]()

    # plot log marginal likelihood estimates
    plot_single_series(OFFSET_PROPOSAL, "Offset proposal", "magenta")
    plot_single_series(BROAD_PROPOSAL, "Broad proposal", "blue")
    gold_standard_lml = mean(get_log_marginal_likelihood_estimates(sir_outputs, "broad", maximum(all_num_particles)))
    xlim = ax[:get_xlim]()
    plt[:plot](xlim, gold_standard_lml * ones(2), color="orange", label="Gold-standard")

    # arrow and label for small penalty
    plt[:text](12, -26, "Small penalty")
    draw_arrow(0.8, 0.73, 0.0, 0.19, ax)

    # plot formatting and labels
    ax[:set_xscale]("log")
    plt[:minorticks_off]()
    ax[:tick_params](axis="both", which="major", labelsize=8)
    ax[:set_xlim](xlim)
    plt[:legend](loc="lower right", fontsize=8.7)
    plt[:xlabel]("Number of particles")
    plt[:title]("Log marginal\nlikelihood (nats)")
    plt[:tight_layout]()
    plt[:savefig](fname)
end

function aide_estimate(gold_standard_sampler::Sampler, target_sampler::Sampler, num_replicates::Int)
    lowers = map((i::Int) -> estimate_elbo(target_sampler, gold_standard_sampler, 1, 1),
                 1:num_replicates)
    uppers = map((i::Int) -> -estimate_elbo(gold_standard_sampler, target_sampler, 1, 1),
                 1:num_replicates)
    return (uppers - lowers)
end

function generate_aide_estimates(problem::Problem, all_num_particles::Vector{Int}, num_replicates::Int)
    gold_standard_scheme = make_broad_proposal_scheme(problem, maximum(all_num_particles))
    gold_standard_sampler = make_sampler(gold_standard_scheme, problem)
    aide_estimates = Dict()
    for num_particles in all_num_particles
        offset_proposal_scheme = make_offset_proposal_scheme(problem, num_particles)
        offset_proposal_sampler = make_sampler(offset_proposal_scheme, problem)
        broad_proposal_scheme = make_broad_proposal_scheme(problem, num_particles)
        broad_proposal_sampler = make_sampler(broad_proposal_scheme, problem)
        aide_estimates[(OFFSET_PROPOSAL, num_particles)] = aide_estimate(gold_standard_sampler, offset_proposal_sampler, num_replicates)
        aide_estimates[(BROAD_PROPOSAL, num_particles)] = aide_estimate(gold_standard_sampler, broad_proposal_sampler, num_replicates)
    end
    return aide_estimates
end

function plot_aide_estimates(aide_estimates::Dict, all_num_particles::Vector{Int}, fname::String)

    function plot_single_series(proposal_name::String, label, color)
        # a vector of a vector of IID estimates
        # there is one vector of IID estimates for each number of particles
        individual_estimates = map((num_particles) -> aide_estimates[(proposal_name, num_particles)],
                                   all_num_particles)
        estimates = map(mean, individual_estimates)
        stderrs = map((data::Vector{Float64}) -> std(data)/sqrt(length(data)), individual_estimates)
        plt[:plot](all_num_particles, estimates, color=color, label=label)
        lower = estimates - stderrs
        upper = estimates + stderrs
        plt[:fill_between](all_num_particles, lower, estimates, color=color, alpha=0.3)
        plt[:fill_between](all_num_particles, estimates, upper, color=color, alpha=0.3)
    end

    width = 2.5
    height = 2.5
    plt[:figure](figsize=(width, height))
    ax = plt[:gca]()

    # plot log marginal likelihood estimates
    plot_single_series(OFFSET_PROPOSAL, "Offset proposal", "magenta")
    plot_single_series(BROAD_PROPOSAL, "Broad proposal", "blue")

    # arrow and label for large penalty
    plt[:text](10, 29, "Large penalty for\nmissing mode")
    draw_arrow(0.8, 0.36, 0.0, -0.22, ax)

    # plot formatting and labels
    ymax = ax[:get_ylim]()[2]
    ax[:set_ylim](0, ymax)
    ax[:set_xscale]("log")
    plt[:minorticks_off]()
    ax[:tick_params](axis="both", which="major", labelsize=8)
    plt[:legend](loc="upper right", fontsize=8.7)
    plt[:xlabel]("Number of particles")
    plt[:title]("AIDE estimate\n(nats)")
    plt[:tight_layout]()
    plt[:savefig](fname)
end

# generate plots
srand(1)
problem = Problem(0.0, 10.0, [0.7, 0.3], [-1, 1], [0.1, 0.1])
exponents = linspace(0, 3, 10)
num_particles_list = Vector{Int}(floor.(10.^exponents))

# generate subfigure (a)
println("Generating density plot...")
sir_outputs_for_kde = run_sir(problem, num_particles_list, 10000)
kd_estimators = kernel_density_estimators(sir_outputs_for_kde)
plot_sir_sampling_densities(problem, kd_estimators, "$PLOT_DIR/bimodal_densities.pdf")

# generate subfigure (b)
println("Generating AIDE estimate plot...")
aide_estimates = generate_aide_estimates(problem, num_particles_list, 100)
plot_aide_estimates(aide_estimates, num_particles_list, "$PLOT_DIR/bimodal_symkl.pdf")

# generate subfigure (c)
println("Generating log marginal likelihood plot...")
sir_outputs_for_lml = run_sir(problem, num_particles_list, 100)
plot_log_marginal_likelihoods(sir_outputs_for_lml, num_particles_list, "$PLOT_DIR/bimodal_lml.pdf")


export Poisson, poisson_aux
export Gauss, gauss_aux
export ScaledGauss, scaled_gauss_aux

"""
    poisson_aux(μ, meas)

Calculates the Poisson loss for `μ` and `meas`.
`μ` can be of larger size than `meas`. In that case
we extract a centered region from `μ` of the same size as `meas`.
"""
function poisson_aux(μ, meas)
    μ = center_extract(μ, size(meas))

    μ[μ .<= 1f-8] .= 1f-8
    return sum(μ .- meas .* log.(μ))
end

 # define custom gradient for speed-up
 # ChainRulesCore offers the possibility to define a backward AD rule
 # which can be used by several different AD systems
function ChainRulesCore.rrule(::typeof(poisson_aux), μ, meas)
    Y = poisson_aux(μ, meas)

    function poisson_aux_pullback(xbar)
        meas_new = copy(μ)
        meas_new = center_set!(meas_new, meas)

        ∇ = xbar .* (one(eltype(μ)) .- meas_new ./ μ)
        ∇[μ .< 1f-8] .= 0
        return zero(eltype(μ)), ∇, zero(eltype(μ))
    end

    return Y, poisson_aux_pullback
end

"""
    Poisson()

Returns a function to calculate Poisson loss
Check the help of `poisson_aux`.
"""
function Poisson()
      return poisson_aux
end

"""
    gauss_aux(μ, meas)

Calculates the Gauss loss for `μ` and `meas`.
`μ` can be of larger size than `meas`. In that case
we extract a centered region from `μ` of the same size as `meas`.
"""
function gauss_aux(μ, meas)
    μ = center_extract(μ, size(meas))
    return sum(abs2.(μ - meas))
end

 # define custom gradient for speed-up
function ChainRulesCore.rrule(::typeof(gauss_aux), μ, meas)
    Y = gauss_aux(μ, meas)
    function gauss_aux_pullback(xbar)
        meas_new = copy(μ)
        meas_new = center_set!(meas_new, meas)
        return zero(eltype(μ)), 2 .* (μ - meas_new), zero(eltype(μ))
    end
    return Y, gauss_aux_pullback
end

"""
    Gauss()

Returns a function to calculate Gauss loss.
Check the help of `gauss_aux`.
"""
function Gauss()
    return gauss_aux
end

"""
    scaled_gauss_aux(μ, meas)
Calculates the scaled Gauss loss for `μ` and `meas`.
`μ` can be of larger size than `meas`. In that case
we extract a centered region from `μ` of the same size as `meas`.
"""
function scaled_gauss_aux(μ, meas)
    μ = center_extract(μ, size(meas))
    μ[μ .<= 1f-8] .= 1f-8
    return sum(1/2f0 .* log.(μ) .+ (meas .- μ).^2 ./ (2 .* μ))
end

 # define custom gradient for speed-up
function ChainRulesCore.rrule(::typeof(scaled_gauss_aux), μ, meas)
    Y = scaled_gauss_aux(μ, meas)
    function scaled_gauss_aux_pullback(xbar)
        meas_new = copy(μ)
        meas_new = center_set!(meas_new, meas)
        ∇ = (μ + μ.^2 - meas_new.^2)./(2 .* μ.^2)
        ∇[μ .<= 1f-8] .= 0
        return zero(eltype(μ)), ∇, zero(eltype(μ))
    end
    return Y, scaled_gauss_aux_pullback
end

"""
    ScaledGauss()

Returns a function to calculate scaled Gauss loss.
Check the help of `scaled_gauss_aux`.
"""
function ScaledGauss()
    return scaled_gauss_aux
end


function sample!(fin_el::Hexahedron, s::Matrix{Float64} , wt::Vector{Float64})
  #
  # This subroutine returns the local coordinates and weighting coefficients
  # of the integrating points.
  #

  root3 = 1.0 / sqrt(3.0)
  r15 = 0.2 * sqrt(15.0)
  nip = size(s,1)

  w = [5.0/9.0, 8.0/9.0, 5.0/9.0]
  v = [5.0/9.0*w; 8.0/9.0*w; 5.0/9.0*w]

  if nip == 1
    s[1,1:3] .= 0.0
    wt[1] = 8.0
  elseif nip == 8
    s[1,1] =  root3
    s[1,2] =  root3
    s[1,3] =  root3
    s[2,1] =  root3
    s[2,2] =  root3
    s[2,3] = -root3
    s[3,1] =  root3
    s[3,2] = -root3
    s[3,3] =  root3
    s[4,1] =  root3
    s[4,2] = -root3
    s[4,3] = -root3
    s[5,1] = -root3
    s[5,2] =  root3
    s[5,3] =  root3
    s[6,1] = -root3
    s[6,2] = -root3
    s[6,3] =  root3
    s[7,1] = -root3
    s[7,2] =  root3
    s[7,3] = -root3
    s[8,1] = -root3
    s[8,2] = -root3
    s[8,3] = -root3
    fill!(wt, 1)
  elseif nip == 14
    b = 0.795822426
    c = 0.758786911
    wt[1:6] = 0.886426593
    wt[7:14] = 0.335180055
    s[1,1] = -b
    s[2,1] = b
    s[3,2] = -b
    s[4,2] = b
    s[5,3] = -b
    s[6,3] = b
    s[7:end,:] = c
    s[7,1] = -c
    s[7,2] = -c
    s[7,3] = -c
    s[8,2] = -c
    s[8,3] = -c
    s[9,1] = -c
    s[9,3] = -c
    s[10,3] = -c
    s[11,1] = -c
    s[11,2] = -c
    s[12,2] = -c
    s[13,1] = -c
  elseif nip == 15
    b        = 1.0
    c        = 0.674199862
    wt[1]    = 1.564444444
    wt[2:7]  = 0.355555556
    wt[8:15] = 0.537777778
    s[2,1] = -b
    s[3,1] = b
    s[4,2] = -b
    s[5,2] = b
    s[6,3] = -b
    s[7,3] = b
    s[8:end,:] = c
    s[8,1] = -c
    s[8,2] = -c
    s[8,3] = -c
    s[9,2] = -c
    s[9,3] = -c
    s[10,1] = -c
    s[10,3] = -c
    s[11,3] = -c
    s[12,1] = -c
    s[12,2] = -c
    s[13,2] = -c
    s[14,1] = -c
  elseif nip == 27
    wt[:] = [5.0/9.0*v; 8.0/9.0*v; 5.0/9.0*v]
    s[1:7:3,1] = -r15
    s[2:8:3,1] = 0.0
    s[3:9:3,1] = r15
    s[1:3,3] = r15
    s[4:6,3] = 0.0
    s[7:9,3] = -r15
    s[1:9,2] = -r15
    s[10:16:3,1] = -r15
    s[11:17:3,1] = 0.0
    s[12:18:3,1] = r15
    s[10:12,3] = r15
    s[13:15,3] = 0.0
    s[16:18,3] = -r15
    s[10:18,2] = 0.0
    s[19:25:3,1] = -r15
    s[20:26:3,1] = 0.0
    s[21:27:3,1] = r15
    s[19:21,3] =  r15
    s[22:24,3] = 0.0
    s[25:27,3] = -r15
    s[19:27,2] =  r15
  else
    println("Wrong number of integrating points for a hexaahedron.")
  end
end


"""

# pcor_test

$(SIGNATURES)

Computes the partial correlation between two variables given a set of other
variables.

### Method
```julia
pcor_test(;
* `u::Vector{Symbol}`                  : Variables used to compute correlation
* `q::Int`                             : Number of variables in conditioning set
* `n::Int`                             : Number of observations
* `S::Matrix`                          : Sample covariance matrix
)
```
where:

  u[1], u[2]: Variables used to compute correlation between, remaining variables
  are the conditioning set

### Returns
```julia
* `res::Float64`                       : Correlation between u[1] and u[2]
```
# Extended help

### Acknowledgements

Original author:                       Giovanni M. Marchetti

Translated to Julia:                   Rob J Goedman

### License

The R package ggm is licensed under License: GPL-2.

The Julia translation is licenced under: MIT.

Part of the api, not exported.
"""
function pcor_test(d::DAG, u::SymbolList, q, n)

  r = pcor(d, u)
  df = n - 2 - q
  tval = r * sqrt(df) / sqrt(1 - r*r)
  pv = 2 * cdf(TDist(df), -abs(tval))
  (pv=tval, df=df, pvalue=pv)

end

export
  pcor_test


include("lib.jl")

function simulate(d, n = 10_000, n_sims = 20_000)
    x = Array(Float64, n)
    y = Array(Float64, n)

    methods = Dict(
        "unadjusted" => unadjusted,
        "trim_shared_threshold" => trim_shared_threshold,
        "trim_separate_thresholds" => trim_separate_thresholds,
        "trim_threshold_from_A" => trim_threshold_from_A,
        "winsorize_shared_threshold" => winsorize_shared_threshold,
        "winsorize_separate_thresholds" => winsorize_separate_thresholds,
        "winsorize_threshold_from_A" => winsorize_threshold_from_A,
    )

    header = (
        "n",
        "method",
        "w",
        "te",
        "se",
        "delta",
        "sigma",
        "pvalue",
        "lower",
        "upper",
    )
    println(join(header, "\t"))

    for te in (0.00, 0.05, 0.10)
        se = sqrt(var(d) / n + var(d) / n)
        for i in 1:n_sims
            gen_data!(x, y, d, te)
            for w in 0.60:0.01:0.99
                τ_x = threshold(x, w)
                τ_y = threshold(y, w)
                τ_xy = threshold(vcat(x, y), w)
                for (method, method_func) in methods
                    x′, y′ = method_func(x, y, τ_x, τ_y, τ_xy)
                    δ, σ, p, lower, upper = analyze(x′, y′)
                    row = (n, method, 1 - w, te, se, δ, σ, p, lower, upper)
                    println(join(row, "\t"))
                end
            end
        end
    end
end

d = LogNormal(0, 1)
simulate(d)


export Interval

struct Interval
    a::Float64
    b::Float64
    Interval(a, b) = a > b ? error("invalid input") : new(a, b)
    Interval(a::Real) = new(a, a)
    Interval(a::Char) = a == '\U2205' ? new(Inf, -Inf) : error("invalid input") # U+2205 is ∅
end

Base.minimum(A::Interval) = A.a
Base.maximum(A::Interval) = A.b

function Base.:(∈)(a::Real, A::Interval)
    if a >= A.a && a <= A.b
        return true
    else
        return false
    end
end

function Base.isempty(A::Interval)
    if A.a > A.b
        return true
    else
        return false
    end
end

function Base.issubset(A::Interval, B::Interval)
    if A.a >= B.a && A.b <= B.b
return true
    else
        return false
    end
end

function Base.intersect(A::Interval, B::Interval)
    if maximum(A) >= minimum(B) && maximum(B) >= minimum(A)
        return Interval(max(minimum(A), minimum(B)), min(maximum(A), maximum(B)))
    else
        return Interval('∅')
    end
end

function Base.show(io::IO, A::Interval)
    if isempty(A)
        print(io, '∅')
    else
        print(io, '\u301a')
        print(io, A.a)
        print(io, ", ")
        print(io, A.b)
        print(io, '\u301b')
    end
end


#=
    Ninjas, Bears, hunters

    Port of the PSI model
    https://github.com/eth-sri/psi/blob/master/test/ninjasBearsHunters.psi

    Similar to Rock, Paper, Scissors but there are some differences:
    - At start there are certain amount of ninjas, bears, and hunters. It
      might be different for each type.
    - In each turn, two of separare types are selected and
      one of them is eliminated according to the elimination rules
      (see below).
    - We continue until there is just one left.
    - The elimination rules are
        * bear kills ninja
        * hunter kills bear
        * ninja kills hunter

   What are the probabilities that the winner is a specific type?

   The PSI model starts with
      10 ninjas, 2 bears, and 5 hunters
   and the the exact probabilities of survival are:
    (probNinjaSurvive,probBearsSurvive,probHuntersSurvive)=
    (0.221480658306459,0.767229372971558,0.0112899687219824)

   It is - at first sight - a little surprising that the probability
   of the ninja as the winner is so low (0.22). But since there are so many
   ninjas, the hunters are killed quite soon which mean that there is no hunter
   left that can kill the bears (which kills the ninjas). Thus the bears have the greatest
   probability of survival (0.77).

   This Turing model:
   Summary Statistics
            parameters      mean       std   naive_se      mcse         ess       rhat   ess_per_sec
               Symbol   Float64   Float64    Float64   Float64    Float64?   Float64?       Float64

      probNinjaSurvive    0.2240    0.4171     0.0132    0.0167    951.6074     0.9992     5228.6121
      probBearsSurvive    0.7590    0.4279     0.0135    0.0151    963.7017     0.9990     5295.0643
    probHuntersSurvive    0.0170    0.1293     0.0041    0.0047   1037.0165     1.0005     5697.8929

   Quantiles
            parameters      2.5%     25.0%     50.0%     75.0%     97.5%
               Symbol   Float64   Float64   Float64   Float64   Float64

      probNinjaSurvive    0.0000    0.0000    0.0000    0.0000    1.0000
      probBearsSurvive    0.0000    1.0000    1.0000    1.0000    1.0000
    probHuntersSurvive    0.0000    0.0000    0.0000    0.0000    0.0000

    Some other examples:
    * start with 3 ninjas, 4 bears, and 3 hunters.
      The probability of survival is
      probNinjaSurvive  : 0.2520
      probBearsSurvive  : 0.2880
      probHuntersSurvive: 0.4600
      (Exact probability from PSI:
       3927004/15690675,2846092/9984975,3402579/7322315 = 0.250276294678209,0.285037468796867,0.464686236524924
       )

    * start with 1 ninjas, 2 bears, and 3 hunters.
      The probability of survival is
      probNinjaSurvive  : 0.4420
      probBearsSurvive  : 0.0810
      probHuntersSurvive: 0.4770

      (Exact probability from PSI:
       339/770,111/1540,751/1540 = 0.44025974025974,0.0720779220779221,0.487662337662338
      )

    Remember;
    * bear kills ninja
    * hunter kills bear
    * ninja kills hunter

    Also see
    * https://en.wikipedia.org/wiki/Rock_paper_scissors
    * It seems to have been popularized by the FedEx commercial: Bear, hunter ninja.
      https://www.youtube.com/watch?v=AhFbbq0zpHY
    * (The PSI model refers to a link that's not accessible.)

=#
using Turing, StatsPlots, DataFrames, Distributions, Random
include("jl_utils.jl")

@model function ninjas_bears_hunters(ninjas=10, bears=2, hunters=5)
    n = ninjas + bears + hunters
    NINJA = 1
    BEAR = 2
    HUNTER = 3
    p1 = tzeros(1)
    p2 = tzeros(1)
    i = 0
    while true
      i += 1

      s = ninjas + bears + hunters
      v = [ninjas,bears,hunters]
      # println("i:$i v:$v")
      if s <= 1 || count(t -> t > 0,v) == 1
         break
      end

      a = [ninjas/s,bears/s,hunters/s]

      push!(p1,0)
      push!(p2,0)
      p1[i] ~ Categorical(simplex(a))
      p2[i] ~ Categorical(simplex(a))
      # true ~ Dirac(p1[i] != p2[i])

      # there must be at least two remaining civilizations.
      if p1[i] != p2[i] && s > 1 && ninjas + bears > 0 && ninjas + hunters > 0 && bears + hunters > 0

         # true ~ Dirac(!(ninjas  == 0 && (p1[i] == NINJA  || p2[i] == NINJA)))
         # true ~ Dirac(!(bears	  == 0 && (p1[i] == BEAR   || p2[i] == BEAR)))
         # true ~ Dirac(!(hunters == 0 && (p1[i] == HUNTER || p2[i] == HUNTER)))

         if		((p1[i] == NINJA && p2[i] == BEAR)   || (p1[i] == BEAR   && p2[i] == NINJA) && ninjas  > 0) ninjas  -=1 # Bear kills ninja
         elseif ((p1[i] == NINJA && p2[i] == HUNTER) || (p1[i] == HUNTER && p2[i] == NINJA) && hunters > 0) hunters -=1 # Ninja kills hunter
         elseif ((p1[i] == BEAR  && p2[i] == HUNTER) || (p1[i] == HUNTER && p2[i] == BEAR)  && bears   > 0) bears   -=1 # Hunter kills bear
         end
      end
   end

   # true ~ Dirac(!(ninjas == 0 && bears == 0 && hunters == 0))
   # true ~ Dirac(ninjas + bears == 0 || ninjas + hunters == 0 || bears + hunters == 0)
   probNinjaSurvive   ~ Dirac(ninjas > 0)
   probBearsSurvive   ~ Dirac(bears > 0)
   probHuntersSurvive ~ Dirac(hunters > 0)
   return
end

num_ninjas=10
num_bears=2
num_hunters=5
# num_ninjas=3
# num_bears=4
# num_hunters=3

# num_ninjas=1
# num_bears=2
# num_hunters=3
model = ninjas_bears_hunters(num_ninjas,num_bears,num_hunters)

# Since we don't do any observations, we can use Prior()
chns = sample(model, Prior(), 1_000)
# chns = sample(model, MH(), 1_000)
# chns = sample(model, PG(5), 1_000)
# chns = sample(model, SMC(), 1_000)
# chns = sample(model, IS(), 1_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns[[:probNinjaSurvive,:probBearsSurvive,:probHuntersSurvive]])

# show_var_dist_pct(chns,:probNinjaSurvive)
# show_var_dist_pct(chns,:probBearsSurvive)
# show_var_dist_pct(chns,:probHuntersSurvive)


using Decimals
using Compat.Test

@testset "Rounding" begin

@test round(Decimal(7.123456), digits=0) == Decimal(7)
@test round(Decimal(7.123456), digits=2) == Decimal(7.12)
@test round(Decimal(7.123456), digits=3) == Decimal(7.123)
@test round(Decimal(7.123456), digits=5) == Decimal(7.12346)
@test round(Decimal(7.123456), digits=6) == Decimal(7.123456)

@test round.(Decimal.([0.1111, 0.2222, 0.8888]), digits=2) == Decimal.([0.11, 0.22, 0.89])

function tet()
    a = parse(Decimal, "1.0000001")
    for i = 1:27
        a *= a
    end
    return a
end

# set DIGITS = 20 (aaljuffali's example)
@test tet() == Decimal(0, 67453047074102193157641340, -20)

end


module Potentials

export PowerLawPotential, mass, density, potential

include("constants.jl")
using .Constants
c = Constants

abstract Potential
abstract SphericalPotential <: Potential

function density(phi::SphericalPotential, r::Real)
    error("Must use subtype of SphericalPotential with defined density!")
end

function mass(phi::SphericalPotential, r::Real)
    error("Numeric integration of density profile not implemented yet.")
end

function potential(phi::SphericalPotential, r::Real)
    M = mass(phi, r)
    return -c.G * M ./ r
end

immutable PowerLawPotential <: SphericalPotential
    #=
    r0 : scale radius
    rho0 : scale density
    alpha : power law index of density
    =#
    r0::Real
    rho0::Real
    alpha::Real
end

function mass(phi::PowerLawPotential, r::Real)
    r0 = phi.r0
    rho0 = phi.rho0
    alpha = phi.alpha
    4pi * rho0 * r0 ^ 3 / (3 + alpha) * (r / r0) .^ (3 + alpha)
end

function density(phi::PowerLawPotential, r::Real)
    r0 = phi.r0
    rho0 = phi.rho0
    alpha = phi.alpha
    rho0 * (r / r0) .^ alpha
end

end


# -----------------------------------------------------------------------
# Radiative Forcing From Methane (based on Myhre et al. equations).
# -----------------------------------------------------------------------

# Create function to account for overlap in CH₄ and N₂O absoprtion bands.
function overlap(M, N)
    d = 1.0 + 2.01e-5*(M*N)^0.75 + 5.31e-15*M*(M*N)^1.52
    return 0.47 * log(d)
end

@defcomp rf_ch4_myhre_fair begin

    N₂O_0           = Parameter()             # Initial (pre-industrial) nitrous oxide concentration (ppb).
    CH₄_0           = Parameter()             # Initial (pre-industrial) methane concentration (ppb).
    scale_CH₄       = Parameter()             # Scaling factor for uncertainty in methane radiative forcing.
    H₂O_share       = Parameter()             # Share of direct methane forcing used to estimate stratospheric water vapor forcing due to methane oxidation.
    CH₄             = Parameter(index=[time]) # Atmospheric methane concentration (ppb).

    forcing_CH₄     = Variable(index=[time])  # Radiative forcing from methane (Wm⁻²).
    forcing_CH₄_H₂O = Variable(index=[time])  # Radiative forcing from stratospheric water vapor due to methane oxidiation (Wm⁻²).

    function run_timestep(p, v, d, t)

        if is_first(t)
            # Set initial radiative forcing to 0.
            v.forcing_CH₄[t] = 0.0
            v.forcing_CH₄_H₂O[t] = 0.0

        else

            # Calculate direct CH₄ radiative forcing.
            direct_rf = 0.036 * (sqrt(p.CH₄[t]) - sqrt(p.CH₄_0)) - (overlap(p.CH₄[t], p.N₂O_0) - overlap(p.CH₄_0, p.N₂O_0))

            # Calculate direct radiative forcing from CH4, accounting for forcing uncertainty.
            v.forcing_CH₄[t] = direct_rf * p.scale_CH₄

            #Calculate indirect CH4 forcing due to production of stratospheric H20 from CH4 oxidation.
            v.forcing_CH₄_H₂O[t] = direct_rf * p.H₂O_share
        end
    end
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -6.019960436071491 <= q <= 11.696817216382863)

                     add_NL_constraint(m, :(1/(1 + exp(-(0.9042935565441375 + 0.31236426273908524*1/(1 + exp(-(0.5868252904954097 + 0.04712499510197299*$(x[1]) + 0.4005456947892885*$(x[2]) + 0.34544434364772814*$(x[3]) + 0.48957180414623247*$(x[4])))) + -0.4128166127064836*1/(1 + exp(-(0.5797949667252156 + -0.7245574143823843*$(x[1]) + -0.08207866517990237*$(x[2]) + 0.4631961331947827*$(x[3]) + 0.9179781125579716*$(x[4])))) + -0.2476169775927386*1/(1 + exp(-(0.8773151656273224 + 0.003673243009740812*$(x[1]) + 0.698013040735662*$(x[2]) + 0.6787695040522581*$(x[3]) + -0.9128403322093082*$(x[4]))))))) + 1/(1 + exp(-(0.8600467319105785 + -0.4006275013725942*1/(1 + exp(-(0.5868252904954097 + 0.04712499510197299*$(x[1]) + 0.4005456947892885*$(x[2]) + 0.34544434364772814*$(x[3]) + 0.48957180414623247*$(x[4])))) + 0.22106887260345687*1/(1 + exp(-(0.5797949667252156 + -0.7245574143823843*$(x[1]) + -0.08207866517990237*$(x[2]) + 0.4631961331947827*$(x[3]) + 0.9179781125579716*$(x[4])))) + 0.7210213948433983*1/(1 + exp(-(0.8773151656273224 + 0.003673243009740812*$(x[1]) + 0.698013040735662*$(x[2]) + 0.6787695040522581*$(x[3]) + -0.9128403322093082*$(x[4]))))))) + 1/(1 + exp(-(0.982756228863964 + 0.34792310275142935*1/(1 + exp(-(0.5868252904954097 + 0.04712499510197299*$(x[1]) + 0.4005456947892885*$(x[2]) + 0.34544434364772814*$(x[3]) + 0.48957180414623247*$(x[4])))) + 0.7924515917531365*1/(1 + exp(-(0.5797949667252156 + -0.7245574143823843*$(x[1]) + -0.08207866517990237*$(x[2]) + 0.4631961331947827*$(x[3]) + 0.9179781125579716*$(x[4])))) + -0.9399735455074207*1/(1 + exp(-(0.8773151656273224 + 0.003673243009740812*$(x[1]) + 0.698013040735662*$(x[2]) + 0.6787695040522581*$(x[3]) + -0.9128403322093082*$(x[4]))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


@polly function kernel_jacobi_1d(tsteps, A, B)
    n = size(A,1)

    for t = 1:tsteps
        for i = 2:(n-1)
            B[i] = 0.33333 * (A[i-1] + A[i] + A[i + 1])
        end
        for i = 2:(n-1)
            A[i] = 0.33333 * (B[i-1] + B[i] + B[i + 1])
        end
    end
end

let
    tsteps = 500
    n = 2000

    A = zeros(Float32, n)
    B = zeros(Float32, n)

    for i = 1:n
        A[i] = (i+ 2) / n
        B[i] = (i+ 3) / n
    end

    SUITE["jacobi-1d"] = @benchmarkable kernel_jacobi_1d($tsteps, A, B) setup = (A = copy($A); B = copy($B))
end


using Printf

# --------------------------
# -- function setup       --
# --------------------------
#
# 基本量と保存量の設定
#
function setup(nx,dx,gamma)
    u   = zeros(nx)               # 速度
    rho = zeros(nx)               # 密度
    p   = zeros(nx)               # 圧力
    e   = zeros(nx)               # エネルギー
    x   = zeros(nx)               # 位置

    qf = zeros(nx,3)              # 基本量
    Qc = zeros(nx,3)              # 保存量

    """
    e=p/(r-1)+rho*u^2/2
    """
    for i in 1:nx
        u[i] = 0.0

        if i <= nx*0.5
            rho[i] = 1.0
            p[i] = 1.0
        elseif i >= nx*0.5
            rho[i] = 0.125
            p[i] = 0.1
        end
        e[i] = p[i]/(gamma-1)+rho[i]*(u[i]^2)/2
        x[i] = i*dx-dx/2
    end

    for i in 1:nx
        qf[i,1] = u[i]
        Qc[i,1] = rho[i]

        qf[i,2] = rho[i]
        Qc[i,2] = u[i]*rho[i]

        qf[i,3] = p[i]
        Qc[i,3] = e[i]

    end
    return x,qf,Qc
end

# --------------------------
# -- function boundary    --
# --------------------------
#
# 境界条件の設定
#
function boundary(nx,gamma,qf,Qc,bdcon)
    if Int(bdcon[1][1]) == 0
        for i in 1:3
            Qc[nx,i] = Qc[nx-1,i]
        end
    elseif Int(bdcon[1][1]) == 1
        rho = bdcon[1][2]
        u = bdcon[1][3]
        p = bdcon[1][4]
        e = p/(gamma-1)+rho*(u^2)/2
        Qc[nx,1] = rho
        Qc[nx,2] = rho*u
        Qc[nx,3] = e
    end

    if Int(bdcon[1][2]) == 0
        for i in 1:3
            Qc[1,i] = Qc[2,i]
        end
    elseif Int(bdcon[1][2]) == 1
        rho = bdcon[2][2]
        u = bdcon[2][3]
        p = bdcon[2][4]
        e = p/(gamma-1)+rho*(u^2)/2
        Qc[1,1] = rho
        Qc[1,2] = rho*u
        Qc[1,3] = e
    end

    return Qc
end

# --------------------------
# -- function GS          --
# --------------------------
#
# gauss-seidel法
#
function GS(Qc,qf,nx,dx,dt,RHS,yacobiA,norm_ok)
    delta_Q = zeros(nx,3)
    delta_Q2 = zeros(nx,3)
    delta_Q_temp = zeros(nx,3)

    sum_b = zeros(3)
    for i in 2:nx-1
        for j in 1:3
            sum_b[j] += abs(RHS[i,j])
        end
    end

    ite=0
    con=0
    while con == 0　# gs iteration
        delta_Q_temp = copy(delta_Q)

        # LDU分解
        L = zeros(nx,3)
        D = zeros(nx)
        U = zeros(nx,3)
        for i in 1:nx
            temp1 = zeros(3,3)
            temp2 = zeros(3)
            for l in 1:3
                for m in 1:3
                    temp1[l,m] = yacobiA[i,l,m]
                end
                temp2[l] = delta_Q[i,l]
            end
            temp = nm_inner_product(temp1,temp2)

            for j in 1:3
                L[i,j] = dt*0.5*temp[j]
                U[i,j] = dt*0.5*temp[j]
            end
            D[i] = dx
        end

        RHS_temp = zeros(nx,3)
        # RHS
        for i in 2:nx-1
            for j in 1:3
                RHS_temp[i,j] = -dt*RHS[i,j]
            end
        end

        # (D+L)Q=Q
        for i in 2:nx-1
            for j in 1:3
                delta_Q[i,j] = (L[i-1,j]+RHS_temp[i,j]) / D[i]
            end
        end

        # (D+U)Q=DQ
        for i in 2:nx-1
            for j in 1:3
                delta_Q[i,j] = delta_Q[i,j] - U[i+1,j]/ D[i]
            end
        end

        # 収束判定
        if (ite+1) % 100 == 0
            sum_b_Ax = zeros(3)

            for i in 2:nx-1
                for j in 1:3
                    sum_b_Ax[j] += abs(delta_Q[i,j]-delta_Q_temp[i,j])
                end
            end

            norm2d = zeros(3)

            for i in 1:3
                norm2d[i] = sum_b_Ax[i]/sum_b[i]
            end

            if norm2d[1] < norm_ok && norm2d[2] < norm_ok && norm2d[3] < norm_ok
                con=1
            end
        end
        ite += 1

        if ite > 1.0e5
            println("Number of iterations exceeded 1.0e5 ")
            throw(UndefVarError(:x))
        end
    end

    # Qcの更新
    for i in 2:nx-1
        for j in 1:3
            Qc[i,j] = Qc[i,j]+delta_Q[i,j]
        end
    end
    return Qc
end

# --------------------------
# -- function cal_RHS     --
# --------------------------
#
# 右辺の計算
# RHSを定義
#
function cal_RHS(Fplus,nx)  # 境界フラックスの計算

    RHS = zeros(nx,3)

    for i in 2:nx-1
        for j in 1:3
            RHS[i,j] = Fplus[i,j]-Fplus[i-1,j]
        end
    end
    return RHS
end

# --------------------------
# -- function fvs         --
# --------------------------
#
# Flux Vector Splitting method
# Fplusを定義
#
function fvs(QcL,QcR,qfL,qfR,nx,gamma)  # FVS法によるフラックスの計算(セル1と2の境界をFplus[1]に格納)

    Fplus = zeros(nx+1,3)

    for i in 1:nx-1
        H = (QcL[i,3]+qfL[i,3])/QcL[i,1]  # エンタルピー
        u = qfL[i,1]
        c = ((gamma-1)*(H-0.5*u^2))^0.5
        b_para = (gamma-1)/c^2
        a_para = 0.5*b_para*u^2
        # i+1/2セルにおけるR,R^-1,Λ,|Λ|
        R, R_inv, Gam, Gam_abs = A_pm(H,u,c,b_para,a_para)
        Ap = nn_inner_product((nn_inner_product(R, Gam+Gam_abs)), R_inv)               # 固有値が正のものを計算

        H = (QcR[i,3]+qfR[i,3])/QcR[i,1]  # エンタルピー
        u = qfR[i,1]
        c = ((gamma-1)*(H-0.5*u^2))^0.5
        b_para = (gamma-1)/c^2
        a_para = 0.5*b_para*u^2
        # i+1/2セルにおけるR,R^-1,Λ,|Λ|
        R, R_inv, Gam, Gam_abs = A_pm(H,u,c,b_para,a_para)
        Am = nn_inner_product((nn_inner_product(R, Gam-Gam_abs)), R_inv)               # 固有値が負のものを計算

        tempL = zeros(3)
        tempR = zeros(3)
        for j in 1:3
            tempL[j] = QcL[i,j]
            tempR[j] = QcR[i,j]
        end

        tempL = nm_inner_product(Ap, tempL)
        tempR = nm_inner_product(Am, tempR)

        for j in 1:3
            Fplus[i,j] = 0.5*(tempL[j] + tempR[j])   # フラックスを計算tempL[j] =
        end
    end
    return Fplus
end

# --------------------------
# -- function A_pm       --
# --------------------------
#
# ヤコビアン行列の正負の計算に向け、固有値等を計算
# A = R^(-1)ΛR としたとき
# R = R
# R_inv = R^(-1)
# Gam = Λ
#
function A_pm(H,u,c,b_para,a_para)  # ヤコビアン行列の固有値もろもろ計算
    R     = [1.0 1.0 1.0;
            u-c u u+c;
            H-u*c 0.5*u^2 H+u*c]
    R_inv = [0.5*(a_para+u/c) 0.5*(-b_para*u-1/c) 0.5*b_para;
            1-a_para b_para*u -b_para;
            0.5*(a_para-u/c) 0.5*(-b_para*u+1/c) 0.5*b_para]
    Gam   = [(u-c) 0.0 0.0; 0.0 u 0.0; 0.0 0.0 (u+c)]
    Gam_abs = [abs(u-c) 0.0 0.0; 0.0 abs(u) 0.0; 0.0 0.0 abs(u+c)]

    return R, R_inv, Gam, Gam_abs
end

# --------------------------
# -- function yacobi_A    --
# --------------------------
#
# ヤコビアン行列の計算
#
function yacobi_A(Qc,qf,gamma,nx)
    yacobiA = zeros(nx,4,4)

    for i in 1:nx
        H = (Qc[i,3]+qf[i,3])/Qc[i,1]  # エンタルピー
        u = qf[i,1]
        c = ((gamma-1)*(H-0.5*u^2))^0.5
        b_para = (gamma-1)/c^2
        a_para = 0.5*b_para*u^2

        R, R_inv, Gam, Gam_abs = A_pm(H,u,c,b_para,a_para)

        temp = nn_inner_product((nn_inner_product(R, Gam)), R_inv)

        for j in 1:3
            for k in 1:3
                yacobiA[i,j,k] = temp[j,k]
            end
        end
    end

    return yacobiA
end

# -------------------------------
# -- function nn_inner_product --
# -------------------------------
#
# n*n行列同士の内積
#
function nn_inner_product(a,b)
    temp=zeros(size(a)[1],size(a)[1])
    for i in 1:size(a)[1]
        for j in 1:size(a)[1]
            for k in 1:size(a)[1]
                temp[i,j] += a[i,k]*b[k,j]
            end
        end
    end
    return temp
end

# -------------------------------
# -- function nm_inner_product --
# -------------------------------
#
# n*m行列と1*m行列の内積
#
function nm_inner_product(a,b)
    temp = zeros(size(a)[1])
    for i in 1:size(a)[1]
        for k in 1:size(a)[2]
            temp[i] += a[i,k]*b[k]
        end
    end
    return temp
end

# -------------------------------
# -- function muscl            --
# -------------------------------
#
# muscl法による補間
#
function muscl(Qc,qf,nx,k_muscl,b_muscl,gamma)

    qfL = zeros(nx+1,3)
    qfR = zeros(nx+1,3)

    for i in 2:nx-2
        for j in 1:3
            dplus_j   = qf[i+1,j]-qf[i,j]
            dminus_j  = qf[i,j]-qf[i-1,j]
            dplus_jp  = qf[i+2,j]-qf[i+1,j]
            dminus_jp = qf[i+1,j]-qf[i,j]

            qfL[i,j] = qf[i,j]+1/4*((1-k_muscl)*minmod(dminus_j,dplus_j,b_muscl)+(1+k_muscl)*minmod(dplus_j,dminus_j,b_muscl))
            qfR[i,j] = qf[i+1,j]-1/4*((1-k_muscl)*minmod(dplus_jp,dminus_jp,b_muscl)+(1+k_muscl)*minmod(dminus_jp,dplus_jp,b_muscl))
        end
    end

    # 境界内側用
    for j in 1:3
        dplus_jp  = qf[3,j]-qf[2,j]
        dminus_jp = qf[2,j]-qf[1,j]
        qfR[1,j] = qf[2,j]-1/4*((1-k_muscl)*minmod(dplus_jp,dminus_jp,b_muscl)+(1+k_muscl)*minmod(dminus_jp,dplus_jp,b_muscl))

        dplus_j      = qf[nx,j]-qf[nx-1,j]
        dminus_j     = qf[nx-1,j]-qf[nx-2,j]
        qfL[nx-1,j] = qf[nx-1,j]+1/4*((1-k_muscl)*minmod(dminus_j,dplus_j,b_muscl)+(1+k_muscl)*minmod(dplus_j,dminus_j,b_muscl))
    end

    QcL = qftoQc(qfL,nx,gamma)
    QcR = qftoQc(qfR,nx,gamma)

    # 境界外側用(境界は風上)
    for i in 1:3
        qfL[1,i] = qf[1,i]
        QcL[1,i] = Qc[1,i]
        qfR[nx-1,i] = qf[nx,i]
        QcR[nx-1,i] = Qc[nx,i]
    end

    return QcL,QcR,qfL,qfR
end

# -------------------------------
# -- function minmod           --
# -------------------------------
#
# 流速制限関数minmod
#
function minmod(x,y,b)
    ans = sign(x)*max(0,min(abs(x),sign(x)*y*b))
    return ans
end

# -------------------------------
# -- function qftoQc           --
# -------------------------------
#
# 基本量から保存量に変換
#
function qftoQc(qf,nx,gamma)
    lo_Qc = zeros(nx,3)
    for i in 1:nx
        lo_Qc[i,1] = qf[i,2]
        lo_Qc[i,2] = qf[i,2]*qf[i,1]
        lo_Qc[i,3] = (qf[i,3]/(gamma-1)+1.0/2.0*qf[i,2]*(qf[i,1]^2))
    end

    return lo_Qc
end

# -------------------------------
# -- function Qctoqf           --
# -------------------------------
#
# 保存量から基本量に変換
#
function Qctoqf(Qc,nx,gamma)  # 保存量から基本量変換
    lo_qf = zeros(nx,3)
    for i in 1:nx
        lo_qf[i,1] = Qc[i,2]/Qc[i,1]
        lo_qf[i,2] = Qc[i,1]
        lo_qf[i,3] = (gamma-1)*(Qc[i,3]-1.0/2.0*Qc[i,1]*((Qc[i,2]/Qc[i,1])^2))
    end
    return lo_qf
end

# -------------------------------
# -- function output_q         --
# -------------------------------
#
# x,rho,u,pの出力
#
function output_q(x,qf,nx,out_dir,out_file_front,stepnum,out_ext)
    fff = out_dir*"/"*out_file_front*string(stepnum)*out_ext
    open(fff,"w") do f
        write(f,"result:x, rho[kg/m^3], u[m/s], p[Pa]\n")
        for i in 2:nx-1
            a = @sprintf("%8.8e", x[i])
            b = @sprintf("%8.8e", qf[i,1])
            c = @sprintf("%8.8e", qf[i,2])
            d = @sprintf("%8.8e", qf[i,3])
            write(f,a*" "*b*" "*c*" "*d*"\n")
        end
    end
end

# -------------------------------
# -- function cre_dir          --
# -------------------------------
#
# ディレクトリの作成
#
function cre_dir(outdir)
    k = 0
    try rm(outdir,recursive=true)
    catch
        mkdir(outdir)
        k = 1
    end

    if k == 0
        mkdir(outdir)
    end
end

# -------------------------------
# -- main                      --
# -------------------------------
function main()
    # --------------------------
    # -- initial value        --
    # --------------------------
    nstep = 300                     # 時間ステップ数
    nx0   = 100                     # 空間ステップ数
    dt    = 0.002                   # 時間刻み幅
    dx    = 0.01                    # 空間刻み幅

    # -- slope limiter --
    k_muscl = 1/3                      # muscl精度
    b_muscl = (3-k_muscl)/(1-k_muscl)

    # -- constant --
    gamma = 1.4                    # 比熱比

    # -- boundary --
    # bd1_con = 0 :流入条件
    #         = 1 :流出条件
    #
    # -- x+ right --
    bd1_con = 0
    bd1_rho = 1.00
    bd1_u   = 0.0
    bd1_p   = 1.0
    bd1_T   = 300.0
    # -- x- left --
    bd2_con = 1
    bd2_rho = 1.0
    bd2_u   = 0.0
    bd2_p   = 1.0
    bd2_T   = 300.0

    # -- itetarive --
    norm_ok = 1.0e-4               # 収束条件

    # -- 出力--
    dir_name       = "sod_central_0.002s_jl"           # 出力フォルダ名
    out_name_front = "time"                 # 出力ファイル名（先頭）
    out_name_back  = "d-3"

    # --------------------------
    # -- prepare cal          --
    # --------------------------
    lbound = 1                     # 仮想境界セル数
    nx = nx0+2*lbound              # 総空間セル数

    bdcon = [[bd1_con, bd1_rho, bd1_u, bd1_p, bd1_T],
            [bd2_con, bd2_rho, bd2_u, bd2_p, bd2_T]]

    # --------------------------
    # -- setup                --
    # --------------------------
    cre_dir(dir_name)
    x,qf,Qc = setup(nx,dx,gamma)

    # --------------------------
    # -- main loop            --
    # --------------------------
    for t in 1:nstep
        print(t)

        Qc = boundary(nx,gamma,qf,Qc,bdcon)
        QcL,QcR,qfL,qfR = muscl(Qc,qf,nx,k_muscl,b_muscl,gamma)

        Fplus = fvs(QcL,QcR,qfL,qfR,nx,gamma)
        RHS = cal_RHS(Fplus,nx)

        yacobiA = yacobi_A(Qc,qf,gamma,nx)
        Qc = GS(Qc,qf,nx,dx,dt,RHS,yacobiA,norm_ok)

        qf=Qctoqf(Qc,nx,gamma)

        stepnum = Int(round(t*dt*1000))
        output_q(x,qf,nx,dir_name,out_name_front,stepnum,out_name_back)
    end
end

# main
main()


module TestCompare
    using Benchmark

    f1() = svd(zeros(100, 100))
    f2() = svd(ones(100, 100))
    f3() = svd(eye(100, 100))

    benchmark(f1, "Linear Algebra", "svd(zeros(100, 100))", 10)
    benchmark(f1, "svd(zeros(100, 100))", 10)

    benchmark(f2, "Linear Algebra", "svd(ones(100, 100))", 10)
    benchmark(f2, "svd(ones(100, 100))", 10)

    compare([f1, f2, f3], 100)
end


#!/usr/bin/env julia

function main()
    if length(ARGS) > 0
        n = parse(Int64, ARGS[1])
        save_n = n
        factors = Int64[]
        d = 2
        while n > 1
            while n % d == 0
                push!(factors, d)
                n = div(n,d)
            end
            d += 1
            if d * d > n
                if n > 1
                    push!(factors, n)
                end
                break
            end
        end

        @printf("The largest prime factor of %d is %d\n",
                save_n, factors[length(factors)])
    else
        error("this program needs a single numeric argument")
    end
end

main()


using NaNMath
using Base.Test

@test isnan(NaNMath.log(-10))
@test isnan(NaNMath.log1p(-100))
@test isnan(NaNMath.pow(-1.5,2.3))
@test isnan(NaNMath.sqrt(-5))
@test NaNMath.sqrt(5) == Base.sqrt(5)
@test NaNMath.sum([1., 2., NaN]) == 3.0
@test isnan(NaNMath.sum([NaN, NaN]))
@test NaNMath.sum(Float64[]) == 0.0
@test NaNMath.sum([1f0, 2f0, NaN32]) === 3.0f0
@test NaNMath.maximum([1., 2., NaN]) == 2.0
@test NaNMath.minimum([1., 2., NaN]) == 1.0
@test NaNMath.mean([1., 2., NaN]) == 1.5
@test NaNMath.var([1., 2., NaN]) == 0.5
@test NaNMath.std([1., 2., NaN]) == 0.7071067811865476


using LinearAlgebraicRepresentation, SparseArrays
Lar = LinearAlgebraicRepresentation;
using IntervalTrees,LinearAlgebra
# using Revise, OhMyREPL

function settestpoints2d(W,EV,FV,f, copEV,copFE) # W by rows
	e = findnz(copFE[f,:])[1][1] # first edge of face f
	v1,v2 = findnz(copEV[e,:])[1] # two (global) verts incident on in
    t = W[v2,:] - W[v1,:]
    n = [-t[2],t[1]]
	p0 = (W[v1,:] + W[v2,:]) ./ 2
	ϵ = 1.0e-4
	ptest1 = p0 + ϵ * n
	ptest2 = p0 - ϵ * n
	return ptest1, ptest2
end

function getinternalpoint2d(W,EV,FV, f, copEV,copFE) # W by rows
	#edges for v1=FV[1][1]
	ptest1, ptest2 = Lar.settestpoints2d(W,EV,FV, f, copEV,copFE)
    edges = [findnz(copEV[e,:])[1] for e in findnz(copFE[f,:])[1]]
    V = convert(Lar.Points,W')
    classify = Lar.pointInPolygonClassification(V,edges)
    if classify(ptest1) == "p_in"
        return ptest1
    elseif classify(ptest2) == "p_in"
        return ptest2
    else
        error("classifying inner point in face $f")
    end
end

function chainbasis2polygons(V,copEV,copFE)
	FE = [findnz(copFE[k,:])[1] for k=1:copFE.m]
	EV = [findnz(copEV[k,:])[1] for k=1:copEV.m]

	FEs = Array{Int64,1}[]
	EVs = Array{Array{Int64,1},1}[]
	FVs = Array{Int64,1}[]
	for f=1:copFE.m
		push!( FEs, collect(Set(cat([e for e in FE[f]]))) )
		# edges in EVs are aggregated by face, in order to answer point-classifications
		push!( EVs, [EV[e] for e in FE[f]]  )
		push!( FVs, collect(Set(cat([EV[e] for e in FE[f]])))  )
	end
	polygons = collect(zip(EVs,FVs,FEs)) # (EV,FV,FFE)s for polygon
	W = convert(Lar.Points,V')
	return W,polygons,FE
end

function internalpoints2d(W,copEV,copFE) # W by rows
	# transform each 2-cell in a solid (via Lar model)
	#-------------------------------------------------------------------------------
	U,pols,FE = Lar.chainbasis2polygons(W,copEV,copFE) # V by rows
	# compute, for each `pol` (3-cell) in `pols`, one `internalpoint`.
	#-------------------------------------------------------------------------------
	internalpoints = []
	for f=1:length(pols)
		(EV,FV,FE) = pols[f]
		#GL.VIEW([ GL.GLFrame, GL.GLLines(V,EV) ]);
		internalpoint = Lar.getinternalpoint2d(W,EV,FV, f, copEV,copFE)
		push!(internalpoints,internalpoint)
	end
	return internalpoints
end

"""
	testinternalpoint2d(V::Lar.Points, EV::Lar.Cells, FV::Lar.Cells)

"""
function testinternalpoint2d(listOfModels)
	function testinternalpoint0(testpoint)
		intersectedfaces = Int64[]
		depot = []
		# actual containment test of ray point in faces within depot
		for (k,model) in enumerate(listOfModels)
            verts,edges = model
			classify = Lar.pointInPolygonClassification(verts,edges)
			inOut = classify(testpoint)
			if inOut == "p_in"
				push!(intersectedfaces,k)
			end
		end
		return intersectedfaces
	end
	return testinternalpoint0
end

################################################################################

function bool2d(assembly)
	# input of affine assembly
	#-------------------------------------------------------------------------------
	# V,EV = Lar.struct2lar(assembly) #TODO proper different method
	V,EV = Lar.struct2lar(assembly)
	cop_EW = convert(Lar.ChainOp, Lar.coboundary_0(EV::Lar.Cells));
	W = convert(Lar.Points, V');
	# generate the 3D space arrangement
	#-------------------------------------------------------------------------------
	W, copEV, copFE = Lar.Arrangement.planar_arrangement(W::Lar.Points, cop_EW::Lar.ChainOp)
	#V,CVs,FVs,EVs = Lar.pols2tria(W, copEV, copFE, copCF)
	innerpoints = Lar.internalpoints2d(W,copEV,copFE[1:end,:])
	# associate internal points to 2-cells
	#-------------------------------------------------------------------------------
	listOfModels = Lar.evalStruct(assembly)
	inputfacenumbers = [length(listOfModels[k][2]) for k=1:length(listOfModels)]
	# test input data for containment of reference points
	#-------------------------------------------------------------------------------
	boolmatrix = BitArray(undef, length(innerpoints), length(listOfModels))
    containmenttest = Lar.testinternalpoint2d(listOfModels)
	for (k,point) in enumerate(innerpoints) # k runs on columns
		cells = containmenttest(point) # contents of columns
		#println(k," ",faces)
		for l in cells
			boolmatrix[k,l] = 1
		end
	end
	return W, copEV, copFE, boolmatrix
end

"""
	boolops( assembly::Lar.Struct, op::Symbol )

User interface to 2d Boolean ops, where `op` symbol ``in`` {`:|`, `:&`, `:-`}.
Return a pair `(V,EV)`.
# Example
```julia
assembly = Lar.Struct([wall, openings])
V,EV = boolops(assembly, :-)
```
"""
function boolops(assembly::Lar.Struct, op::Symbol)
    W, copEV, copFE, boolmatrix = Lar.bool2d(assembly)
    boolvars = [boolmatrix[:,k] for k=1:size(boolmatrix,2)]
    if eval(op) == -
        solution = boolvars[1] .& .!(.|(boolvars[2:end]...))
    else
        operator = eval(op)
        solution = operator.(boolvars...)
    end
    chain1d = sparse(copFE') * Int.(solution)
    EV = Lar.cop2lar(copEV)
    EVop = [ev for (k,ev) in enumerate(EV) if abs(chain1d[k])==1 ]
    V = convert(Lar.Points,W')
    return V,EVop
end


using FirstOrderSolvers: conjugategradient!, CGdata

Random.seed!(2)
# a = randn(1000,1000)
# A = a'a
A = rand(1000,1000)
A = A'*A
b = randn(1000)
x0 = randn(1000)

x = copy(x0)
#anrm = sqrt(sum(abs2, A, 2))[:]
#A = A./anrm
#b = b./anrm
function test(A, b, x0, x, cgdata)
    x .= x0
    conjugategradient!(x, A, b, cgdata.r, cgdata.p, cgdata.z)
end

cgdata = CGdata(similar(b), similar(b), similar(b))
conjugategradient!(x, A, b, cgdata.r, cgdata.p, cgdata.z, max_iters = 100)
n0 = norm(A*x-b)    # 48
conjugategradient!(x, A, b, cgdata.r, cgdata.p, cgdata.z, max_iters = 5000)
n1 = norm(A*x-b)    # 3e-6

@test n1 < 1e-5

xcopy = x .+ 1e-5.*randn(size(x))
n2 = norm(A*xcopy-b) #0.9
conjugategradient!(xcopy, A, b, cgdata.r, cgdata.p, cgdata.z, max_iters = 100)
n3 = norm(A*xcopy-b) #3e-5

@test n3 < 10*n2


A = [  1.0  -2.0   3.0  -4.0  5.0;
       2.0  -1.0   0.0  -1.0  3.0;
      -1.0   0.0   4.0  -3.0  2.0;
      -1.0  -1.0  -1.0   1.0  3.0]
b = [1.0, 2.0, 3.0, 4.0]

m, n = size(A)

f = Translate(SqrNormL2(), -b)
f2 = LeastSquares(A, b)
lam = 0.1*norm(A'*b, Inf)
g = NormL1(lam)

x_star = [-3.877278911564627e-01, 0, 0, 2.174149659863943e-02, 6.168435374149660e-01]

## Nonfast/Nonadaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 150
println(sol)

# testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

# Nonfast/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 300
println(sol)

# Fast/Nonadaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2, fast=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 100
println(sol)

# testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

# Fast/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, adaptive=true, fast=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 200
println(sol)

# ZeroFPR/Nonadaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 15
println(sol)

#testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

# ZeroFPR/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 15

# PANOC/Nonadaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 20
println(sol)

# testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

## PANOC/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 20
println(sol)

# Douglas-Rachford

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.DRS(x0; f=f2, g=g, gamma=10.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 30
println(sol)

#testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

#####################
# Complex Variables #
#####################

Random.seed!(123)
m, n = 10,5

A = randn(m,n)+im*randn(m,n)
b = randn(m)+im*randn(m)

f = Translate(SqrNormL2(), -b)
f2 = LeastSquares(A, b)
lam = 0.01*norm(A'*b, Inf)
g = NormL1(lam)

x0 = zeros(n)+im*zeros(n)
it, x_star, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2, verbose = 0)

## Nonfast/Nonadaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 200
println(sol)

# testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

# Nonfast/Adaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 250
println(sol)

# Fast/Nonadaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2, fast=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 120
println(sol)

# testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

# Fast/Adaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, adaptive=true, fast=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 120
println(sol)

# ZeroFPR/Nonadaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 15
println(sol)

#testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

# ZeroFPR/Adaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 15

# PANOC/Nonadaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, gamma=1.0/opnorm(A)^2)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 20
println(sol)

# testing solver already at solution
@time it, x = ProximalAlgorithms.run!(sol)
@test it == 1

## PANOC/Adaptive

x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 20
println(sol)

#############################################
# Real Variables  mapped to Complex Numbers #
#############################################

Random.seed!(123)
n = 2^6
A = AbstractOperators.DFT(n)[1:div(n,2)]      # overcomplete dictionary

x = sprandn(n,0.5)
b = fft(x)[1:div(n,2)]

#f = Translate(LogisticLoss(ones(n)), -b)
f = Translate(SqrNormL2(), -b)
lam = 0.01*norm(A'*b,Inf)
g = NormL1(lam)

x0 = zeros(n)
it, x_star, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, verbose = 0, tol = 1e-8)

# Nonfast/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, tol=1e-6, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 50
println(sol)

# Fast/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, tol=1e-6, adaptive=true, fast=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 50
println(sol)

# ZeroFPR/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, tol=1e-6, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 15
println(sol)

# PANOC/Adaptive

x0 = zeros(n)
@time it, x, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, tol=1e-6, adaptive=true)
@test norm(x - x_star, Inf) <= 1e-4
@test it < 20
println(sol)


struct Gencost
    "cost model, 1 = piecewise linear, 2 = polynomial "
    MODEL::Vector{Int64}
    "startup cost in US dollars"
    STARTUP::Vector{Float64}
    "shutdown cost in US dollars"
    SHUTDOWN::Vector{Float64}
    "number breakpoints in piecewise linear cost function, or number of coefficients in polynomial cost function"
    NCOST::Vector{Int64}
    "parameters defining total cost function begin in this col
        (MODEL = 1) : p0, f0, p1, f1, ..., pn, fn
             where p0 < p1 < ... < pn and the cost f(p) is defined
             by the coordinates (p0,f0), (p1,f1), ..., (pn,fn) of
             the end/break-points of the piecewise linear cost
        (MODEL = 2) : cn, ..., c1, c0
             n+1 coefficients of an n-th order polynomial cost fcn,
             starting with highest order, where cost is
             f(p) = cn*p^n + ... + c1*p + c0"
    COST::Vector{Vector{Float64}}
end

function extract_gencost(mpc)
    gencost = mpc["gencost"]
    ng = size(gencost, 1)
    coef = Vector{Vector{Float64}}([gencost[i,:][5:end] for i in 1:ng])
    return Gencost([gencost[:, i] for i in 1:4]..., coef)
end

"""
Shortcut for returning gencost info for one of the cases listed by
`casenames()`.
"""
function extract_gencost(cname::AbstractString)
    mpc = loadcase(cname; describe=false)
    extract_gencost(mpc)
end


using Combinatorics
using StatsBase

# Bounds: puzzle input
const LOWER = 387638
const UPPER = 919123

function int_from_digits(digits::AbstractVector)
    range = length(digits)-1:-1:0
    sum(@. digits * 10^range)
end

function countcombs(f)
    combs = Iterators.filter(f, with_replacement_combinations(1:9, 6))
    # Just iterate to avoid materializing a large array
    n = 0
    for c in combs
        n += 1
    end
    n
end

println("Part 1 answer: ", countcombs() do c
    issorted(c) && any(iszero, diff(c)) && LOWER ≤ int_from_digits(c) ≤ UPPER
end)

println("Part 2 answer: ", countcombs() do c
    issorted(c) && LOWER ≤ int_from_digits(c) ≤ UPPER && any(==(2), last(rle(c)))
end)