using IntervalArithmetic
using Test

setformat(:standard, decorations=true, sigfigs=6)

setprecision(128)

@testset "Parse string to Interval" begin
    @test parse(Interval{Float64}, "1") == Interval(1, 1)
    @test parse(Interval{Float64}, "[1, 2]") == Interval(1, 2)
    @test parse(Interval{Float64}, "[-0x1.3p-1, 2/3]") == @interval(-0x1.3p-1, 2/3)
    @test parse(Interval{Float64}, "[1,+infinity]") == Interval(1.0, Inf)
    @test parse(Interval{Float64}, "[1.234e5,Inf]") == Interval(123400.0, Inf)
    @test parse(Interval{Float64}, "[]") == ∅
    @test parse(Interval{Float64}, "[,]") == entireinterval(Float64)
    @test parse(Interval{Float64}, "[ entire ]") == entireinterval(Float64)
    @test parse(Interval{Float64}, "3.56?1") == Interval(0x3.8ccccccccccccp+0, 0x3.91eb851eb8520p+0)
    @test parse(Interval{Float64}, "3.56?1e2") == Interval(355.0, 357.0)
    @test parse(Interval{Float64}, "3.560?2") == Interval(0x3.8ed916872b020p+0, 0x3.8fdf3b645a1ccp+0)
    @test parse(Interval{Float64}, "3.56?") == Interval(0x3.8e147ae147ae0p+0, 0x3.90a3d70a3d70cp+0)
    @test parse(Interval{Float64}, "3.560?2u") == Interval(0x3.8f5c28f5c28f4p+0, 0x3.8fdf3b645a1ccp+0)
    @test parse(Interval{Float64}, "-10?") == Interval(-10.5, -9.5)
    @test parse(Interval{Float64}, "-10?u") == Interval(-10.0, -9.5)
    @test parse(Interval{Float64}, "-10?12") == Interval(-22.0, 2.0)
    @test parse(Interval{Float64}, "0.0?d") == Interval(-0.05, 0.0)
    @test parse(Interval{Float64}, "2.5?d") == Interval(0x1.3999999999999p+1, 2.5)
    @test parse(Interval{Float64}, "0.000?5d") == Interval(-0.005, 0.0)
    @test parse(Interval{Float64}, "2.500?5d") == Interval(0x1.3f5c28f5c28f5p+1, 2.5)
    @test parse(Interval{Float64}, "2.5??d") == Interval(-Inf, 2.5)
    @test parse(Interval{Float64}, "2.500?5ue4") == Interval(0x1.86ap+14, 0x1.8768p+14)
    @test parse(Interval{Float64}, "2.500?5de-5") == Interval(0x1.a2976f1cee4d5p-16, 0x1.a36e2eb1c432dp-16)
    @test parse(Interval{Float64}, "3.1416?1") == Interval(0x3.24395810624dcp+0, 0x3.24467381d7dc0p+0)

    @test parse(Interval{BigFloat}, "1") == @biginterval(1)
    @test parse(Interval{BigFloat}, "[-0x1.3p-1, 2/3]") == @biginterval(-0x1.3p-1, 2/3)

    @test parse(Interval{Float64}, "[Empty]") == emptyinterval(Float64)
    @test parse(Interval{BigFloat}, "[Empty]") == emptyinterval(BigFloat)

    @test parse(Interval{Float64}, "3 ±  4") == (-1)..7
    @test parse(Interval{Float64}, "0.2 ± 0.1") == 0.1..0.3
    @test parse(Interval{BigFloat}, "0.2 ± 0.1") == big"0.2" ± big"0.1"
end

@testset "Parse string to DecoratedInterval" begin
    @test parse(DecoratedInterval{Float64}, "[3]") == DecoratedInterval(3)

    @test parse(DecoratedInterval{Float64}, "[3, 4]") == DecoratedInterval(3, 4)

    @test parse(DecoratedInterval{Float64}, "[3, 4]_dac") == DecoratedInterval(3, 4, dac)

    @test parse(DecoratedInterval{Float64}, "[ ]") == DecoratedInterval(∅, trv)

    @test parse(DecoratedInterval{Float64}, "[entire]") == DecoratedInterval(Interval(-Inf, Inf), dac)

    @test parse(DecoratedInterval{Float64}, "[,]") == DecoratedInterval(entireinterval(Float64), dac)
    @test parse(DecoratedInterval{Float64}, "3.56?1") == DecoratedInterval(Interval(0x3.8ccccccccccccp+0, 0x3.91eb851eb8520p+0), com)
    @test parse(DecoratedInterval{Float64}, "3.56?1e2") == DecoratedInterval(Interval(355.0, 357.0), com)
    @test parse(DecoratedInterval{Float64}, "3.560?2") == DecoratedInterval(Interval(0x3.8ed916872b020p+0, 0x3.8fdf3b645a1ccp+0), com)
    @test parse(DecoratedInterval{Float64}, "3.56?") == DecoratedInterval(Interval(0x3.8e147ae147ae0p+0, 0x3.90a3d70a3d70cp+0), com)
    @test parse(DecoratedInterval{Float64}, "3.56?e2") == DecoratedInterval(Interval(3.555e2, 3.565e2), com)
    @test parse(DecoratedInterval{Float64}, "3.560?2u") == DecoratedInterval(Interval(0x3.8f5c28f5c28f4p+0, 0x3.8fdf3b645a1ccp+0), com)
    @test parse(DecoratedInterval{Float64}, "-10?") == DecoratedInterval(Interval(-10.5, -9.5), com)
    @test parse(DecoratedInterval{Float64}, "-10?e2") == DecoratedInterval(Interval(-10.5e2, -9.5e2), com)
    @test parse(DecoratedInterval{Float64}, "-10?u") == DecoratedInterval(Interval(-10.0, -9.5), com)
    @test parse(DecoratedInterval{Float64}, "-10?12") == DecoratedInterval(Interval(-22.0, 2.0), com)
    @test parse(DecoratedInterval{Float64}, "-10??u") == DecoratedInterval(Interval(-10.0, Inf), dac)
    @test parse(DecoratedInterval{Float64}, "-10??") == DecoratedInterval(Interval(-Inf, Inf), dac)
    @test isnai(parse(DecoratedInterval{Float64}, "[nai]"))
    @test parse(DecoratedInterval{Float64}, "3.56?1_def") == DecoratedInterval(Interval(0x3.8ccccccccccccp+0, 0x3.91eb851eb8520p+0), def)
end


function csr_to_sparse_test()
    i = [1;2;3]
    j = [3;4;4]
    v = [8;9;10]
    (rp,ci,ai,m) = sparse_to_csr(i,j,v)
    (nzi,nzj,nzv) = csr_to_sparse(rp,ci,ai)
    A = sparse(nzi,nzj,nzv,length(rp)-1,maximum(ci))

    # more tests added here
    # clique to sparse test
    rp = collect(1:5:26)
    ci = vec(reshape(repmat(1:5,5,1)',25,1))
    ai = ones(Int64,25)
    A = csr_to_sparse_matrix(rp,ci,ai,5,5)
    if !isequal(full(A),ones(5,5))
        error("csr_to_sparse_test failed")
    end

    # 100 random trials
    for t = 1:100
        A = sprand(100,80,0.01)
        (rp,ci,ai) = sparse_to_csr(A)
        A2 = csr_to_sparse_matrix(rp,ci,ai,100,80)
        if ~isequal(A,A2)
            error("csr_to_sparse_test failed")
        end
    end
    return true
end


using RDatasets, DataArrays, Gadfly

plot(dataset("car", "SLID"), x=:Wages, color=:Language,
     Geom.histogram(position=:dodge))


function rand_weighted_index(cumulative_weights::Array{Float64, 1})
	val::Float64 = rand()
	idx::Int64 = 1
	while cumulative_weights[idx] < val
		idx += 1
	end
	return idx
end


using Interpolations
using Base.Threads
using Printf
using FFTW
using NetCDF
using Base.Filesystem

const G = begin
    #              ∞     (-1)^n * (2n + 1)
    # G(r) ≈ 1/π * Σ ——————————————————————————
    #             n=0 [(2n + 1)^2 + r^2]^(3/2)
    #
    # r ∈ [0; √(800)]
    Δr = .01
    l_r = 0:Δr:(sqrt(800) + Δr)

    function Σ(r::Float64)
        Σ = 0.
        for n ∈ 10000:-1:0
            frac = (2n + 1) / ((2n + 1)^2 + r^2)^1.5
            if n & 0x1 != 0 # n is odd, subtract
                Σ -= frac
            else # n is even, add
                Σ += frac
            end
        end

        return Σ
    end

    l_G = [1/Float64(π) * Σ(r) for r in l_r]

    itp = interpolate(l_G, BSpline(Constant()))
    extrapolate(scale(itp, l_r), Flat())
end

function precalculate_σ(h_min :: Float64, h_max :: Float64, Δx :: Float64, Δy :: Float64; n_h :: Float64 = 20.)
    Δh = (h_max - h_min) / 10
    if Δh == 0.; Δh = 1.; end
    l_h = max(0.0001, h_min):Δh:(h_max + Δh)

    #                       h^2
    # σ = ———————————————————————————————————————
    #       Σ   Σ  G(√[(nΔx)^2 + (mΔy)^2]/h)ΔxΔy
    #      n∈N m∈M
    #
    # N = {n ∈ Z| |nΔx| ≤ n_h * h}
    # M = {m ∈ Z| |mΔy| ≤ n_h * h}
    #
    # As the domain of n and m are point symmetric around 0, only ~1/4th of the sum element have to be evaluated:
    # * One quadrant (n > 0, m > 0), multiplied by 4
    # * One half of the x-axis (n > 0), multiplied by 2
    # * One hald of the y-axis (m > 0), multiplied by 2
    # * The point at the origin (n, m) = (0, 0)

    function σ(h)
        σ_inv = 0.
        n_max = ceil(Int, (n_h * h / Δx))
        m_max = ceil(Int, (n_h * h / Δy))

        N = n_max:-1:1
        M = m_max:-1:1

        for n ∈ N, m ∈ M
            σ_inv += 4*G(sqrt((n*Δx)^2 + (m*Δy)^2) / h)
        end

        for n ∈ N
            σ_inv += 2*G(n*Δx / h)
        end

        for m ∈ M
            σ_inv += 2*G(m*Δy / h)
        end

        σ_inv += G(0)

        return h^2 / (σ_inv * Δx * Δy)
    end

    l_σ = [σ(h) for h ∈ l_h]

    itp = interpolate(l_σ, BSpline(Cubic(Flat(OnGrid()))))
    return extrapolate(scale(itp, l_h), Flat())
end

function apply_kajiura!(b::AbstractArray{Float64, 2}, d::AbstractArray{Float64, 2}, η::AbstractArray{Float64, 2},
                h_min :: Float64, h_max :: Float64, Δx :: Float64, Δy :: Float64; water_level :: Float64 = 0., n_h :: Float64 = 20.,
                σ = precalculate_σ(h_min, h_max, Δx, Δy; n_h=n_h))

    filter_nx_half = ceil(Int, n_h * h_max / Δx / 2)
    filter_ny_half = ceil(Int, n_h * h_max / Δy / 2)

    nx = size(η, 2)
    ny = size(η, 1)

    filter = zeros(Float64, (ny, nx))

    @assert filter_nx_half < nx
    @assert filter_ny_half < ny

    println("  Computing filter matrix...")
    Threads.@threads for x ∈ -filter_nx_half:filter_nx_half
        for y ∈ -filter_ny_half:filter_ny_half
            filter[((ny + y) % ny) + 1, ((nx + x) % nx) + 1] = G(sqrt((x * Δx)^2 + (y * Δy)^2) / h_max)
        end
    end

    println("  Computing displacement matrix...")
    Threads.@threads for x ∈ 1:size(η, 2)
        for y ∈ 1:size(η, 1)
            h_yx = max(0., water_level - b[y, x]) # height 0 on land

            η[y, x] =
                if h_yx != 0.
                    σ(h_yx) * Δx * Δy / h_yx^2 * d[y, x]
                else
                    0. # No displacement where no water is (Kajiura not applicable)
                end
        end
    end

    println("  Computing filter FFT...")
    Filter = fft(filter)

    println("  Computing displacement FFT...")
    Η = fft(η)

    println("  Computing product...")
    Η .*= Filter

    println("  Computing IFFT...")
    η_complex = reinterpret(Float64, ifft(Η))
    copyto!(η, @view η_complex[1:2:end-1])
end

function main()
    if length(ARGS) != 3
        println("Usage: julia ./kajiura.jl <in_file.nc> <out_file.nc> <timestep_end>")
        return
    end

    println("Using $(nthreads()) threads.")

    in_filename = ARGS[1]
    out_filename = ARGS[2]
    t_end = parse(Int, ARGS[3])
    nc = NetCDF.open(in_filename, mode=NC_NOWRITE)
    d = nc["d"]

    lx = nc["x"]
    ly = nc["y"]

    Δx = lx[2] - lx[1]
    Δy = ly[2] - ly[1]

    nx = length(lx)
    ny = length(ly)

    l_times = nc["time"]
    n_times = length(l_times)
    last_timestamp = l_times[end]

    if t_end == -1 || t_end > n_times - 1
	    t_end = n_times - 1
    end

    println("Processing $nx × $ny cells over $(t_end + 1) timesteps.")

    b = Array{Float64, 2}(undef, (ny, nx))
    ncread!(in_filename, "b", b)
    current_disp = zeros(ny, nx)
    current_η_diff = Array{Float64, 2}(undef, (ny, nx))
    current_d_diff = Array{Float64, 2}(undef, (ny, nx))

    println("Creating output file")

    isfile(out_filename) && rm(out_filename)
    timeatts = Dict("units" => "seconds")
    xatts = Dict("units" => "m")
    yatts = Dict("units" => "m")

    nccreate(out_filename, "eta_diff", "y", ly, yatts, "x", lx, xatts, "time", l_times[1:t_end+1], timeatts)
    nccreate(out_filename, "d_diff", "y", "x", "time")

    nccreate(out_filename, "b", "y", "x")
    ncwrite(b, out_filename, "b")

    for t in 1:t_end + 1
        println("Working on timestep $(t - 1) of $(t_end)")
        current_η_diff .= 0.
        current_d_diff .= d[:,:,t] .- current_disp
        apply_kajiura!(b .+ current_disp, current_d_diff, current_η_diff, -maximum(b), -minimum(b), Δx, Δy)
        current_disp = d[:,:,t]

        println("  Writing output for timestep")
        ncwrite(current_d_diff, out_filename, "d_diff", start=[1,1,t], count=[-1,-1,1])
        ncwrite(current_η_diff, out_filename, "eta_diff", start=[1,1,t], count=[-1,-1,1])
    end
end

println("Using $(nthreads()) threads.")

Pkg.precompile()
println("Done.")

main()


using CompScienceMeshes, BEAST

Γ = meshcuboid(1.0, 1.0, 1.0, 0.15)
X = raviartthomas(Γ)

κ,  η  = 1.0, 1.0
κ′, η′ = 1.4κ, η/1.4

γ, γ′ = im*κ, im*κ′
T,  K  = MWSingleLayer3D(γ),  MWDoubleLayer3D(γ)
T′, K′ = MWSingleLayer3D(γ′), MWDoubleLayer3D(γ′)

o, x, y, z = euclidianbasis(3)
d = normalize(x+y+z); p = normalize(d × z)
E = planewavemw3d(direction=d, polarization=p, wavenumber=κ)
H = -1/(im*κ*η)*curl(E)

e, h = (n × E) × n, (n × H) × n

@hilbertspace j m
@hilbertspace k l

α, α′ = 1/η, 1/η′
pmchwt = @discretise(
    (η*T+η′*T′)[k,j] +      (K+K′)[k,m] -
         (K+K′)[l,j] + (α*T+α′*T′)[l,m] == h[k] + e[l],
    j∈X, m∈X, k∈X, l∈X)

u = solve(pmchwt)

## Post-processing
u1 = u[1:numfunctions(X)]
fcr, geo = facecurrents(u1,X)

include(Pkg.dir("CompScienceMeshes","examples","plotlyjs_patches.jl"))
p = patch(geo, real.(norm.(fcr)))
PlotlyJS.plot(p)


using LazySets: translate, ispermutation

for N in [Float64, Rational{Int}, Float32]
    # random line
    rand(Line)

    # construction
    l1 = Line(from=N[0, 1], to=N[1, 1]) # two points in the line
    l2 = Line(N[0, 1], N[1, 0]) # point and direction

    # construction given a 2d direction and offset
    ll = Line(N[0, 1], N(1)) # y == 1
    @test N[0, 1] ∈ ll && N[1, 1] ∈ ll
    ll = Line(N[1, 0], N(1)) # x == 1
    @test N[1, 0] ∈ ll && N[1, 1] ∈ ll
    ll = Line(N[1, -1], N(0)) # x == y
    @test N[0, 0] ∈ ll && N[1, 1] ∈ ll

    # the lines are the same modulo the sign of the normal vector
    @test l1.p ≈ l2.p && l1.d ≈ -l2.d

    # dimension
    @test dim(l1) == 2

    # support function
    @test ρ(N[0, 1], l1) == N(1)
    @test ρ(N[1, 0], l1) == N(Inf)

    # support vector
    @test σ(N[0, 1], l1) == N[0, 1]
    @test_throws ArgumentError σ(N[1, 0], l1) == N(Inf)

    # boundedness
    @test !isbounded(l1)

    # universality
    @test !isuniversal(l1)

    # isempty
    @test !isempty(l1)

    # an_element and membership
    an_element(l1) ∈ l1

    # constraints_list
    if N <: AbstractFloat
        @test ispermutation(constraints_list(l2),
                            [HalfSpace(N[0, 1], N(1)),    # y <= 1
                             HalfSpace(N[0, -1], N(-1))]) # y >= 1
    end

    # translation
    @test translate(l2, N[0, 1]) == Line(N[0, 2], N[1, 0])

    # distance
    distance(N[1, 0], l1) == N(1)
    distance(l1, N[1, 0]) == N(1)
    distance(Singleton(N[1, 0]), l1) == N(1)
    distance(l1, Singleton(N[1, 0])) == N(1)

    # concrete linear map
    if N <: AbstractFloat
        mirror = N[-1 0; 0 1]
        l = Line(from=N[0, 1], to=N[1, 1])
        @test isequivalent(linear_map(mirror, l), l)
        rot = N[0 -1; 1 0] # π/2 ccw rotation
        @test isequivalent(linear_map(rot, l), Line(N[-1, 0], N[0, 1]))
    end

    # projection
    L = Line(N[1, 2, 3], N[1, 0, 0])
    @test project(L, [1]) == Universe{N}(1)
    @test project(L, [2]) == Singleton(N[2])
    @test project(L, [3]) == Singleton(N[3])
    @test project(L, [1, 2]) == Line(N[1, 2], N[1, 0])
    @test project(L, [1, 3]) == Line(N[1, 3], N[1, 0])
    @test project(L, [2, 3]) == Singleton(N[2, 3])
end


using AstroLib: kepler_solver

struct KeplerianOrbit <: AbstractOrbit
    a
    ecc
    period
    ρ_star
    r_star
    m_star
    m_planet
    m_total
    n
    a_star
    a_planet
    K0
    M0
    t_periastron
    t0
    tref
    incl
    sin_incl
    cos_incl
end

const G_grav = 2942.2062175044193

function KeplerianOrbit(; period, m_planet=0, r_star=1, m_star=1, ρ_star = 3 * m_star / (4 * π * r_star^3), ecc=nothing, t_periastron=0, t0 = 0 + t_periastron)
    a = cbrt(G_grav * (m_star + m_planet) * period^2 / (4 * π^2))
    m_total = m_star + m_planet
    n = 2 * π / period
    a_star = a * m_planet / m_total
    a_planet = -a * m_star / m_total
    K0 = n * a / m_total
    if ecc === nothing
        M0 = 0.5 * π
        incl_factor = 1
    else
        error("not implemented")
    end
    incl = 0.5 * π
    sin_incl, cos_incl = sincos(incl)

    tref = t_periastron - t0
    return KeplerianOrbit(a,
                          ecc,
                          period,
                          ρ_star,
                          r_star,
                          m_star,
                          m_planet,
                          m_total,
                          n,
                          a_star,
                          a_planet,
                          K0,
                          M0,
                          t_periastron,
                          t0,
                          tref,
                          incl,
                          sin_incl,
                          cos_incl)
end

function relative_position(orbit::KeplerianOrbit, t)
    sinf, cosf = get_true_anomaly(orbit, t)
    if orbit.ecc === nothing
        r = -orbit.a
    else
        r = -orbit.a * (1 - orbit.ecc^2) / (1 + orbit.ecc * cosf)
    end
    return rotate_vector(orbit, r * cosf, r * sinf)
end

function get_true_anomaly(orbit::KeplerianOrbit, t)
    M = orbit.n * ((t - orbit.t0) - orbit.tref)
    return kepler(M, orbit.ecc)
end

function kepler(M, ecc)
    E = kepler_solver(M, ecc)
    return sincos(E)
end
kepler(M, ::Nothing) = sincos(M)

function rotate_vector(orbit::KeplerianOrbit, x, y)
    # rotate about z0 axis by ω
    if orbit.ecc === nothing
        x1, y1 = x, y
    else
        # x1 = orbit.cosω * x - orbit.sinω * y
        # y1 = orbit.sinω * x + orbit.cosω * y
    end

    # rotate about x1 axis by -incl
    x2 = x1
    y2 = orbit.cos_incl * y1
    Z = -orbit.sin_incl * y1

    return SA[x2, y2, Z]

    # rotate about z2 axis by Ω
    # if orbit.Ω === nothing
    #     return SA[x2, y2, Z]
    # end
    # X = orbit.cosΩ * x2 - orbit.sinΩ * y2
    # Y = orbit.sinΩ * x2 + orbit.cosΩ * y2
    # return SA[X, Y, Z]
end


using DelimitedFiles
using LinearAlgebra
using FastGaussQuadrature
include("Quadrature.jl")
include("standardQuadrature.jl")

# gauss
for n=2:2:100
	p,w = Quadrature(n,1)
	p = [p; p[:,1] p[:,2] -p[:,3]]
	w = [w;w]

	m = [p w]
	display(n)
	name = string(n)*"_gauss.txt"
	writedlm(name,m)
end
#
## octa
#for n=2:2:50
#	p,w = Quadrature(n,2)
#	m = [p w]
#	display(n)
#	name = string(n)*"_s_octa.txt"
#	writedlm(name,m)
#end
#
## ico
#for n=2:2:50
#	p,w = Quadrature(n,3)
#	m = [p w]
#	display(n)
#	name = string(n)*"_s_ico.txt"
#	writedlm(name,m)
#end


@testset "c2i00a" begin
    rc2i = c2i00a(2400000.5, 53736.0)
    @test rc2i[1,1] ≈ 0.9999998323037165557 atol=1e-12
    @test rc2i[1,2] ≈ 0.5581526348992140183e-9 atol=1e-12
    @test rc2i[1,3] ≈ -0.5791308477073443415e-3 atol=1e-12
    @test rc2i[2,1] ≈ -0.2384266227870752452e-7 atol=1e-12
    @test rc2i[2,2] ≈ 0.9999999991917405258 atol=1e-12
    @test rc2i[2,3] ≈ -0.4020594955028209745e-4 atol=1e-12
    @test rc2i[3,1] ≈ 0.5791308472168152904e-3 atol=1e-12
    @test rc2i[3,2] ≈ 0.4020595661591500259e-4 atol=1e-12
    @test rc2i[3,3] ≈ 0.9999998314954572304 atol=1e-12
end

@testset "c2i00b" begin
    rc2i = c2i00b(2400000.5, 53736.0)
    @test rc2i[1,1] ≈ 0.9999998323040954356 atol=1e-12
    @test rc2i[1,2] ≈ 0.5581526349131823372e-9 atol=1e-12
    @test rc2i[1,3] ≈ -0.5791301934855394005e-3 atol=1e-12
    @test rc2i[2,1] ≈ -0.2384239285499175543e-7 atol=1e-12
    @test rc2i[2,2] ≈ 0.9999999991917574043 atol=1e-12
    @test rc2i[2,3] ≈ -0.4020552974819030066e-4 atol=1e-12
    @test rc2i[3,1] ≈ 0.5791301929950208873e-3 atol=1e-12
    @test rc2i[3,2] ≈ 0.4020553681373720832e-4 atol=1e-12
    @test rc2i[3,3] ≈ 0.9999998314958529887 atol=1e-12
end

@testset "c2i06a" begin
    rc2i = c2i06a(2400000.5, 53736.0)
    @test rc2i[1,1] ≈ 0.9999998323037159379 atol=1e-12
    @test rc2i[1,2] ≈ 0.5581121329587613787e-9 atol=1e-12
    @test rc2i[1,3] ≈ -0.5791308487740529749e-3 atol=1e-12
    @test rc2i[2,1] ≈ -0.2384253169452306581e-7 atol=1e-12
    @test rc2i[2,2] ≈ 0.9999999991917467827 atol=1e-12
    @test rc2i[2,3] ≈ -0.4020579392895682558e-4 atol=1e-12
    @test rc2i[3,1] ≈ 0.5791308482835292617e-3 atol=1e-12
    @test rc2i[3,2] ≈ 0.4020580099454020310e-4 atol=1e-12
    @test rc2i[3,3] ≈ 0.9999998314954628695 atol=1e-12
end

@testset "c2ibpn" begin
    rbpn = [9.999962358680738e-1 -2.516417057665452e-3 -1.093569785342370e-3;
            2.516462370370876e-3 9.999968329010883e-1 4.006159587358310e-5;
            1.093465510215479e-3 -4.281337229063151e-5 9.999994012499173e-1]
    rc2i = c2ibpn(2400000.5, 50123.9999, rbpn)
    @test rc2i[1,1] ≈ 0.9999994021664089977 atol=1e-12
    @test rc2i[1,2] ≈ -0.3869195948017503664e-8 atol=1e-12
    @test rc2i[1,3] ≈ -0.1093465511383285076e-2 atol=1e-12
    @test rc2i[2,1] ≈ 0.5068413965715446111e-7 atol=1e-12
    @test rc2i[2,2] ≈ 0.9999999990835075686 atol=1e-12
    @test rc2i[2,3] ≈ 0.4281334246452708915e-4 atol=1e-12
    @test rc2i[3,1] ≈ 0.1093465510215479000e-2 atol=1e-12
    @test rc2i[3,2] ≈ -0.4281337229063151000e-4 atol=1e-12
    @test rc2i[3,3] ≈ 0.9999994012499173103 atol=1e-12
    @test_throws ArgumentError c2ibpn(2400000.5, 50123.9999, rbpn[1:2,:])
end

@testset "c2s" begin
    t, p = c2s([100.,-50.,25.])
    @test t ≈ -0.4636476090008061162 atol=1e-15
    @test p ≈ 0.2199879773954594463 atol=1e-15
    @test_throws ArgumentError c2s([100.,-50.])
end

@testset "c2ixy" begin
    x = 0.5791308486706011000e-3
    y = 0.4020579816732961219e-4
    rc2i = c2ixy(2400000.5, 53736., x, y)
    @test rc2i[1,1] ≈ 0.9999998323037157138 atol=1e-12
    @test rc2i[1,2] ≈ 0.5581526349032241205e-9 atol=1e-12
    @test rc2i[1,3] ≈ -0.5791308491611263745e-3 atol=1e-12
    @test rc2i[2,1] ≈ -0.2384257057469842953e-7 atol=1e-12
    @test rc2i[2,2] ≈ 0.9999999991917468964 atol=1e-12
    @test rc2i[2,3] ≈ -0.4020579110172324363e-4 atol=1e-12
    @test rc2i[3,1] ≈ 0.5791308486706011000e-3 atol=1e-12
    @test rc2i[3,2] ≈ 0.4020579816732961219e-4 atol=1e-12
    @test rc2i[3,3] ≈ 0.9999998314954627590 atol=1e-12
end

@testset "c2ixys" begin
    x =  0.5791308486706011000e-3
    y =  0.4020579816732961219e-4
    s = -0.1220040848472271978e-7
    rc2i = c2ixys(x, y, s)
    @test rc2i[1,1] ≈ 0.9999998323037157138 atol=1e-12
    @test rc2i[1,2] ≈ 0.5581984869168499149e-9 atol=1e-12
    @test rc2i[1,3] ≈ -0.5791308491611282180e-3 atol=1e-12
    @test rc2i[2,1] ≈ -0.2384261642670440317e-7 atol=1e-12
    @test rc2i[2,2] ≈ 0.9999999991917468964 atol=1e-12
    @test rc2i[2,3] ≈ -0.4020579110169668931e-4 atol=1e-12
    @test rc2i[3,1] ≈ 0.5791308486706011000e-3 atol=1e-12
    @test rc2i[3,2] ≈ 0.4020579816732961219e-4 atol=1e-12
    @test rc2i[3,3] ≈ 0.9999998314954627590 atol=1e-12
end

@testset "c2t00a" begin
    tta = 2400000.5
    uta = 2400000.5
    ttb = 53736.0
    utb = 53736.0
    xp = 2.55060238e-7
    yp = 1.860359247e-6
    rc2t = c2t00a(tta, ttb, uta, utb, xp, yp)
    @test rc2t[1,1] ≈ -0.1810332128307182668 atol=1e-12
    @test rc2t[1,2] ≈ 0.9834769806938457836 atol=1e-12
    @test rc2t[1,3] ≈ 0.6555535638688341725e-4 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834768134135984552 atol=1e-12
    @test rc2t[2,2] ≈ -0.1810332203649520727 atol=1e-12
    @test rc2t[2,3] ≈ 0.5749801116141056317e-3 atol=1e-12
    @test rc2t[3,1] ≈ 0.5773474014081406921e-3 atol=1e-12
    @test rc2t[3,2] ≈ 0.3961832391770163647e-4 atol=1e-12
    @test rc2t[3,3] ≈ 0.9999998325501692289 atol=1e-12
end

@testset "c2t00b" begin
    tta = 2400000.5
    uta = 2400000.5
    ttb = 53736.0
    utb = 53736.0
    xp = 2.55060238e-7
    yp = 1.860359247e-6
    rc2t = c2t00b(tta, ttb, uta, utb, xp, yp)
    @test rc2t[1,1] ≈ -0.1810332128439678965 atol=1e-12
    @test rc2t[1,2] ≈ 0.9834769806913872359 atol=1e-12
    @test rc2t[1,3] ≈ 0.6555565082458415611e-4 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834768134115435923 atol=1e-12
    @test rc2t[2,2] ≈ -0.1810332203784001946 atol=1e-12
    @test rc2t[2,3] ≈ 0.5749793922030017230e-3 atol=1e-12
    @test rc2t[3,1] ≈ 0.5773467471863534901e-3 atol=1e-12
    @test rc2t[3,2] ≈ 0.3961790411549945020e-4 atol=1e-12
    @test rc2t[3,3] ≈ 0.9999998325505635738 atol=1e-12
end

@testset "c2t06a" begin
    tta = 2400000.5
    uta = 2400000.5
    ttb = 53736.0
    utb = 53736.0
    xp = 2.55060238e-7
    yp = 1.860359247e-6
    rc2t = c2t06a(tta, ttb, uta, utb, xp, yp)
    @test rc2t[1,1] ≈ -0.1810332128305897282 atol=1e-12
    @test rc2t[1,2] ≈ 0.9834769806938592296 atol=1e-12
    @test rc2t[1,3] ≈ 0.6555550962998436505e-4 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834768134136214897 atol=1e-12
    @test rc2t[2,2] ≈ -0.1810332203649130832 atol=1e-12
    @test rc2t[2,3] ≈ 0.5749800844905594110e-3 atol=1e-12
    @test rc2t[3,1] ≈ 0.5773474024748545878e-3 atol=1e-12
    @test rc2t[3,2] ≈ 0.3961816829632690581e-4 atol=1e-12
    @test rc2t[3,3] ≈ 0.9999998325501747785 atol=1e-12
end

@testset "c2tcio" begin
    c = [0.9999998323037164738 0.5581526271714303683e-9 -0.5791308477073443903e-3;
         -0.2384266227524722273e-7 0.9999999991917404296 -0.4020594955030704125e-4;
         0.5791308472168153320e-3 .4020595661593994396e-4 0.9999998314954572365]
    era = 1.75283325530307
    p = [0.9999999999999674705 -0.1367174580728847031e-10 0.2550602379999972723e-6;
         0.1414624947957029721e-10 0.9999999999982694954 -0.1860359246998866338e-5;
         -0.2550602379741215275e-6 0.1860359247002413923e-5 0.9999999999982369658]
    rc2t = c2tcio(c, era, p)
    @test rc2t[1,1] ≈ -0.1810332128307110439 atol=1e-12
    @test rc2t[1,2] ≈ 0.9834769806938470149 atol=1e-12
    @test rc2t[1,3] ≈ 0.6555535638685466874e-4 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834768134135996657 atol=1e-12
    @test rc2t[2,2] ≈ -0.1810332203649448367 atol=1e-12
    @test rc2t[2,3] ≈ 0.5749801116141106528e-3 atol=1e-12
    @test rc2t[3,1] ≈ 0.5773474014081407076e-3 atol=1e-12
    @test rc2t[3,2] ≈ 0.3961832391772658944e-4 atol=1e-12
    @test rc2t[3,3] ≈ 0.9999998325501691969 atol=1e-12
    @test_throws ArgumentError c2tcio(c[1:2,:], era, p)
    @test_throws ArgumentError c2tcio(c, era, p[1:2,:])
end

@testset "c2teqx" begin
    c = [0.9999989440476103608 -0.1332881761240011518e-2 -0.5790767434730085097e-3;
         0.1332858254308954453e-2 0.9999991109044505944 -0.4097782710401555759e-4;
         0.5791308472168153320e-3 0.4020595661593994396e-4 0.9999998314954572365]
    gst = 1.754166138040730516
    p = [0.9999999999999674705 -0.1367174580728847031e-10 0.2550602379999972723e-6;
         0.1414624947957029721e-10 0.9999999999982694954 -0.1860359246998866338e-5;
         -0.2550602379741215275e-6 0.1860359247002413923e-5 0.9999999999982369658]
    rc2t = c2teqx(c, gst, p)
    @test rc2t[1,1] ≈ -0.1810332128528685730 atol=1e-12
    @test rc2t[1,2] ≈ 0.9834769806897685071 atol=1e-12
    @test rc2t[1,3] ≈ 0.6555535639982634449e-4 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834768134095211257 atol=1e-12
    @test rc2t[2,2] ≈ -0.1810332203871023800 atol=1e-12
    @test rc2t[2,3] ≈ 0.5749801116126438962e-3 atol=1e-12
    @test rc2t[3,1] ≈ 0.5773474014081539467e-3 atol=1e-12
    @test rc2t[3,2] ≈ 0.3961832391768640871e-4 atol=1e-12
    @test rc2t[3,3] ≈ 0.9999998325501691969 atol=1e-12
    @test_throws ArgumentError c2teqx(c[1:2,:], gst, p)
    @test_throws ArgumentError c2teqx(c, gst, p[1:2,:])
end

@testset "c2tpe" begin
    tta = 2400000.5
    uta = 2400000.5
    ttb = 53736.0
    utb = 53736.0
    deps =  0.4090789763356509900
    dpsi = -0.9630909107115582393e-5
    xp = 2.55060238e-7
    yp = 1.860359247e-6
    rc2t = c2tpe(tta, ttb, uta, utb, dpsi, deps, xp, yp)
    @test rc2t[1,1] ≈ -0.1813677995763029394 atol=1e-12
    @test rc2t[1,2] ≈ 0.9023482206891683275 atol=1e-12
    @test rc2t[1,3] ≈ -0.3909902938641085751 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834147641476804807 atol=1e-12
    @test rc2t[2,2] ≈ -0.1659883635434995121 atol=1e-12
    @test rc2t[2,3] ≈ 0.7309763898042819705e-1 atol=1e-12
    @test rc2t[3,1] ≈ 0.1059685430673215247e-2 atol=1e-12
    @test rc2t[3,2] ≈ 0.3977631855605078674 atol=1e-12
    @test rc2t[3,3] ≈ 0.9174875068792735362 atol=1e-12
end

@testset "c2txy" begin
    tta = 2400000.5
    uta = 2400000.5
    ttb = 53736.0
    utb = 53736.0
    x = 0.5791308486706011000e-3
    y = 0.4020579816732961219e-4
    xp = 2.55060238e-7
    yp = 1.860359247e-6
    rc2t = c2txy(tta, ttb, uta, utb, x, y, xp, yp)
    @test rc2t[1,1] ≈ -0.1810332128306279253 atol=1e-12
    @test rc2t[1,2] ≈ 0.9834769806938520084 atol=1e-12
    @test rc2t[1,3] ≈ 0.6555551248057665829e-4 atol=1e-12
    @test rc2t[2,1] ≈ -0.9834768134136142314 atol=1e-12
    @test rc2t[2,2] ≈ -0.1810332203649529312 atol=1e-12
    @test rc2t[2,3] ≈ 0.5749800843594139912e-3 atol=1e-12
    @test rc2t[3,1] ≈ 0.5773474028619264494e-3 atol=1e-12
    @test rc2t[3,2] ≈ 0.3961816546911624260e-4 atol=1e-12
    @test rc2t[3,3] ≈ 0.9999998325501746670 atol=1e-12
end

@testset "cal2jd" begin
    dmj0, dmj = cal2jd(2003, 6, 1)
    @test dmj0 ≈ 2400000.5 atol=1e-9
    @test dmj ≈ 52791.0 atol=1e-9
    @test_throws ERFAException cal2jd(-1e9, 6, 1)
    @test_throws ERFAException cal2jd(2021, 13, 1)
    @test_throws ERFAException cal2jd(2021, 6, 32)
end

@testset "cp" begin
    a = randn(3)
    b1 = ERFA._cp(a)
    b2 = copy(a)
    b1[1] = 1.0
    b2[1] = 1.0
    @test a != b1
    @test a != b2
    @test b1 == b2
end

@testset "cpv" begin
    a = [randn(3), randn(3)]
    b1 = ERFA._cpv(a)
    b2 = deepcopy(a)
    b1[1][1] = 1.0
    b2[1][1] = 1.0
    @test a != b1
    @test a != b2
    @test b1 == b2
end

@testset "cr" begin
    a = randn(3, 3)
    b1 = ERFA._cr(a)
    b2 = copy(a)
    b1[1,1] = 1.0
    b2[1,1] = 1.0
    @test a != b1
    @test a != b2
    @test b1 == b2
end


using LinearAlgebra, JuMP, Mosek, MosekTools, SCS, LowRankApprox, Random, Dates
using TSVD, Distributions, StatsBase

include("src/BnB.jl")
include("src/SPCP.jl")


using ROCKS
using Random
using Distributions

Random.seed!(888)
const x = rand(Uniform(-5, 5), 1_000_000)
const logit = -3.0 .+ 0.5 .* x .+ rand(Normal(0, 0.1), length(x))
const prob = @. 1.0 / (1.0 + exp(-logit))
const target = rand(length(x)) .<= prob

kk = kstest(target, prob)

cc = concordance(target, prob)

perf = bcdiag(target, prob)

kp = ksplot(perf)
png(kp, "ksplot.png")

rp = rocplot(perf)
png(rp, "rocplot.png")

bp = biasplot(perf)
png(bp, "biasplot.png")

ap = accuracyplot(perf)
png(ap, "accuracyplot.png")

lp = liftcurve(perf)
png(lp, "liftcurve.png")

clp = cumliftcurve(perf)
png(clp, "cumliftcurve.png")

liftable(perf)

cumliftable(perf)


@with_kw struct FastGrid
    resolution::Float64 = 1.0
    #tight::Bool         = false
end

# This is the main function
function solve(solver::FastGrid, problem::Problem) #original
    result = true
    delta = solver.resolution

    center = problem.input.center
    radius = problem.input.radius[1]

    (W, b) = (problem.network.layers[1].weights, problem.network.layers[1].bias)

    np = pyimport("numpy")
    q, e, p = np.linalg.svd(W)

    input = forward_affine_map(solver, W, b, problem.input)

    lower, upper = low(input), high(input)
    n_hypers_per_dim = BigInt.(max.(ceil.(Int, (upper-lower) / delta), 1))

    k_1 = size(W, 1)
    k_0 = size(W, 2)

    hps = Array{Hyperplane, 1}(undef, k_0-k_1)

    for i in k_1+1:k_0
        hps[i-k_1] = Hyperplane(p[i,:], dot(p[i,:], center))
    end

    S = hps[1]
    if k_0-k_1 > 1
        for i in 2:k_0-k_1
            S = intersection(S, hps[i])
        end
    end

    C = vcat(W, -W)
    d = zeros(2k_1)

    #kb = p[k_1+1:k_0,:] #kernel basis
    #kc = kb * center

    #= Al = zeros(k_0, k_0)
    Au = zeros(k_0, k_0)

    dl = zeros(k_1)
    du = zeros(k_1)

    bl = zeros(k_0)
    bu = zeros(k_0)

    pl = zeros(k_0)
    pu = zeros(k_0) =#

    #count3 = BigInt(0)
    count4 = BigInt(0)
    #println("All: " * string(prod(n_hypers_per_dim)) * " - " * string(prod(n_hypers_per_dim.-2)))

    # preallocate work arrays
    local_lower, local_upper, CI = similar(lower), similar(lower), similar(lower)
    for i in 1:prod(n_hypers_per_dim)
        n = i
        for j in firstindex(CI):lastindex(CI)
            n, CI[j] = fldmod1(n, n_hypers_per_dim[j])
        end
        @. local_lower = lower + delta * (CI - 1)
        @. local_upper = min(local_lower + delta, upper)
        hyper = Hyperrectangle(low = local_lower, high = local_upper)

        d = vcat(local_upper - b, b - local_lower)

        P_i = HPolyhedron(C, d)

        inter = intersection(problem.input, P_i)

        if isempty(inter) == false
            #count3 += 1
            O_i = box_approximation(intersection(P_i, S))
            inner = true
            for j in 1:k_0
                if low(O_i)[j] ≤ low(problem.input)[j]
                    inner = false
                    break
                end

                if high(problem.input)[j] ≤ high(O_i)[j]
                    inner = false
                    break
                end
            end

            if inner == false
                reach = forward_network(solver, problem.network, hyper)
                count4 += 1
                if !issubset(reach, problem.output)
                    result = false
                end
            end
        end
    end

    println("Verified: " * string(count4))

    if result
        return BasicResult(:holds)
    end
    return BasicResult(:violated)
end

function forward_network(solver::FastGrid, nnet::Network, input::Hyperrectangle)
    layers = nnet.layers
    act = layers[1].activation
    reach = Hyperrectangle(low = act.(low(input)), high = act.(high(input)))

    for i in 2:length(layers)
        reach = forward_layer(solver, layers[i], reach)
    end
    return reach
end

# This function is called by forward_network
function forward_layer(solver::FastGrid, L::Layer, input::Hyperrectangle)
    (W, b, act) = (L.weights, L.bias, L.activation)
    center = zeros(size(W, 1))
    gamma  = zeros(size(W, 1))
    for j in 1:size(W, 1)
        node = Node(W[j,:], b[j], act)
        center[j], gamma[j] = forward_node(solver, node, input)
    end
    return Hyperrectangle(center, gamma)
end

function forward_affine_map(solver::FastGrid, W::Matrix, b::Vector, input::Hyperrectangle)
    center = W * input.center + b
    radius = abs.(W) * input.radius
    return Hyperrectangle(center, radius)
end

function forward_node(solver::FastGrid, node::Node, input::Hyperrectangle)
    output    = node.w' * input.center + node.b
    deviation = sum(abs.(node.w) .* input.radius)
    β    = node.act(output)  # TODO expert suggestion for variable name. beta? β? O? x?
    βmax = node.act(output + deviation)
    βmin = node.act(output - deviation)
    return ((βmax + βmin)/2, (βmax - βmin)/2)
    #if solver.tight
        #return ((βmax + βmin)/2, (βmax - βmin)/2)
    #else
        #return (β, max(abs(βmax - β), abs(βmin - β)))
    #end
end

#distance of point to hyperPlane cx = d
function distance(point::Vector, c::Vector, d::Real)
    if length(point) == length(c)
        return abs(sum(point .* c) - d)/sqrt(sum(abs2.(c)))
    else
        error("Dimesion dismatch for point and constraint")
    end
end


# Script to calculate the availabily of solar power over the year
using DataFrames, FeatherFiles, Statistics

const SUN_AREA = 16.0 # area of the solar panels in m²

function calc_average_power(df)
    return SUN_AREA * mean(df.Power_Generation)
end

function calc_peak_power(df)
    return SUN_AREA * maximum(df.Power_Generation)
end

function calc_capacity_factor(df)
    return 100 * calc_average_power(df) / calc_peak_power(df)
end

function calc_rel_availability(df, min_power)
    minpower = df[df.Power_Generation * SUN_AREA .> min_power, :Power_Generation]
    return 100 * size(minpower, 1) / size(df, 1)
end

if !isdefined(@__MODULE__, :df) df=DataFrame(load("data/solaryear.feather")) end

av = calc_average_power(df)
println("Average solar power [W]:                     ", round(av, digits=1))
println("Capacity factor [%]:                          ", round(calc_capacity_factor(df), digits=1))
println("Percentage of solar power availability   [%]: ", round(calc_rel_availability(df, 2.0), digits=1))
println("Percentage of time for more than 50% av. [%]: ", round(calc_rel_availability(df, av/2), digits=1))
println("Percentage of time for more than average [%]: ", round(calc_rel_availability(df, av), digits=1))


## test_show.jl : test show commands

@info "Testing display commands"

@show Matern(1,1)
@show Exponential(1)
@show SquaredExponential(1)
@show AnisotropicExponential([1000 0; 0 1000])
@show CovarianceFunction(2,Matern(1.0,2))
@show Cholesky()
@show Spectral()
@show KarhunenLoeve(10)
@show CirculantEmbedding()
@show GaussianRandomField(CovarianceFunction(1,Exponential(1)),Cholesky(),0:0.1:1)
@show SeparableCovarianceFunction(Exponential(1),Exponential(1))

@info " Done testing display commands"


#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#
# Problem set 5 solutions
# Written by Tyler Ransom
# Commented by Giuseppe Grasso
# Recording available in Notability
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#

using Random
using LinearAlgebra
using Statistics
using Optim
using DataFrames
using DataFramesMeta
using CSV
using HTTP
using GLM
cd("/Users/peppegrass/Documents/GitHub/fall-2020/ProblemSets/PS5-ddc/") # GG: sets working directory
pwd() ## GG: prints working directory
readdir() # GG: equivalent to -ls- to see elements of working directory

# read in function to create state transitions for dynamic model
include("create_grids.jl")

function allwrap()

    #=
    GG: In this problem set, we will explore a simpliﬁed version of the Rust (1987, Econometrica) bus engine
    replacement model. Let’s start by reading in the data.
    =#

# GG: I - PRELIMINARIES: Loading data and reshaping them

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 1: Reshaping the data
    # Reshape the data into “long” panel format, calling your long dataset df_long.
    # I have included code on how to do this in the PS5starter.jl ﬁle that accompanies this problem set.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # load in the data
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdataBeta0.csv"
    df = CSV.read(HTTP.get(url).body)

    # create bus id variable
    df = @transform(df, bus_id = 1:size(df,1))

    #---------------------------------------------------
    # reshape from wide to long (must do this twice be-
    # cause DataFrames.stack() requires doing it one
    # variable at a time)
    #---------------------------------------------------
    # first reshape the decision variable
    dfy = @select(df, :bus_id,:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20,:RouteUsage,:Branded)
    dfy_long = DataFrames.stack(dfy, Not([:bus_id,:RouteUsage,:Branded]))
    rename!(dfy_long, :value => :Y)
    dfy_long = @transform(dfy_long, time = kron(collect([1:20]...),ones(size(df,1))))
    select!(dfy_long, Not(:variable))

    # next reshape the odometer variable
    dfx = @select(df, :bus_id,:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20)
    dfx_long = DataFrames.stack(dfx, Not([:bus_id]))
    rename!(dfx_long, :value => :Odometer)
    dfx_long = @transform(dfx_long, time = kron(collect([1:20]...),ones(size(df,1))))
    select!(dfx_long, Not(:variable))

    # join reshaped df's back together
    df_long = leftjoin(dfy_long, dfx_long, on = [:bus_id,:time])
    sort!(df_long,[:bus_id,:time])

# GG: II - STATIC ESTIMATION

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 2: Estimate a STATIC version of the model
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    #=
    GG: The model we would like to estimate is Harold Zurcher’s decision to run buses in his ﬂeet.
    Zurcher’s FLOW UTILITY OF RUNNING (i.e. not replacing) a bus is

    (1): $u_{1}(x_{1t},b)= θ_{0}+ θ_{1}x_{1t} + θ_{2}b$ [see assignment text]

    where x_{1t} is the mileage on the bus’s odometer (in 10,000s of miles) and b is a dummy variable indicating
    whether the bus is branded (meaning its manufacturer is high-end). The choice set is {0,1} where 0 denotes
    replacing the engine.
    Estimate the θ parameters assuming Zurcher is completely myopic.
    This amounts to estimating a simple binary logit model.
    (Note: you may estimate this any way you wish. I would recommend using the GLM package, but you may also
    use Optim with your own log likelihood function.)
    =#

    θ̂_glm = glm(@formula(Y ~ Odometer + Branded), df_long, Binomial(), LogitLink()) # GG: estimating using GLM package. Pooled binary logit using reshaped long data
    println(θ̂_glm)

#= GG: III - DYNAMIC ESTIMATION
Now I will walk you through how to estimate the DYNAMIC version of this model using BACKWARDS RECURSION.
With discount factor β, the DIFFERENTIATED CONDITIONAL VALUE FUNCTION for running the bus (relative to replacing it) is

(2): $ v_{1t}(x_{t},b) - v_{0t}(x_{t},b) =
        θ_{0} + θ_{1}x_{1t} + θ_{2}b                    [GG: same as static so far, flow utility]
        + β ∫ V_{t+1}(x_{t+1}, b) dF(x_{t+1} | x_{t})     [GG: recursive component]
        $

where V_{t+1} is the VALUE FUNCTION and the integral is over TRANSITIONS in the mileage states x_{t}.

We will APPROXIMATE the INTEGRAL with a SUMMATION, which means that we will specify a DISCRETE MASS FUNCTION for f(x_{t+1}  |x_{t}).
This probability mass function depends on the current odometer reading (x_{1t}),
whether the engine is newly replaced (i.e. d_{t−1}=0),
and on the value of another state variable x_2 which measures the usage intensity of the bus’s route
(i.e. high values of x_2 imply a low usage intensity and vice versa).
We discretize the mileage transitions into 1,250-mile bins (i.e. 0.125 units of x-{1t}).
We specify x_2 as a discrete uniform distribution ranging from 0.25 to 1.25 with 0.01 unit increments.

Formally, we are DISCRETELY (but not discreetly!) APPRXIMATING an EXPONENTIAL DISTRIBUTION:

(3): [see assignment text]

You will not need to program (3); I will provide code for this part. Under this formulation, (2) can be written as:

(4): [see assignment text]

Finally, we can simplify (4) since we know that V_{t+1 = log( ∑_{k} exp(v_{k,t+1}) ) when we assume
that unobserved utility is drawn from a T1EV distribution (as we do here):

(5): [see assignment text]

Estimation of our dynamic model now requires two steps:

#1 SOLVING THE MODEL
First, we need to solve the value functions for a given value of our parameters θ.
The way we do this is by BACKWARDS RECURSION. We know that V_{t+1} = 0 in our ﬁnal period (i.e. when t = T).
Then we work backwards to obtain the future value at every possible state in our model.
This will include many states that do not actually show up in our data.

# 2 ESTIMATING THE MODEL
Second, once we’ve solved the value functions, we use maximum likelihood to estimate the parameters θ.
The log likelihood function in this case is simply:
(6)-(7): [see assignment text]

Now estimate the θ’s assuming that Zurcher discounts the future with discount factor β = 0.9.
I will walk you through speciﬁc steps for how to do this
=#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3a: Read in the data for the dynamic model.
    # This can be found at the same URL as listed at the top of p. 2, but remove the "Beta0" from the CSV ﬁlename.
    # Rather than reshaping the data to “long” format as in question 1, we want to keep the data in “wide” format.
    # Thus, columns :Y1 through :Y20 should be converted to an array labeled Y which has
    # dimension 1000 × 20 where N = 1000 and T = 20. And similarly for columns starting with :Odo and :Xst.
    # Variables :Xst* and :Zst keep track of which discrete bin of the f j ’s the given observation falls into.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # [ Load in the data ]
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdata.csv"
    df = CSV.read(HTTP.get(url).body)
    Y = Matrix(df[:,[:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20]])
    X = Matrix(df[:,[:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20]])
    Z = Vector(df[:,:RouteUsage])
    B = Vector(df[:,:Branded])
    N = size(Y,1)
    T = size(Y,2)
    Xstate = Matrix(df[:,[:Xst1,:Xst2,:Xst3,:Xst4,:Xst5,:Xst6,:Xst7,:Xst8,:Xst9,:Xst10,:Xst11,:Xst12,:Xst13,:Xst14,:Xst15,:Xst16,:Xst17,:Xst18,:Xst19,:Xst20]])
    Zstate = Vector(df[:,:Zst])

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3b: Generate state transition matrices
    # Construct the state transition matrices, which are the f_j’s in (3). To do so, simply run the following code:
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    zval,zbin,xval,xbin,xtran = create_grids() # GG: creates grid points based on points given. See .jl file pf function
    # GG: returns bin numbers and associated values for z (rout usage) and x (odometer reading); and transition values for the X

    #=
    GG: zval and xval are the grids deﬁned at the bottom of p. 2, which respectively correspond to the route usage
    and odometer reading. zbin and xbin are the number of bins in zval and xval, respectively.
    xtran is a (zbin*xbin)×xbin Markov transition matrix that gives the probability of falling into each x_{1,t+1} bin
    given values of x_{1,t} and x_2 , according to the formula in (3).

    Note: A Markov transition matrix is a matrix where each row sums to 1 and moving from e.g. column 1 to column 4
    within a row gives the probability of moving from state 1 to state 4. Check out the Wikipedia page for more information
    =#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3c: Compute the future value terms for all possible states of the model
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # • First, initialize the future value array, which should be a 3-dimensional array of zeros.
    #   The size of the ﬁrst dimension should be the total number of grid points (i.e. the number of rows of xtran).
    #   The second dimension should be 2, which is the possible outcomes of :Branded.
    #   The third dimension should be T+1.
    # • Now write four nested for loops over each of the possible states:
    #   - Loop backwards over t from T+1 to 1
    #   – Loop over the two possible brand states {0,1}
    #   – Loop over the possible permanent route usage states (i.e. from 1 to zbin)
    #   – Loop over the possible odometer states (i.e. from 1 to xbin)
    # • Inside all of the for loops, make the following calculations
    #   – Create an object that marks the row of the transition matrix that we need to be looking at
    #     (based on the loop values of the two gridded state variables).
    #     This will be x + (z-1)*xbin (where x indexes the mileage bin and z indexes the route usage bin),
    #     given how the xtran matrix was constructed in the create grids() function.
    #   – Create the conditional value function for driving the bus (v_{1t}) based on the values of the state
    #     variables in the loop (not the values observed in the data). For example, for the mileage (x_{1t}),
    #     you should plug in xval[x] rather than :Odo. The difﬁcult part of the conditional value function is
    #     the discrete summation over the state transitions. For this, you need to grab the appropriate row
    #     (and all columns) of the xtran matrix, and then take the dot product with that and the all possible x_{1t}
    #     rows of the FV matrix. You should end up with something like xtran[row,:]'*FV[(z-1)*xbin+1:z*xbin,b+1,t+1]
    #     where b indexes the branded dummy and t indexes time periods.
    #   – Now create the conditional value function for replacing the engine (v_{0t}). For this, we repeat the same
    #     process as with v_{1t} except the θ’s are NORMALIZED to be 0. The code for the expected future value is
    #     the same as for v_{1t} with the exception that mileage resets to 0 after replacement, so instead of grabbing
    #     xtran[row,:] we want xtran[1+(z-1)*xbin,:].
    #   – Finally, update the future value array in period t by storing β log( exp(v_{0t}) + exp(v—1t}) ) in the t_th
    #     slice of the 3rd dimension of the array. This will be the new future value term for period t−1.
    #     Remember to set β = 0.9
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    @views @inbounds function likebus(α,β,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval)

        # FV in period T+1 = 0 since this is a finite horizon problem
        FV=zeros(zbin*xbin,2,T+1) # GG: initialize future values; initialize the future value array, which should be a 3-dimensional array of zeros.
        # GG: 2 columns is Branded=1 vs Branded=0

        # First loop: solve the backward recursion problem given the values of α
        # This will give the future value for *all* possible states that we might visit
        # This is why the FV array does not have an individual dimension
        for t=T:-1:1  # GG: Loop backwards over t from T+1 to 1; going backwards over time every period
            for b=0:1   # GG: Loop over the two possible brand states {0,1}; based on branded or not
                for z=1:zbin # GG: Loop over the possible permanent route usage states (i.e. from 1 to zbin); which bin am I in the x2 state variable
                    for x=1:xbin # GG: Loop over the possible odometer states (i.e. from 1 to xbin); which hin am I in the X1t state variable
                        # inputs to FV

                        # GG: object that marks the row of the transition matrix; to figurate the transition probabilities
                        row = x + (z-1)*xbin # which row of xtran should we look at? depends on milage bin x and route usage z

                        # GG: conditional value function for driving the bus (v_{1t})

                        v1  = α[1] + α[2]*xval[x] + α[3]*b + xtran[         row,:]⋅FV[(z-1)*xbin+1:z*xbin,b+1,t+1] # mileage bin is x, route usage z is permanent
                        # GG: flow utility component + summation over state transitions
                        # GG: set of probs xtran[row,:] DOT PRODUCT Future Values (FV)

                        # GG: conditional value function for replacing the engine (v_{0t})
                        v0  =                                xtran[1+(z-1)*xbin,:]⋅FV[(z-1)*xbin+1:z*xbin,b+1,t+1] # the engine got replaced => mileage is 0, so first bin

                        # FV is discounted log sum of the exponentiated conditional value functions #GG: Updating Future Value terms!!!
                        FV[row,b+1,t] = β*log(exp(v1) + exp(v0))
                        #GG: that's where the recursion comes in!
                    end
                end
            end
        end

        #:::::::::::::::::::::::::::::::::::::::::::::::::::
        # Question 3d: Construct the log likelihood using the future value terms from the previous step and only
        # using the observed states in the data. This will entail a for loop over buses and time periods.
        #:::::::::::::::::::::::::::::::::::::::::::::::::::
        # • Initialize the log likelihood value to be 0. (We will iteratively add to it as we loop over observations in the data)
        # • Create a variable that indexes the state transition matrix rows for the case where the bus has been replaced.
        #   This will be the same 1+(z-1)*xbin as in the conditional value function v 0t above. However, we need to plug in :Zst from the data rather than a hypothetical value z.
        # • Create a variable that indexes the state transition matrix rows for the case where the bus has not been replaced.
        #   This will be the same x + (z-1)*xbin as in v1t  above, except we substitute :Xst and :Zst for x and z.
        # • Now create the ﬂow utility component of v_{1t} − v_{0t} using the actual observed data on mileage and branding.
        # • Next, we need to add the appropriate discounted future value to round out our calculation of v_{1t} − v_{0t}.
        #   Here, we can difference the f_j’s as in (5). You should get something like (xtran[row1,:].-xtran[row0,:])'*FV[row0:row0+xbin-1,B[i]+1,t+1]
        # • Finally, create the choice probabilities for choosing each option as written in (7) and then create the
        #   log likelihood according to the summation in (6).

        # Second loop: form the likelihood given the future values implied by the previous α guess
        # Here, we will take the state-specifc FV's calculated in the first loop
        # But we will only use in the likelihood those state value that are actually visited
        like = 0 # GG: Initialize the log likelihood value to be 0.
        for i=1:N
            # GG: vector of indices in the ACTUAL DATA (no replacement)
            row0 = (Zstate[i]-1)*xbin+1 # this is the same argument as the index of xtran in v0 above, but we use the actual Z
            #z2   = (Zstate[i]-1)*xbin+1 # this is the same argument as the 1st part of the 1st index of the FV array above, but using actual Z
            #z3   = z2+xbin-1            # this is the same argument as the 2nd part of the 1st index of the FV array above, but using actual Z
            for t=1:T
                # GG: vector of indices in the ACTUAL DATA (replacement)
                row1  = Xstate[i,t] + (Zstate[i]-1)*xbin                                                                # this is the same as row in the first loop, except we use the actual X and Z
                v1    = α[1] + α[2]*X[i,t] + α[3]*B[i] + (xtran[row1,:].-xtran[row0,:])⋅FV[row0:row0+xbin-1,B[i]+1,t+1] # using the same formula as v1 above, except we use observed values of X and B, and we difference the transitions
                dem   = 1 + exp(v1)
                like -= ( (Y[i,t]==1)*v1 ) - log(dem) # negative of the binary logit likelihood
            end
        end
        return like
    end

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3e: Wrap all of the code you wrote in (c) and (d) into a function and set up the function so that it can be
    # passed to Optim. For example, you will need to return the negative of the log likelihood and you will need to have
    # the ﬁrst argument be the θ vector that we are trying to estimate
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: See row 213

    θ_start = rand(3)
    θ_true  = [2; -.15; 1]
    # how long to evaluate likelihood function once?
    println("Timing (twice) evaluation of the likelihood function")
    @time likebus(θ_start,0.9,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval)
    @time likebus(θ_start,0.9,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval)
    # estimate likelihood function
    θ̂_optim = optimize(a -> likebus(a,0.9,Y,B,N,T,X,Z,Zstate,Xstate,xtran,zbin,xbin,xval), θ_true, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true))
    θ̂_ddc = θ̂_optim.minimizer
    println(θ̂_ddc)

    return nothing
end

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3f: On the same line as the function, prepend the function declaration with the macros so that your code
    # says @views @inbounds function myfun() rather than function myfun(). This will give you more performant code.
    # On my machine, it cut the computation time in half.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: See row 213
    # @inbounds turns off dimensions checkin by Julia; consumes time, if your sure comformability is right you can turn it off to save time
    # @views treats items as immutable, makes things more efficiently
    # GG: You gotta make sure your functions work firt. Then you use these options to make things faster
    # GG: in Matlab it takes about 25mins; in Julia only 5mins; if we used the TwiceDifferentiable object, it'd be even faster

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3g: Wrap all of your code in an empty function as you’ve done with other problem sets
#:::::::::::::::::::::::::::::::::::::::::::::::::::
allwrap()

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3h: Try executing your script to estimate the likelihood function. This took about 4 minutes
# on my machine when I started from the estimates of the static model in Question 2.
#:::::::::::::::::::::::::::::::::::::::::::::::::::

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3i: Pat yourself on the back and grab a beverage of your choice, because that was a lot of work!
#:::::::::::::::::::::::::::::::::::::::::::::::::::


# Monoid / Semigroup
# ==================

"""
    neutral
    neutral(::Type)
    neutral(_default_return_value) = neutral

Neutral element for `⊕`, also called "identity element". `neutral` is a function which can give you
the neutral element for a concrete type, or alternatively you can use it as a singleton value which combines with everything.

By default `neutral(type)` will return the generic `neutral` singleton. You can override it for your specific type to have a more specific neutral value.

We decided for name `neutral` according to https://en.wikipedia.org/wiki/Identity_element. Alternatives seem inappropriate
  - "identity" is already taken
  - "identity_element" seems to long
  - "I" is too ambiguous
  - "unit" seems ambiguous with physical units

Following Laws should hold
--------------------------

Left Identity

      ⊕(neutral(T), t::T) == t

Right Identity

      ⊕(t::T, neutral(T)) == t
"""
function neutral end
neutral(T::Type) = neutral  # we use the singleton type itself as the generic neutral value
neutral(a) = neutral(typeof(a))

"""
    combine(::T, ::T)::T  # overload this
    ⊕(::T, ::T)::T  # alias \\oplus
    combine(a, b, c, d, ...)  # using combine(a, b) internally

Associcative combinator operator.

The symbol `⊕` (\\oplus) following http://hackage.haskell.org/package/base-unicode-symbols-0.2.3/docs/Data-Monoid-Unicode.html

# Following Laws should hold

Associativity

    ⊕(::T, ⊕(::T, ::T)) == ⊕(⊕(::T, ::T), ::T)
"""
function combine end
const ⊕ = combine

combine(a, b, c, more...) = foldl(⊕, more, init=(a⊕b)⊕c)

# supporting `neutral` as generic neutral value
combine(::typeof(neutral), b) = b
combine(a, ::typeof(neutral)) = a
combine(::typeof(neutral), ::typeof(neutral)) = neutral

for reduce ∈ [:reduce, :foldl, :foldr]
  reduce_monoid = Symbol(reduce, "_monoid")
  _reduce_monoid = Symbol("_", reduce_monoid)

  @eval begin
    """
        $($reduce_monoid)(itr; init=TypeClasses.neutral)

    Combines all elements of `itr` using the initial element `init` if given and `TypeClasses.combine`.
    """
    function $reduce_monoid(itr; init=neutral)
      # taken from Base._foldl_impl

      # Unroll the while loop once to hopefully infer the element type at compile time
      y = iterate_named(itr)
      isnothing(y) && return init
      v = combine(init, y.value)
      while true
          y = iterate_named(itr, y.state)
          isnothing(y) && break
          v = combine(v, y.value)
      end
      return v
    end
  end
end

# Alternative
# ===========

"""
    orelse(a, b)  # overload this
    ⊘(a, b)  # alias \\oslash
    orelse(a, b, c, d, ...)  # using orelse(a, b) internally

Implements an alternative logic, like having two options a and b, taking the first valid one.
We decided for "orelse" instead of "alternatives" to highlight the intrinsic asymmetry in choosing.

The operator ⊘ (\\oslash) is choosen to have an infix operator which is similar to \\oplus, however clearly distinguishable, asymmetric, and somehow capturing a choice semantics.
The slash actually is used to indicate choice (at least in some languages, like German), and luckily \\oslash exists (and is not called \\odiv).
"""
function orelse end
const ⊘ = orelse

orelse(a, b, c, more...) = foldl(⊘, more, init=(a⊘b)⊘c)


# Right division
@inline /(tvec::TransposedVector, mat::AbstractMatrix) = transpose(transpose(mat) \ transpose(tvec))

# Base only has specializations of A_rdiv_Bc and A_rdiv_Bt
@inline A_rdiv_Bt(tvec::TransposedVector, mat::AbstractMatrix) = transpose(mat \ transpose(tvec))

@inline A_rdiv_Bc(tvec::TransposedVector, mat::AbstractMatrix) = ctranspose(mat  \ ctranspose(tvec))

# These others fall back to applying the transpose, which is perfect for us:
#   Ac_rdiv_B
#   Ac_rdiv_Bc
#   At_rdiv_B
#   At_rdiv_Bt


x0 = ones(6)
c(x) = [sum(x)]
G(x) = [x[1];x[3]]
H(x) = [x[2];x[4]]
admpcc = ADMPCCModel(rosenbrock,  x0,
                     lvar = fill(-10.0,size(x0)), uvar = fill(10.0,size(x0)),
                     y0 = [0.0], c = c, lcon = [-Inf], ucon = [6.],
                     G = G, H = H, lccG = zeros(2), lccH = zeros(2))

@test obj(admpcc, admpcc.meta.x0)  == 0.0
@test grad(admpcc, admpcc.meta.x0) == zeros(6)
@test issparse(hess(admpcc, admpcc.meta.x0))
@test issparse(jac(admpcc, admpcc.meta.x0))
@test cons(admpcc, admpcc.meta.x0) == [6.0, 1.0, 1.0, 1.0, 1.0]
@test viol(admpcc, admpcc.meta.x0) == [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]

@test hessG(admpcc, admpcc.meta.x0, zeros(2)) == zeros(6,6)
@test hGprod(admpcc, admpcc.meta.x0, zeros(2), ones(6)) == zeros(6)
@test issparse(jacG(admpcc, admpcc.meta.x0))
@test consG(admpcc, admpcc.meta.x0) == ones(2)

@test hessH(admpcc, admpcc.meta.x0, zeros(2)) == zeros(6,6)
@test hHprod(admpcc, admpcc.meta.x0, zeros(2), ones(6)) == zeros(6)
@test sum(jacH(admpcc, admpcc.meta.x0)) == 2
@test consH(admpcc, admpcc.meta.x0) == ones(2)

@test jprod(admpcc, admpcc.meta.x0, zeros(6)) == zeros(5)
@test jtprod(admpcc, admpcc.meta.x0, zeros(5)) == zeros(6)

@test size(hess(admpcc, admpcc.meta.x0)) == (6,6)
@test hess(admpcc, admpcc.meta.x0, zeros(5), obj_weight = 0.0) == zeros(6,6)
@test hprod(admpcc, admpcc.meta.x0, zeros(6)) == zeros(6)
@test hprod(admpcc, admpcc.meta.x0, zeros(5), ones(6), obj_weight = 0.0) == zeros(6)

@test hessH_structure(admpcc) == hessG_structure(admpcc)
@test length(hessH_coord(admpcc, admpcc.meta.x0, zeros(2))) == admpcc.meta.nvar * (admpcc.meta.nvar + 1) / 2
@test length(hessG_coord(admpcc, admpcc.meta.x0, zeros(2))) == admpcc.meta.nvar * (admpcc.meta.nvar + 1) / 2


# Lebesgue measure

export Lebesgue
struct Lebesgue{X} <: AbstractMeasure{X} end

Lebesgue(X) = Lebesgue{X}()

baseMeasure(μ::Lebesgue{X}) where {X} = μ


import Base: convert
import SymEngine: convert, free_symbols

"""
   is_linearcombination(L::Basic)

Return whether the expression `L` is a linear combination of its symbols.

### Input

- `L` -- expression

### Output

`true` if `L` is a linear combination or false otherwise.

### Examples

```jldoctest
julia> using SpaceExParser: is_linearcombination

julia> is_linearcombination(:(2*x1 - 4))
true

julia> is_linearcombination(:(6.1 - 5.3*f - 0.1*g))
true

julia> is_linearcombination(:(2*x1^2))
false

julia> is_linearcombination(:(x1^2 - 4*x2 + x3 + 2))
false
```
"""
function is_linearcombination(L::Basic)
    all(isempty.(free_symbols.(diff.(L, free_symbols(L)))))
end

is_linearcombination(L::Expr) = is_linearcombination(convert(Basic, L))

"""
    is_halfspace(expr::Expr)

Return wheter the given expression corresponds to a halfspace.

### Input

- `expr` -- a symbolic expression

### Output

`true` if `expr` corresponds to a halfspace or `false` otherwise.

### Examples

```jldoctest
julia> using SpaceExParser: is_halfspace

julia> all(is_halfspace.([:(x1 <= 0), :(x1 < 0), :(x1 > 0), :(x1 >= 0)]))
true

julia> is_halfspace(:(2*x1 <= 4))
true

julia> is_halfspace(:(6.1 <= 5.3*f - 0.1*g))
true

julia> is_halfspace(:(2*x1^2 <= 4))
false

julia> is_halfspace(:(x1^2 > 4*x2 - x3))
false

julia> is_halfspace(:(x1 > 4*x2 - x3))
true
```
"""
function is_halfspace(expr::Expr)::Bool

    # check that there are three arguments
    # these are the comparison symbol, the left hand side and the right hand side
    if (length(expr.args) != 3) || !(expr.head == :call)
        return false
    end

    # check that this is an inequality
    if !(expr.args[1] in [:(<=), :(<), :(>=), :(>)])
        return false
    end

    # convert to symengine expressions
    lhs, rhs = convert(Basic, expr.args[2]), convert(Basic, expr.args[3])

    # check if the expression defines a halfspace
    return is_linearcombination(lhs) && is_linearcombination(rhs)
end

"""
    is_hyperplane(expr::Expr)

Return wheter the given expression corresponds to a hyperplane.

### Input

- `expr` -- a symbolic expression

### Output

`true` if `expr` corresponds to a halfspace or `false` otherwise.

### Examples

```jldoctest
julia> using SpaceExParser: is_hyperplane

julia> is_hyperplane(:(x1 = 0))
true

julia> is_hyperplane(:(2*x1 = 4))
true

julia> is_hyperplane(:(6.1 = 5.3*f - 0.1*g))
true

julia> is_hyperplane(:(2*x1^2 = 4))
false

julia> is_hyperplane(:(x1^2 = 4*x2 - x3))
false

julia> is_hyperplane(:(x1 = 4*x2 - x3))
true
```
"""
function is_hyperplane(expr::Expr)::Bool

    # check that there are three arguments
    # these are the comparison symbol, the left hand side and the right hand side
    if (length(expr.args) != 2) || !(expr.head == :(=))
        return false
    end

    # convert to symengine expressions
    lhs = convert(Basic, expr.args[1])

    if :args in fieldnames(typeof(expr.args[2]))
        # treats the 4 in :(2*x1 = 4)
        rhs = convert(Basic, expr.args[2].args[2])
    else
        rhs = convert(Basic, expr.args[2])
    end

    # check if the expression defines a hyperplane
    return is_linearcombination(lhs) && is_linearcombination(rhs)
end

"""
    convert(::Type{HalfSpace{N}}, expr::Expr; vars=nothing) where {N}

Return a `LazySet.HalfSpace` given a symbolic expression that represents a halfspace.

### Input

- `expr` -- a symbolic expression
- `vars` -- (optional, default: `nothing`): set of variables with respect to which
            the gradient is taken; if nothing, it takes the free symbols in the given expression

### Output

A `HalfSpace`, in the form `ax <= b`.

### Examples

```jldoctest convert_halfspace
julia> using LazySets: HalfSpace

julia> convert(HalfSpace, :(x1 <= -0.03))
HalfSpace{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(HalfSpace, :(x1 < -0.03))
HalfSpace{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(HalfSpace, :(x1 > -0.03))
HalfSpace{Float64, Vector{Float64}}([-1.0], 0.03)

julia> convert(HalfSpace, :(x1 >= -0.03))
HalfSpace{Float64, Vector{Float64}}([-1.0], 0.03)

julia> convert(HalfSpace, :(x1 + x2 <= 2*x4 + 6))
HalfSpace{Float64, Vector{Float64}}([1.0, 1.0, -2.0], 6.0)
```

You can also specify the set of "ambient" variables, even if not
all of them appear:

```jldoctest convert_halfspace
julia> using SymEngine: Basic

julia> convert(HalfSpace, :(x1 + x2 <= 2*x4 + 6), vars=Basic[:x1, :x2, :x3, :x4])
HalfSpace{Float64, Vector{Float64}}([1.0, 1.0, 0.0, -2.0], 6.0)
```
"""
function convert(::Type{HalfSpace{N}}, expr::Expr; vars::Vector{Basic}=Basic[]) where {N}

    @assert is_halfspace(expr) "the expression :(expr) does not correspond to a halfspace"

    # check sense of the inequality, assuming < or <= by default
    got_geq = expr.args[1] in [:(>=), :(>)]

    # get sides of the inequality
    lhs, rhs = convert(Basic, expr.args[2]), convert(Basic, expr.args[3])

    # a1 x1 + ... + an xn + K [cmp] 0 for cmp in <, <=, >, >=
    eq = lhs - rhs
    if isempty(vars)
        vars = free_symbols(eq)
    end
    K = subs(eq, [vi => zero(N) for vi in vars]...)
    a = convert(Basic, eq - K)

    # convert to numeric types
    K = convert(N, K)
    a = convert(Vector{N}, diff.(a, vars))

    if got_geq
        return HalfSpace(-a, K)
    else
        return HalfSpace(a, -K)
    end
end

# type-less default halfspace conversion
convert(::Type{HalfSpace}, expr::Expr; vars::Vector{Basic}=Basic[]) = convert(HalfSpace{Float64}, expr; vars=vars)

"""
    convert(::Type{Hyperplane{N}}, expr::Expr; vars=nothing) where {N}

Return a `LazySet.Hyperplane` given a symbolic expression that represents a hyperplane.

### Input

- `expr` -- a symbolic expression
- `vars` -- (optional, default: `nothing`): set of variables with respect to which
            the gradient is taken; if nothing, it takes the free symbols in the given expression

### Output

A `Hyperplane`, in the form `ax = b`.

### Examples

```jldoctest convert_hyperplane
julia> using LazySets: Hyperplane

julia> convert(Hyperplane, :(x1 = -0.03))
Hyperplane{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(Hyperplane, :(x1 + 0.03 = 0))
Hyperplane{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(Hyperplane, :(x1 + x2 = 2*x4 + 6))
Hyperplane{Float64, Vector{Float64}}([1.0, 1.0, -2.0], 6.0)
```

You can also specify the set of "ambient" variables in the hyperplane, even if not
all of them appear:

```jldoctest convert_hyperplane
julia> using SymEngine: Basic

julia> convert(Hyperplane, :(x1 + x2 = 2*x4 + 6), vars=Basic[:x1, :x2, :x3, :x4])
Hyperplane{Float64, Vector{Float64}}([1.0, 1.0, 0.0, -2.0], 6.0)
```
"""
function convert(::Type{Hyperplane{N}}, expr::Expr; vars::Vector{Basic}=Basic[]) where {N}

    @assert is_hyperplane(expr) "the expression :(expr) does not correspond to a Hyperplane"

    # get sides of the inequality
    lhs = convert(Basic, expr.args[1])

    # treats the 4 in :(2*x1 = 4)
    rhs = :args in fieldnames(typeof(expr.args[2])) ? convert(Basic, expr.args[2].args[2]) : convert(Basic, expr.args[2])

    # a1 x1 + ... + an xn + K = 0
    eq = lhs - rhs
    if isempty(vars)
        vars = free_symbols(eq)
    end
    K = subs(eq, [vi => zero(N) for vi in vars]...)
    a = convert(Basic, eq - K)

    # convert to numeric types
    K = convert(N, K)
    a = convert(Vector{N}, diff.(a, vars))

    return Hyperplane(a, -K)
end

# type-less default Hyperplane conversion
convert(::Type{Hyperplane}, expr::Expr; vars::Vector{Basic}=Basic[]) = convert(Hyperplane{Float64}, expr; vars=vars)

"""
    free_symbols(expr::Expr, set_type::Type{LazySet})

Return the free symbols in an expression that represents a given set type.

### Input

- `expr` -- symbolic expression

### Output

A list of symbols, in the form of SymEngine `Basic` objects.

### Examples

```jldoctest
julia> using SpaceExParser: free_symbols

julia> using LazySets: HalfSpace

julia> free_symbols(:(x1 <= -0.03), HalfSpace)
1-element Vector{SymEngine.Basic}:
 x1

julia> free_symbols(:(x1 + x2 <= 2*x4 + 6), HalfSpace)
3-element Vector{SymEngine.Basic}:
 x2
 x1
 x4
```
"""
function free_symbols(expr::Expr, set_type::Type{<:HalfSpace})
    # get sides of the inequality
    lhs, rhs = convert(Basic, expr.args[2]), convert(Basic, expr.args[3])
    return free_symbols(lhs - rhs)
end

function free_symbols(expr::Expr, set_type::Type{<:Hyperplane})
    # get sides of the inequality
    lhs = convert(Basic, expr.args[1])

    # treats the 4 in :(2*x1 = 4)
    rhs = :args in fieldnames(typeof(expr.args[2])) ? convert(Basic, expr.args[2].args[2]) : convert(Basic, expr.args[2])
    return free_symbols(lhs - rhs)
end

function free_symbols(expr::Expr)
    if is_hyperplane(expr)
        return free_symbols(expr, HyperPlane)
    elseif is_halfspace(expr)
        return free_symbols(expr, Halfspace)
    else
        error("the free symbols for the expression $(expr) is not implemented")
    end
end


# ---------------------- #
# Piecewise linear Basis #
# ---------------------- #

struct Lin <: BasisFamily end

struct LinParams{T<:AbstractVector} <: BasisParams
    breaks::T
    evennum::Int
    function LinParams{T}(breaks::T, evennum::Int) where T
        n = length(breaks)  # 28
        !(issorted(breaks)) && error("breaks should be increasing")

        if evennum != 0
            if length(breaks) == 2
                breaks = range(breaks[1], stop=breaks[2], length=evennum)
            else
                if length(breaks) < 2
                    error("breaks must have at least 2 elements")
                end

                if any(abs.(diff(diff(breaks))) .> 5e-15*mean(abs.(breaks)))
                    error("Breaks not evenly spaced")
                end
                evennum = length(breaks)
            end
        end
        new{T}(breaks, evennum)
    end
end

LinParams(breaks::T, evennum::Int=0) where {T<:AbstractVector} = LinParams{T}(breaks, evennum)

# constructor to take a, b, n and form linspace for breaks
LinParams(n::Int, a::Real, b::Real) =
    LinParams(range(a, stop=b, length=n), 0)

## BasisParams interface
# define these methods on the type, the instance version is defined over
# BasisParams
family(::Type{T}) where {T<:LinParams} = Lin
family_name(::Type{T}) where {T<:LinParams} = "Lin"
SparseArrays.issparse(::Type{T}) where {T<:LinParams} = true
function Base.eltype(::Type{LinParams{T}}) where T
    elT = eltype(T)
    elT <: Integer ? Float64 : elT
end

# methods that only make sense for instances
Base.min(p::LinParams) = minimum(p.breaks)
Base.max(p::LinParams) = maximum(p.breaks)
Base.length(p::LinParams) = length(p.breaks)

function Base.show(io::IO, p::LinParams)
    m = string("Piecewise linear interpoland parameters ",
               "from $(p.breaks[1]), $(p.breaks[end])")
    print(io, m)
end

nodes(p::LinParams) = collect(p.breaks)

function derivative_op(p::LinParams, x, order::Int=1)
    breaks, evennum = p.breaks, p.evennum

    newbreaks = collect(breaks)
    Tb = typeof(newbreaks)
    n = length(breaks)
    D = Array{SparseMatrixCSC{basis_eltype(p, x),Int}}(undef, abs(order))

    for i in 1:order
        d = 1.0 ./ diff(newbreaks)
        d = sparse([1:n-1; 1:n-1], [1:n-1; 2:n], [-d; d], n-1, n)
        if i > 1
            D[i] = d*D[i-1]
        else
            D[1] = d
        end
        newbreaks = ((newbreaks[1:end-1]+newbreaks[2:end])/2)::Tb
        n = n-1
    end

    for i in -1:-1:order
        newbreaks = ([dot([3, -1], newbreaks[1:2]);
                      (newbreaks[1:end-1]+newbreaks[2:end]);
                       dot([-1, 3], newbreaks[end-1:end])]/2)::Tb
        d = diff(newbreaks)'
        n = n+1
        d = tril(repeat(d, n, 1), -1)
        if i<-1
            D[-i] = d*D[-i-1]
        else
            D[1] = d
        end
        #adjustment to make value at original left endpoint equal 0
        if evennum > 0
            temp = evalbase(LinParams(newbreaks, length(newbreaks)),
                            breaks[1], 0)*D[-i]
        else
            temp = evalbase(LinParams(newbreaks, 0), breaks[1], 0)*D[-i]
        end
        D[-i] = D[-i]-repeat(temp, length(newbreaks), 1)
    end

    params = LinParams(newbreaks, evennum > 0 ? evennum : 0)

    return D, params

end

function _prep_evalbase(p::LinParams, x::Real)
    n = length(p.breaks)
    if p.evennum != 0
        _ind = fix((x-p.breaks[1])*((n-1)/(p.breaks[end]-p.breaks[1])))
        ind = clamp(_ind+1, 1, n-1)
    else
        ind = lookup(p.breaks, x, 3)
    end
    1, n, [ind]
end

function _prep_evalbase(p::LinParams, x::AbstractArray)
    m = size(x, 1)
    n = length(p.breaks)

    # Determine the maximum index of
    #   the breakpoints that are less than or equal to x,
    #   (if x=b use the index of the next to last breakpoint).
    if p.evennum != 0
        ind = fix((x .- p.breaks[1]).*((n-1)./(p.breaks[end]-p.breaks[1]))) .+ 1
        clamp!(ind, 1, n-1)
    else
        ind = lookup(p.breaks, x, 3)
    end

    return m, n, ind
end

function evalbase(
        ::Type{SparseMatrixCSC}, p::LinParams,
        x::Union{Real,AbstractArray}=nodes(p), order::Int=0
    )
    # 46-49
    if order != 0
        D, params = derivative_op(p, x, order)
        B = evalbase(SparseMatrixCSC, params, x, 0)
        return B*D[end]
    end

    m, n, ind = _prep_evalbase(p, x)
    z = Array{basis_eltype(p, x)}(undef, length(x))
    for i in 1:length(x)
        z[i] = (x[i]-p.breaks[ind[i]])/(p.breaks[ind[i]+1]-p.breaks[ind[i]])
    end

    out = sparse(vcat(1:m, 1:m), vcat(ind, ind .+ 1), vcat(1 .- z, z), m, n)
    return out
end

function evalbase(
        ::Type{SplineSparse}, p::LinParams,
        x::Union{Real,AbstractArray}=nodes(p), order::Int=0
    )
    # 46-49
    if order != 0
        error("derivatives un-supported right now")
    end

    m, n, ind = _prep_evalbase(p, x)

    T = basis_eltype(p, x)
    z = Array{T}(undef, 2*length(x))
    for i in 1:length(x)
        ix = 2i
        z[ix] = (x[i]-p.breaks[ind[i]])/(p.breaks[ind[i]+1]-p.breaks[ind[i]])
        z[ix-1] = 1 - z[ix]
    end

    return SplineSparse{T,Int,1,2}(n, z, ind)
end

function evalbase(
        p::LinParams, x::Union{Real,AbstractArray}=nodes(p),
        order::Int=0
    )
    evalbase(SparseMatrixCSC, p, x, order)
end

function evalbase(p::LinParams, x, order::AbstractArray{Int})
    out = Array{SparseMatrixCSC{basis_eltype(p,x),Int}}(undef, size(order))

    for I in eachindex(order)
        out[I] = evalbase(p, x, order[I])
    end
    return out
end


include("box_test.jl")
include("ball_test.jl")
include("simplex_test.jl")
include("affine_test.jl")


using DSGE
using HDF5
path = dirname(@__FILE__)

# Test in model optimization
custom_settings = Dict{Symbol, Setting}(
    :date_forecast_start  => Setting(:date_forecast_start, quartertodate("2015-Q4")))
m = AnSchorfheide(custom_settings = custom_settings, testing = true)

file = "$path/../reference/optimize.h5"
# For regenerating test file
#=params_test = h5read(file, "params")
data_test = h5read(file, "data") =#

x0 = h5read(file, "params")
data = h5read(file, "data")'
minimizer = h5read(file, "minimizer")
minimum = h5read(file, "minimum")
H_expected = h5read(file, "H")

# See src/estimate/estimate.jl
update!(m, x0)
n_iterations = 3

x0 = Float64[p.value for p in m.parameters]

out, H = optimize!(m, data; iterations=n_iterations)

# Re-generate test file
#=h5open("$path/../reference/optimize.h5", "w") do file
    file["params"] = params_test
    file["data"] = data_test
    file["minimizer"] = out.minimizer
    file["minimum"] = out.minimum
    file["H"] = H
end=#

@testset "Check optimize minimizers are the same" begin
    @test @test_matrix_approx_eq minimizer out.minimizer
    @test minimum ≈ out.minimum atol=5e-7
    @test @test_matrix_approx_eq H_expected H
end

nothing


using DataStructures: Queue, enqueue!, dequeue!, isempty

include("../LoadGraph/LoadGraph.jl")
using .LoadGraph: load_adjacency_list, AdjacencyList

function main(dir_path="../../cleaned_data", dtype::Type=UInt32)
    for file in filter(x->endswith(x, ".txt"), readdir(dir_path))
        if file != "friendster.txt"
        println("=========== ", file, " ===========")
        graph = load_adjacency_list("$dir_path/$file", dtype)

        @time all_cc = all_connected_components(graph)
        println("Number of connected components: ", length(all_cc))
        fraction = maximum(length(cc) for cc in all_cc) / length(graph)
        println("Fraction of nodes in the largest CC: ",
                round(100*fraction, digits=2), "%")

        @time diam_lb = diameter_lb(graph)
        println("Diameter lower bound : ", diam_lb)
        println()
        end
    end
end

function BFS(graph::AdjacencyList{T}, node::T) where T
    fifo = Queue{T}()
    enqueue!(fifo, node)
    marks = Dict{T, Bool}([(node, true)])
    distances = Dict{T, Int}([(node, 0)])

    connected_component = Vector{T}()
    while ! isempty(fifo)
        u = dequeue!(fifo)
        push!(connected_component, u)
        for v in graph[u]
            if ! haskey(marks, v)
                enqueue!(fifo, v)
                marks[v] = true
                distances[v] = distances[u] + 1
            end
        end
    end
    return [connected_component, distances]
end

function all_connected_components(graph::AdjacencyList{T}) where T
    connected_components = Vector{T}[]
    marks = Dict{T, Bool}()
    for key in keys(graph)
        if length(marks) == length(graph)
            break
        elseif ! haskey(marks, key) # We check that the node is not marked yet
            # We mark this new node
            marks[key] = true
            cc, _ = BFS(graph, key)

            # We mark the nodes that are in the connected component returned by BFS
            for node in cc
                marks[node] = true
            end
            push!(connected_components, cc)
        end
    end

    return connected_components
end

"""Computes a lower bound of the diameter of a graph.

The distance is the number of edges between two nodes.
"""
function diameter_lb(graph::AdjacencyList)
    node0 = highest_degree_node(graph)
    node1, _ = furthest_node(graph, node0)
    _, dist = furthest_node(graph, node1)
    return maximum(values(dist))
end

function highest_degree_node(graph::AdjacencyList{T}) where T
    max_deg = 0
    arg_max_degree = collect(Iterators.take(keys(graph), 1))[1]
    for (node, neighbour) in graph
        if length(neighbour) > max_deg
            arg_max_degree = node
        end
    end
    return arg_max_degree
end

function furthest_node(graph::AdjacencyList{T}, node::T) where T
    #get the connected components of node and distances from node
    cc_node, dist = BFS(graph, node)
    furthest = cc_node[end]#get the node that is the furthest from node
    return furthest, dist
end

main()


include("./util.jl")

struct On end
struct Off end
struct Toggle end

Op = Union{On,Off,Toggle}

parse_op(str::AbstractString)::Op =
    if str == "turn on"
        On()::Op
    elseif str == "turn off"
        Off()::Op
    elseif str == "toggle"
        Toggle()::Op
    else
        throw(DomainError("Unknown operation $str"))
    end

struct Instruction{T<:Op}
    op::T
    start::Tuple{Int,Int}
    finish::Tuple{Int,Int}
end

function process!(ins::Instruction{On}, arr::Array{Bool,2})::Nothing
    x0, y0 = ins.start
    x1, y1 = ins.finish
    arr[x0:x1, y0:y1] .= true
    nothing
end

function process!(ins::Instruction{Off}, arr::Array{Bool,2})::Nothing
    x0, y0 = ins.start
    x1, y1 = ins.finish
    arr[x0:x1, y0:y1] .= false
    nothing
end

function process!(ins::Instruction{Toggle}, arr::Array{Bool,2})::Nothing
    x0, y0 = ins.start
    x1, y1 = ins.finish
    arr[x0:x1, y0:y1] = broadcast(!, arr[x0:x1, y0:y1])
    nothing
end

function process!(ins::Instruction{On}, arr::Array{Int,2})::Nothing
    x0, y0 = ins.start
    x1, y1 = ins.finish
    arr[x0:x1, y0:y1] .+= 1
    nothing
end

function process!(ins::Instruction{Off}, arr::Array{Int,2})::Nothing
    x0, y0 = ins.start
    x1, y1 = ins.finish
    arr[x0:x1, y0:y1] .-= 1
    arr[x0:x1, y0:y1] = broadcast(v -> max(v, 0), arr[x0:x1, y0:y1])
    nothing
end

function process!(ins::Instruction{Toggle}, arr::Array{Int,2})::Nothing
    mod_ins = Instruction(On(), ins.start, ins.finish)
    process!(mod_ins, arr)
    process!(mod_ins, arr)
    nothing
end

function Instruction(str::AbstractString)::Instruction
    re =
        r"(?<op>turn on|turn off|toggle) (?<source_x>\d+),(?<source_y>\d+) through (?<target_x>\d+),(?<target_y>\d+)"
    m = match(re, str)
    if isnothing(m)
        throw(DomainError("No match for $str"))
    end
    Instruction(
        parse_op(m[:op]),
        (parse(Int, m[:source_x]), parse(Int, m[:source_y])),
        (parse(Int, m[:target_x]), parse(Int, m[:target_y])),
    )
end

instructions = map(Instruction, split(util.load_input(6), '\n'))

state = zeros(Bool, (1000, 1000))

for instruction in instructions
    process!(instruction, state)
end

println("Part 1: $(sum(state))")

new_state = zeros(Int, (1000, 1000))
for instruction in instructions
    process!(instruction, new_state)
end
println("Part 2: $(sum(new_state))")


# ---
# title: 480. Sliding Window Median
# id: problem480
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Sliding Window
# link: <https://leetcode.com/problems/sliding-window-median/description/>
# hidden: true
# ---
#
# Median is the middle value in an ordered integer list. If the size of the list
# is even, there is no middle value. So the median is the mean of the two middle
# value.
#
# Examples:
#
# `[2,3,4]` , the median is `3`
#
# `[2,3]`, the median is `(2 + 3) / 2 = 2.5`
#
# Given an array _nums_ , there is a sliding window of size _k_ which is moving
# from the very left of the array to the very right. You can only see the _k_
# numbers in the window. Each time the sliding window moves right by one
# position. Your job is to output the median array for each window in the
# original array.
#
# For example,
# Given _nums_ = `[1,3,-1,-3,5,3,6,7]`, and _k_ = 3.
#
#
#
#     Window position                Median
#     ---------------               -----
#     [1  3  -1] -3  5  3  6  7       1
#      1 [3  -1  -3] 5  3  6  7       -1
#      1  3 [-1  -3  5] 3  6  7       -1
#      1  3  -1 [-3  5  3] 6  7       3
#      1  3  -1  -3 [5  3  6] 7       5
#      1  3  -1  -3  5 [3  6  7]      6
#
#
# Therefore, return the median sliding window as `[1,-1,-1,3,5,6]`.
#
# **Note:**
# You may assume `k` is always valid, ie: `k` is always smaller than input
# array's size for non-empty array.
# Answers within `10^-5` of the actual value will be accepted as correct.
#
#
## @lc code=start
using LeetCode

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


#=
## Network Dynamics without NetworkDynamics.jl
=#
#src In this example we want to explore the same problem as in [Integration with `NetworkDynamics.jl`](@ref).
#src But this time without `NetworkDynamics.jl` ...
#=
In this example we model a Kuramoto system on a complex network.
=#
using LightGraphs
using BlockSystems
using ModelingToolkit
using OrdinaryDiffEq
using Plots

#=
Dynamical systems on complex networks are best modelled on directed graphs, because in
general an interaction of node i and j may be asymmetric, e.g. when node i influences
node j but not vice versa.

Symmetries of the system such as (anti-)symmetric coupling functions may justify other
approaches, but those are of limited scope, for a detailed discussion see
[our paper](https://arxiv.org/abs/2012.12696), Section II.D.

In the following each edge is represented by a function:

```
     e₁₂ = f(1,2)
     .--->---.
   (1)       (2)
     ˙---<---˙
     e₂₁ = f(2,1)
```

Each edge function sees the values of the connected nodes, its source and its destination.
Nodes are modeled as functions as well and see all their *incoming* edges.

The goal is to generate IOBlocks for edges and vertices based on a given graph.

While defining the edge block is straightforward, we have to make the vertex blocks a bit
special: since MTK does not support vector inputs yet we need a special IOBlock which
depends on the number of connected edges.
=#
function gen_edge_block(name)
    @parameters t src(t) dst(t) K
    @variables out(t)
    IOBlock([out ~ K*sin(src-dst)], [src, dst], [out], name=Symbol(name))
end

## append subscript, `sybscript(:foo, 2) => :foo₂`
subscript(s, i) = Symbol(s, Char(0x02080 + i))

function gen_vertex_block(n_edges, name)
    @parameters t ω
    @variables ϕ(t)
    D = Differential(t)
    ## the way array variables work changed. This is a hack to retrieve the old behvaviour of
    ## this closely mimics the old @parameters edge[1:n_edges](t)
    edge = Num[]
    for i in 1:n_edges
        symname = subscript(:edge, i)
        append!(edge, @parameters $symname(t))
    end

    IOBlock([D(ϕ) ~ ω + (+)(edge...)],
            [edge...],
            [ϕ],
            name=Symbol(name))
end
nothing #hide

#=
For simplicity our graph will be a simple ring network specified with LightGraphs.jl

=#
N = 8
g = SimpleDiGraph(watts_strogatz(N,2,0)) # ring network
nothing #hide

# First we generate a list of all edge-blocks because they don't depend on the vertices.

edgelist = [(i=i, src=e.src, dst=e.dst, block=gen_edge_block("e_$(e.src)_$(e.dst)"))
            for (i, e) in enumerate(edges(g))]
edge_blocks = [e.block for e in edgelist]
nothing #hide

#=
Now we can generate vertex blocks based on their number of incoming edges.
We will also create the connections
```
e_i_k.out => node.edge₁
e_j_k.out => node.edge₂
...
```
for all edges that go from vertices i or j to vertex k.
=#
vert_blocks = IOBlock[]
connections = Pair[]

for i in vertices(g)
    ## collect the incoming edges for each node
    edges = filter(e -> e.dst == i, edgelist)

    node = gen_vertex_block(length(edges), "v$i")
    push!(vert_blocks, node)

    ## each node has the open inputs edge₁, edge₂, ...
    ## we need to connect the outputs of the edge-blocks to the
    ## inputs of the node like edge_j_to_1.out => node.edge₁
    for (i, edge) in enumerate(edges)
        node_input_i = getproperty(node, subscript(:edge, i))
        push!(connections, edge.block.out => node_input_i)
    end
end

# Once the vertices are generated we can plug the edges' src and dst to the output of the #
# corresponding vertex, in this case its the oscillators angle ϕ
for edge in edgelist
    push!(connections, vert_blocks[edge.src].ϕ => edge.block.src)
    push!(connections, vert_blocks[edge.dst].ϕ => edge.block.dst)
end

#=
We want the `connect_system` to get rid of the algebraic states for the edges. Therefore
we have to provide a list of `outputs` which only contains the outputs of the vertices. By
doing so the edge outputs will become internal `istates` of the `IOSystem` and upon
connection may be reduced.
=#
outputs = [block.ϕ for block in vert_blocks]

network = IOSystem(connections, vcat(vert_blocks, edge_blocks), outputs=outputs)

networkblock = connect_system(network, verbose=false)
nothing #hide

# As the output shows the system has be reduced to just N equations.
# Well now we can generate the functions...
gen = generate_io_function(networkblock,
                           f_states=[v.ϕ for v in vert_blocks],
                           f_params=vcat([v.ω for v in vert_blocks],
                                         [e.K for e in edge_blocks]),
                           warn=false);
nothing #hide

# ... enclose the `f_ip` to get rid of the empty `inputs` field...
odefun(du, u, p, t) = gen.f_ip(du, u, (), p, t)
nothing #hide

# ... set the starting conditions ...
ω = collect(1:N)./N
ω .-= sum(ω)/N
K = [3.0 for i in edge_blocks]
p = (ω..., K...)

x0 = collect(1:N)./N
x0 .-= sum(x0)./N
nothing #hide

#src using BenchmarkTools
#src @btime $odefun($x0, $x0, $p, 0.0)

# ... and solve the system!
tspan = (0., 10.)
prob = ODEProblem(odefun, x0, tspan, p)
sol = solve(prob, Tsit5())
plot(sol, ylabel="ϕ")


struct SymbolicSATProblem
    variables # ::Vector{<:Sym}
    clauses # ::Vector{<:Term}
    index::IndexDict
    p::SATProblem
end

function process(d::IndexDict, var::Sym)
    if !haskey(d.index, var)
        error("$c is not in the problem")
    end

    return [ d.index[var] ]
end

"Output a clause c from a Dict of variables c, as a vector of integers"
function process(d::IndexDict, clause::Term)

    if clause.f == ¬
        return [ -(d.index[clause.arguments[1]]) ]

    elseif clause.f == ∨
        return reduce(vcat, process.(Ref(d), clause.arguments))

    else
        error("Malformed")
    end

end

function process(d::IndexDict, clauses)
    return process.(Ref(d.index), clauses)
end

function SymbolicSATProblem(variables::Vector{<:Sym}, symbolic_clauses)

    d = IndexDict(variables)

    clauses = process.(Ref(d), symbolic_clauses)

    p = SATProblem(length(variables), clauses)

    # @show variables
    # @show symbolic_clauses
    # @show d
    # @show p

    return SymbolicSATProblem(variables, symbolic_clauses, d, p)
end

SymbolicSATProblem(m::Model) = SymbolicSATProblem(boolean_variables(m), clauses(m))

function decode(vars, results::Vector{Int})
    return Dict(v => (i < 0 ? false : true) for (v, i) in zip(vars, results))
end

function solve(p::SymbolicSATProblem)
    status, results = solve(p.p)

    (status == :unsat) && return status, missing

    return status, decode(p.variables, results)

end


struct Exponential <: Distribution{Float64} end

"""
    exponential(rate::Real)

Sample a `Float64` from the exponential distribution with rate parameter `rate`.
"""
const exponential = Exponential()

function Gen.logpdf(::Exponential, x::Real, rate::Real)
    scale = 1/rate
    Distributions.logpdf(Distributions.Exponential(scale), x)
end

function Gen.logpdf_grad(::Exponential, x::Real, rate::Real)
    scale = 1/rate
    x_grad = Distributions.gradlogpdf(Distributions.Exponential(scale), x)
    rate_grad = 1/rate - x
    (x_grad, rate_grad)
end

function Gen.random(::Exponential, rate::Real)
    scale = 1/rate
    rand(Distributions.Exponential(scale))
end

is_discrete(::Exponential) = false

(::Exponential)(rate) = random(Exponential(), rate)

Gen.has_output_grad(::Exponential) = true
Gen.has_argument_grads(::Exponential) = (true,)

export exponential


import VoronoiCells
import Deldir
using GeometryBasics
using Plots

# Number of generators: Deldir does not go all the way
Nsmall = [1000; 2000:2000:30_000]
Nbig = 40_000:10_000:100_000
#= Nsmall = [1000; 2000:2000:10_000] =#
#= Nbig = [10_000, 20_000] =#
N = vcat(Nsmall, Nbig)

VCtime = Vector{Float64}(undef, length(N))
Dtime = similar(VCtime)
fill!(Dtime, NaN)

# Compile methods
Deldir.voronoiarea(rand(8), rand(8))
VoronoiCells.voronoicells(rand(8), rand(8), VoronoiCells.Rectangle(Point2(0, 0), Point2(1, 1)))

# Simulation window
W = [0.0 ; 10.0 ; 0.0 ; 10.0]
rect = VoronoiCells.Rectangle(Point2(0, 0), Point2(10, 10))

for (idx, n) in enumerate(N)
	println(n)

	local x = W[1] .+ W[2]*rand(n)
	local y = W[3] .+ W[4]*rand(n)

    if n <= Nsmall[end]
		Dtime[idx] = @elapsed Deldir.voronoiarea(x, y, W)
	end

    local points = [Point2(x[k], y[k]) for k in 1:n]

    VCtime[idx] = @elapsed begin
        tess = VoronoiCells.voronoicells(points, rect)
        VoronoiCells.voronoiarea(tess)
    end

    GC.gc()
end

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

x = div.(N, 1000)
maxy = 5 * ceil(maximum(filter(!isnan, Dtime)) / 5)

scatter(x, Dtime, label = "Deldir")
scatter!(x, VCtime, label = "VoronoiCells")
plot!(xlabel = "number of points in 1000s", ylabel = "time in seconds", xticks = 0:20:100, yticks = 0:5:maxy)
plot!(tickfont = font(13), legendfont = font(12))

savefig("comparison.png")


using LinearAlgebra
using PyPlot
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCellDG
include("../../../../test/useful_routines.jl")
include("transformation-elasticity-solver.jl")

function angular_position(points)
    cpoints = points[1, :] + im * points[2, :]
    return rad2deg.(angle.(cpoints))
end

function interface_maxnorm_stress_error(polyorder, nelmts, penaltyfactor)
    width = 1.0
    K1, K2 = 247.0, 192.0    # Pa
    mu1, mu2 = 126.0, 87.0   # Pa
    lambda1 = lame_lambda(K1, mu1)
    lambda2 = lame_lambda(K2, mu2)
    stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)

    theta0 = -0.067
    transfstress =
        CutCellDG.plane_strain_transformation_stress(lambda1, mu1, theta0)

    meshwidth = [width, width]
    numqp = required_quadrature_order(polyorder) + 2

    dx = width / nelmts
    penalty = penaltyfactor / dx * 0.5 * (lambda1 + mu1 + lambda2 + mu2)

    basis = TensorProductBasis(2, polyorder)
    interfacecenter = [0.0, 0.0]
    interfaceradius = 0.5
    outerradius = 2.0
    analyticalsolution = AnalyticalSolution(
        interfaceradius,
        outerradius,
        interfacecenter,
        lambda1,
        mu1,
        lambda2,
        mu2,
        theta0,
    )

    mesh, cellquads, facequads, interfacequads = construct_mesh_and_quadratures(
        meshwidth,
        nelmts,
        basis,
        interfacecenter,
        interfaceradius,
        numqp,
    )

    nodaldisplacement = nodal_displacement(
        mesh,
        basis,
        cellquads,
        facequads,
        interfacequads,
        stiffness,
        theta0,
        analyticalsolution,
        penalty,
    )

    refseedpoints, spatialseedpoints, seedcellids =
        CutCellDG.seed_zero_levelset_with_interfacequads(interfacequads, mesh)
    normals = CutCellDG.collect_interface_normals(interfacequads, mesh)
    tangents = CutCellDG.rotate_90(normals)

    spatialpoints = spatialseedpoints[1, :, :]
    referencepoints = refseedpoints
    referencecellids = seedcellids

    productstress = product_stress_at_reference_points(
        nodaldisplacement,
        basis,
        stiffness,
        transfstress,
        theta0,
        referencepoints,
        referencecellids,
        mesh,
    )
    parentstress = parent_stress_at_reference_points(
        nodaldisplacement,
        basis,
        stiffness,
        referencepoints,
        referencecellids,
        mesh,
    )

    productradialtraction =
        CutCellDG.traction_force_at_points(productstress, normals)
    parentradialtraction =
        CutCellDG.traction_force_at_points(parentstress, normals)

    productradialstress = vec(sum(productradialtraction .* normals, dims = 1))
    parentradialstress = vec(sum(parentradialtraction .* normals, dims = 1))

    exactproductradialstress =
        shell_radial_stress(analyticalsolution, interfaceradius)
    exactparentradialstress = core_in_plane_stress(analyticalsolution)

    productradialstressmaxnormerror =
        maximum(abs.(productradialstress .- exactproductradialstress)) /
        abs(exactproductradialstress)
    parentradialstressmaxnormerror =
        maximum(abs.(parentradialstress .- exactparentradialstress)) /
        abs(exactparentradialstress)

    productcircumferentialtraction =
        CutCellDG.traction_force_at_points(productstress, tangents)
    productcircumferentialstress =
        vec(sum(productcircumferentialtraction .* tangents, dims = 1))

    exactproductcircumferentialstress =
        shell_circumferential_stress(analyticalsolution, interfaceradius)

    productcircumferentialstressmaxnormerror =
        maximum(
            abs.(
                productcircumferentialstress .-
                exactproductcircumferentialstress,
            ),
        ) / abs(exactproductcircumferentialstress)

    return [
        parentradialstressmaxnormerror,
        productradialstressmaxnormerror,
        productcircumferentialstressmaxnormerror,
    ]
end

polyorder = 3
penaltyfactor = 1e2
powers = [3, 4, 5]
nelmts = [2^p+1 for p in powers]
dx = 1.0 ./ nelmts

stresserror = [
    interface_maxnorm_stress_error(polyorder, ne, penaltyfactor) for
    ne in nelmts
]

stresserror = vcat([s' for s in stresserror]...)

using CSV, DataFrames

grid = [string(ne)*" x "*string(ne) for ne in nelmts]

df = DataFrame(
    "Grid" => grid,
    "Parent Radial Stress Error" => stresserror[:, 1],
    "Product Radial Stress Error" => stresserror[:, 2],
    "Product Circumferential Stress Error" => stresserror[:,3]
)

folderpath = "examples/traction-continuity-with-transformation/cylindrical-particle/cylindrical-bc/"
filename = "polyorder-"*string(polyorder)*".csv"
CSV.write(folderpath*filename,df)


type SwapForwardMappings end

function coterminal_swap_forward_jacobian(::SwapForwardMappings, cs::CurveState)
  n = cs.numberOfRates
  f = cs.forwardRates
  tau = cs.rateTaus

  a = Vector{Float64}(n)
  for k in eachindex(a)
    a[k] = discount_ratio(cs, k, n+1) - 1.0
  end

  jacobian = zeros(n, n)
  for i = 1:n, j = i:n
    bi = coterminal_swap_annuity(cs, n+1, i)
    bj = coterminal_swap_annuity(cs, n+1, j)
    jacobian[i, j] = tau[j] / coterminal_swap_annuity(cs, j+1, i) + tau[j] / (1.0 + f[j] * tau[j]) * (-a[j] * bi + a[i] * bj) / (bi * bi)
  end

  return jacobian
end

function coterminal_swap_zed_matrix(sfm::SwapForwardMappings, cs::CurveState, displacement::Float64)
  n = cs.numberOfRates
  zMatrix = coterminal_swap_forward_jacobian(sfm, cs)

  f = cs.forwardRates
  sr = coterminal_swap_rates!(cs)

  for i = 1:n, j = i:n
    zMatrix[i, j] *= (f[j] + displacement) / (sr[i] + displacement)
  end

  return zMatrix
end