"""

"""

using DifferentialEquations
using FindSteadyStates
using Plots
using ModelingToolkit
using LabelledArrays
using Parameters
## Build model
function rober(du, u, p, t)
    y1, y2, y3 = u
    k1, k2, k3 = p

    du[1] = -k1*y1 + k3*y2*y3
    du[2] = k1*y1 - k2*y2^2. - k3*y2*y3
    du[3] = k2*y2^2.

end

function rober_jac(J,u,p,t)
    y₁,y₂,y₃ = u
    k₁,k₂,k₃ = p
    J[1,1] = k₁ * -1
    J[2,1] = k₁
    J[3,1] = 0
    J[1,2] = y₃ * k₃
    J[2,2] = y₂ * k₂ * -2 + y₃ * k₃ * -1
    J[3,2] = y₂ * 2 * k₂
    J[1,3] = k₃ * y₂
    J[2,3] = k₃ * y₂ * -1
    J[3,3] = 0
    nothing
  end

## Solve without jacobian
prob = ODEProblem(rober,[1.0,0.0,0.0],(0.0,1e5),(0.04,3e7,1e4))
sol = solve(prob)

## SOlve with jacobian
de = modelingtoolkitize(prob)

jac = eval(ModelingToolkit.generate_jacobian(de)[2])

f = ODEFunction(rober, jac=jac)
prob_jac = ODEProblem(f,[1.0,0.0,0.0],(0.0,1e5),(0.04,3e7,1e4))

@show prob_jac

sol = solve(prob_jac)

@show rober_jac(zeros(3,3), [0.0184367966813563, 7.51141018910664e-8, 0.9815631282045398],(0.04,3.,10.),(0.0,10.))

jac(zeros(3,3),[0.0184367966813563, 7.51141018910664e-8, 0.9815631282045398],(0.0,1e5),(0.04,0.,0.) )
## Plotting
#plot(sol,tspan=(1e-2,1e5),xscale=:log10)

##

using ForwardDiff
using LinearAlgebra

function d(x)
    return -exp.(x)
end

jac = ForwardDiff.jacobian(d, [0.5])

eigvals(jac)

##

@with_kw struct input_func
    func
    p
end

function (self::input_func)(x)
    return self.func(x, x, self.p, (0.0,1.0))
end

function get_jac(de_func, u)
    dj = input_func(de_func.func, de_func.p)
    return ForwardDiff.jacobian(dj, convert(Vector,u))
end

function get_jac2(de_func, u)

    return ForwardDiff.jacobian!(u->de_func(u,u,de_func.p, (0.0,1.0)), u)
end

function bistable_ode!(du, u, p ,t)
	s1, s2 = u
	K1, K2, k1, k2, k3, k4, n1 , n2  = p
	du[1] = k1 / (1 + (s2/K2)^n1) - k3*s1
	du[2] = k2/  (1 + (s1/K1)^n2) - k4*s2
end

p_ = convert(Vector, LVector(K1=1., K2=1., k1=20., k2=20., k3=5., k4=5., n1= 4., n2=4.))

u_1 = convert(Vector, LVector(s1=3., s2=1.) )
u_2 = convert(Vector, LVector(s1=1., s2=3.) )

dss = DEsteady(func=bistable_ode!, u0=u_1, p=p_, method=Tsit5())

sol = solve(dss)

get_jac2(dss, sol.u)

##

function bistable_ode!(du, u, p ,t)
	s1, s2 = u.s1, u.s2
	K1, K2, k1, k2, k3, k4, n1 , n2  = p.K1, p.K2, p.k1, p.k2, p.k3, p.k4, p.n1 , p.n2
	du[1] = k1 / (1 + (s2/K2)^n1) - k3*s1
	du[2] = k2/  (1 + (s1/K1)^n2) - k4*s2
end

p_ = LVector(K1=1., K2=1., k1=20., k2=20., k3=5., k4=5., n1= 4., n2=4.)

u_1 = LVector(s1=3., s2=1.)
u_2 = LVector(s1=1., s2=3.)

##

function bistable_ode!(du, u, p ,t)
	s1, s2 = u
	K1, K2, k1, k2, k3, k4, n1 , n2  = p
	du[1] = k1 / (1 + (s2/K2)^n1) - k3*s1
	du[2] = k2/  (1 + (s1/K1)^n2) - k4*s2
end

p_ = convert(Vector, LVector(K1=1., K2=1., k1=20., k2=20., k3=5., k4=5., n1= 4., n2=4.))

u_1 = convert(Vector, LVector(s1=3., s2=1.) )
u_2 = convert(Vector, LVector(s1=1., s2=3.) )

prob = SteadyStateProblem(bistable_ode!, u_1, p_)
de = modelingtoolkitize(prob)
jac = eval(ModelingToolkit.generate_jacobian(de)[2])
f = ODEFunction(bistable_ode!, jac=jac)
prob_jac_ = SteadyStateProblem(bistable_ode!, u_1 ,p_)

sol = solve(prob)

##
prob = SteadyStateProblem(bistable_ode!, [1.2,1.2], p_)
sol = solve(prob, SSRootfind())
jac_init = zeros(2,2 )
jac(jac_init, sol.u, p_, (1.0,23.5))

@show jac_init
@show eigvals(jac_init)


function all_close(x::Array, y::Array, rtol::Float64=1.e-5, atol::Float64=1.e-8)
    """
    Returns True if two arrays are element-wise equal within a tolerance.

    References
    ==========
    This function, and docstring, are taken directly from numpy.allclose

    See that documentation for more info

    """
    xinf = isinf(x)
    yinf = isinf(y)
    if any(xinf) || any(yinf)
        if !all(xinf .== yinf)
            return false
        end

        if !all(x[xinf] .== y[yinf])
            return false
        end

        x = x[~xinf]
        y = y[~xinf]
    end

    return all(.<=(abs(x - y), atol + rtol * abs(y)))
end

reload("../stable_reg.jl")
using MAT

# load in matlab results
ml = matread("num_stab_approxMATLAB.mat")

# Grab data
X = ml["X"]
Y = ml["Y"]

# compare to matlab answers

# OLS
@assert all_close(OLS(X, Y), ml["b_ols_yes"])
@assert all_close(OLS(X, Y, false), ml["b_ols_no"])

# SVD
@assert all_close(LS_SVD(X, Y), ml["b_ols_yes"])
@assert all_close(LS_SVD(X, Y, false), ml["b_ols_no"])

# LAD-PP
@assert all_close(LAD_PP(X, Y), ml["b_ladpp_yes"])
@assert all_close(LAD_PP(X, Y, false), ml["b_ladpp_no"])

# LAD-DP
@assert all_close(LAD_DP(X, Y), ml["b_laddp_yes"])
@assert all_close(LAD_DP(X, Y, false), ml["b_laddp_no"])

# Regularized least square Tikhonov
@assert all_close(RLS_T(X, Y), ml["b_ols_yes"])

# Regularized SVD Tikhonov
@assert all_close(RLS_TSVD(X, Y), ml["b_ols_yes"])


export bm_count, string4, code_check, random_code

"""
`code_check(c::Int)` checks that `c` is a valid code
(between 0000 and 9999 inclusive). Returns `true`
if valid and throws an error if not.
"""
@inline function code_check(c::Int)::Bool
    (c >= 0) && (c <= 9999)
end

"""
`bad_code_message(c)` creates a string giving an error message
to be used when `c` is invalid.
"""
@inline function bad_code_message(c::Int)::String
    "Code $c is invalid; must be in the range 0000 to 9999."
end

"""
`string4(c::Int)` renders the code number `c` as a four-digit `String`.
```
julia> string4(17)
"0017"
```
"""
function string4(c::Int)::String
    code_check(c)
    cc = reverse(digits(c, pad = 4))
    result = ""
    for i = 1:4
        result *= "$(cc[i])"
    end
    return result
end

"""
`bm_count(c::Int,d::Int)` counts the number of bullets and maggots
when comparing the codes `c` and `d`; this is returned as a
tuple `(b,m)`.
"""
@inline function bm_count(c::Int, d::Int)::Tuple{Int,Int}
    if bm_flag
        return bm_table_fetch(c, d)
    else
        return slow_bm_count(c, d)
    end
end

function slow_bm_count(c::Int, d::Int)::Tuple{Int,Int}

    # break codes into an array of digits
    cc = digits(c, pad = 4)
    dd = digits(d, pad = 4)

    # flags to show if that digit is available to be matched
    c_avail = [true, true, true, true]
    d_avail = [true, true, true, true]

    # counts of bullets & maggots
    b = 0
    m = 0

    # count bullets
    for i = 1:4
        if cc[i] == dd[i]  # we have a bullet; mark digits as used
            b += 1
            c_avail[i] = false
            d_avail[i] = false
        end
    end

    # count maggots
    for i = 1:4
        for j = 1:4
            if i != j && c_avail[i] && d_avail[j] && cc[i] == dd[j]
                m += 1
                c_avail[i] = false
                d_avail[j] = false
            end
        end
    end

    return b, m
end

@inline function random_code()
    mod(rand(Int), 10000)
end

# various first-guess servers

vec2int(v::Vector{Int}) = 1000 * v[1] + 100 * v[2] + 10 * v[3] + v[4]

function rand_digit(exclude::Vector{Int} = Int[])
    d = rand(0:9)
    while in(d, exclude)
        d = rand(0:9)
    end
    return d
end

function guess_distinct()
    v = -ones(Int, 4)
    for j = 1:4
        v[j] = rand_digit(v)
    end
    return vec2int(v)
end

function guess_one_pair()
    v = -ones(Int, 4)
    v[1] = rand_digit()
    v[2] = v[1]
    for i = 3:4
        v[i] = rand_digit(v)
    end
    shuffle!(v)
    return vec2int(v)
end

function guess_two_pairs()
    v = -ones(Int, 4)
    v[1] = rand_digit()
    v[2] = v[1]
    v[3] = rand_digit(v)
    v[4] = v[3]
    shuffle!(v)
    return vec2int(v)
end

guess_any() = rand(0:9999)

export guess_one_pair, guess_two_pairs, guess_any, guess_distinct


using StaticArrays

@testset "StaticArrays Size" begin
    @test Size(static(2:4)) === Size(3)
    @test Size(static(3)) === Size(3)
    @test Size(static(3), static(4)) === Size(3, 4)
    @test static(Size(3)) === (static(3),)
    @test typeof(Size()) === Size{()} # Issue #8
    @test typeof(Size(())) === Size{()} # Issue #8
end

@testset "static length indexing" begin
    @test SOneTo(StaticOneTo(3)) === SOneTo(3)
    @test StaticOneTo(SOneTo(3)) === StaticOneTo(3)
    @test StaticArrays.SUnitRange(LengthUnitRange(3:4)) === StaticArrays.SUnitRange(3,4)
    @test LengthUnitRange(StaticArrays.SUnitRange(3,4)) === static(3):static(4)
    @test staticlength(StaticArrays.SUnitRange(3,4)) === static(3):static(4)

    for x in ([1, 2, 3, 4], @MVector([1, 2, 3, 4]), [1 3; 2 4], [1 2; 3 4]', @MMatrix([1 3; 2 4]), @MMatrix([1 2; 3 4])')
        xi = x[static(1):static(2)]
        v = @MVector([1,2])
        @test xi == v
        @test typeof(xi) == typeof(v)
        xi = @stat x[1:2]
        @test xi == v
        @test typeof(xi) == typeof(v)

        # Test of trystatic(lastindex, x)
        xi = @stat x[1:end-2]
        @test xi == v
        if x isa MArray
            @test xi isa MArray
        else
            @test !(xi isa MArray)
        end
    end

    for x in (@SVector([1, 2, 3, 4]), )
        xi = x[static(1):static(2)]
        v = @SVector([1,2])
        @test xi == v
        @test typeof(xi) == typeof(v)

        f(t) = @stat t[2:end-1]
        @test f(x) === SVector((2,3))
        Test.@inferred f(x)
        Test.@inferred f((1,2,3,4))

        @stat g(t,k) = t[k .+ (0:1)]
        @test g(x, 2) === SVector((2,3))
        Test.@inferred g((1,2,3,4), 2)
    end

    @test SVector(@stat 1:2) === SVector{2}([1,2])
    @test SVector(i^2 for i in @stat 1:2) === SVector{2}([1,4])

    f(k) = SVector(@stat 1:k)
    @test f(static(2)) === SVector{2}([1,2])
    Test.@inferred(f(static(2)))

    g(k) = SVector(i^2 for i in @stat 1:k)
    @test g(static(2)) === SVector{2}([1,4])
    Test.@inferred(g(static(2)))

    @test SVector(i^2 for i in static(1):static(3)) === SVector{3}([1, 4, 9])

    @test CleanNamespace.get_one_to(SVector(1, 2, 3), static(2)) === SVector(1, 2)
    @test CleanNamespace.get_one_to(MVector(1, 2, 3), static(2)) == MVector(1, 2)
    @test CleanNamespace.get_one_to(MVector(1, 2, 3), static(2)) isa MVector
end

@inline function testmul(A::SMatrix, B::SMatrix)
    SArray((Base.@_inline_meta; +(Tuple(A[i,k]*B[k,j] for k in static(axes(A,2)))...))
        for i in static(axes(A,1)), j in static(axes(B,2)))
end

@testset "SMatrix generators" begin
    A = SArray(i + j for i = static(2:2:4), j = static(10:10:30))
    @test A === @SMatrix [12 22 32; 14 24 34]
    B = SArray(10i + j for i = static(1:3), j = static(4:5))
    @test B === @SMatrix [14 15; 24 25; 34 35]
    @test testmul(A, B) == [1784  1850; 1928  2000]
    Test.@inferred testmul(A, B)
end

@testset "SMatrix from view" begin
    A = rand(5,5)
    v = view(A, range(2, length=static(3)), range(3, length=static(2)))
    S = Test.@inferred SMatrix(v)
    @test S isa SMatrix
    @test S == A[2:4,3:4]
end


function plot_model_coef(s::Vector{NamedTuple},  pars::Vector{Symbol};
  mnames=String[], fig="", title="")

  levels = length(s) * (length(pars) + 1)
  colors = [:blue, :red, :green, :darkred, :black, :grey, :darkblue, :cyan, :darkred]

  xmin = 0; xmax = 0.0
  for i in 1:length(s)
    for par in pars
      syms = Symbol.(keys(s[i]))
      if Symbol(par) in syms
        mp = s[i][par].mean
        xmin = min(xmin, s[i][par].lower)
        xmax = max(xmax, s[i][par].upper)
      end
    end
  end

  ylabs = String[]
  for j in 1:length(s)
    for i in 1:length(pars)
      l = length(String(pars[i]))
      str = repeat(" ", 9-l) * String(pars[i])
      append!(ylabs, [str])
    end
    l = length(mnames[j])
    str = mnames[j] * repeat(" ", 9-l)
    append!(ylabs, [str])
  end

  ys = [string(ylabs[i]) for i = 1:length(ylabs)]
  p = plot(xlims=(xmin, xmax), leg=false, framestyle=:grid)
  title!(title)
  yran = range(1, stop=length(ylabs), length=length(ys))
  yticks!(yran, ys)

  line = 0
  for mindx in 1:length(s)
    line += 1
    hline!([line] .+ length(pars), color=:darkgrey, line=(2, :dash))
    for (pindx, par) in enumerate(pars)
      line += 1
      syms = Symbol.(keys(s[mindx]))
      if Symbol(par) in syms
        ypos = (line - 1)
        mp = s[mindx][Symbol(par)].mean
        lower = s[mindx][Symbol(par)].lower
        upper = s[mindx][Symbol(par)].upper
        plot!([lower, upper], [ypos, ypos], leg=false, color=colors[pindx])
        scatter!([mp], [ypos], color=colors[pindx])
      end
    end
  end
  if length(fig) > 0
    savefig(p, fig)
  end
  (s, p)
end

function plot_model_coef(df::DataFrame, pars::Vector{Symbol};
    mname="", fig="", title="")

  if length(mname) > 9
    mname = mname[1:9]
  end

  s = Vector{NamedTuple}(undef, 1)
  for mindx in 1:1
    m, l, u = estimparam(df)
    d = Dict{Symbol, NamedTuple}()
    for (indx, par) in enumerate(names(df))
      d[Symbol(par)] = (mean=m[indx], lower=l[indx], upper=u[indx])
    end
    s[mindx] =   (; d...)
  end

  plot_model_coef(s, pars; mnames=[mname], fig, title)
end

function plot_model_coef(p::Pair{String, DataFrame}, pars::Vector{Symbol};
    fig="", title="")

  mname = p.first
  if length(mname) > 9
    mname = mname[1:9]
  end

  s = Vector{NamedTuple}(undef, 1)
  for mindx in 1:1
    m, l, u = estimparam(p.second)
    d = Dict{Symbol, NamedTuple}()
    for (indx, par) in enumerate(names(p.second))
      d[Symbol(par)] = (mean=m[indx], lower=l[indx], upper=u[indx])
    end
    s[mindx] =   (; d...)
  end

  plot_model_coef(s, pars; mnames=[mname], fig, title)
end

export
  plot_model_coef


"""
        _cols_by_rows(rows_index,cols_index)

Returns a vector of rows where each row contains
a vector of its column indices.
"""
function _cols_by_rows(rows_index,cols_index)
    nrows = maximum(rows_index)
    cols_by_rows = [eltype(rows_index)[] for _ in 1:nrows]
    for (i,j) in zip(rows_index,cols_index)
        push!(cols_by_rows[i],j)
    end
    return cols_by_rows
end

"""
        _rows_by_cols(rows_index,cols_index)

Returns a vector of columns where each column contains
a vector of its row indices.
"""
function _rows_by_cols(rows_index,cols_index)
    return _cols_by_rows(cols_index,rows_index)
end

"""
        matrix2graph(sparse_matrix, [partition_by_rows::Bool=true])

A utility function to generate a graph from input
sparse matrix, columns are represented with vertices
and 2 vertices are connected with an edge only if
the two columns are mutually orthogonal.

Note that the sparsity pattern is defined by structural nonzeroes, ie includes
explicitly stored zeros.
"""
function matrix2graph(sparse_matrix::SparseMatrixCSC{<:Number, Int}, partition_by_rows::Bool=true)

    (rows_index, cols_index, _) = findnz(sparse_matrix)

    ncols = size(sparse_matrix, 2)
    nrows = size(sparse_matrix, 1)

    num_vtx = partition_by_rows ? nrows : ncols

    inner = SimpleGraph(num_vtx)

    if partition_by_rows
        rows_by_cols = _rows_by_cols(rows_index,cols_index)
        @inbounds for (cur_row,cur_col) in zip(rows_index,cols_index)
            if !isempty(rows_by_cols[cur_col])
                for next_row in rows_by_cols[cur_col]
                    if next_row < cur_row
                        add_edge!(inner, cur_row, next_row)
                    end
                end
            end
        end
    else
        cols_by_rows = _cols_by_rows(rows_index,cols_index)
        @inbounds for (cur_row,cur_col) in zip(rows_index,cols_index)
            if !isempty(cols_by_rows[cur_row])
                for next_col in cols_by_rows[cur_row]
                    if next_col < cur_col
                        add_edge!(inner, cur_col, next_col)
                    end
                end
            end
        end
    end
    return VSafeGraph(inner)
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 7cbf25aa-3def-11eb-0c74-af66e15f45ca
s_ex = """mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
mem[8] = 11
mem[7] = 101
mem[8] = 0"""

# ╔═╡ 70c2c378-3e01-11eb-34c7-d17201765158
lines = readlines("14.txt")

# ╔═╡ 8fe5aaf2-3def-11eb-2708-d395cd38ffde
lines_ex = split(s_ex, "\n")

# ╔═╡ 0a9f638a-3dfc-11eb-3c1a-d3d4f9eab152
function maskedvalue(mask, value)
	bin_arr = lpad(string(value, base=2), 36, "0")
	masked = join([m == 'X' ? v : m for (m, v) in zip(mask, bin_arr)])
	parse(Int64, masked, base=2)
end

# ╔═╡ 40bb7a5e-3e05-11eb-046e-d9cef0009c51
function flipidx(binarr, idx, value="1")
	binarr[1:idx-1] * value * binarr[idx+1:end]
end

# ╔═╡ 7f768d70-3e04-11eb-0a0c-91bb5677385a
function maskedaddresses(mask, value)
	N = length(mask)
	bin_arr = lpad(string(value, base=2), N, "0")
	floating_bits = [i for i in 1:N if mask[i] == 'X']
	masked = join([m == '0' ? v : m for (m, v) in zip(mask, bin_arr)])
	addrs = [masked]
	for idx in floating_bits
		new_addrs = []
		for addr in addrs
			@assert addr[idx] == 'X'
			push!(new_addrs, flipidx(addr, idx, "0"))
			push!(new_addrs, flipidx(addr, idx, "1"))
		end
		addrs = new_addrs
	end
	# floating_bits, masked, addrs
	[parse(Int64, s, base=2) for s in addrs]
end

# ╔═╡ 9d343d2c-3def-11eb-155e-6ff7501e4fb6
function parseline(line, mask, memory::Dict)
	m = match(r"mask = ([X01]+)", line)
	if m !== nothing
		mask = m.captures[1]
		return mask, memory
	end
	m = match(r"mem\[([0-9]+)\] = ([0-9]+)", line)
	if m !== nothing
		mem_idx = parse(Int64, m.captures[1])
		value = parse(Int64, m.captures[2])
		memory[mem_idx] = maskedvalue(mask, value)
		return mask, memory
	end
	println("Unparsed `$(line)`")
end

# ╔═╡ fe31a3d8-3e00-11eb-3b37-cdc824855e51
begin
	mask = join("X" for i=1:36)
	memory = Dict()
	for line in lines
		mask, memory = parseline(line, mask, memory)
	end
	sum(values(memory))
end

# ╔═╡ e665342a-3e09-11eb-2e93-fb0813ae24af
function parseline2(line, mask, memory::Dict)
	m = match(r"mask = ([X01]+)", line)
	if m !== nothing
		mask = m.captures[1]
		return mask, memory
	end
	m = match(r"mem\[([0-9]+)\] = ([0-9]+)", line)
	if m !== nothing
		mem_idx = parse(Int64, m.captures[1])
		value = parse(Int64, m.captures[2])
		for addr in maskedaddresses(mask, mem_idx)
			memory[addr] = value
		end
		return mask, memory
	end
	println("Unparsed `$(line)`")
end

# ╔═╡ c0ee620a-3e09-11eb-2ded-7d1ffe2c0243
begin
	mask2 = join("X" for i=1:36)
	memory2 = Dict()
	for line in lines
		mask2, memory2 = parseline2(line, mask2, memory2)
	end
	sum(values(memory2))
end

# ╔═╡ Cell order:
# ╠═7cbf25aa-3def-11eb-0c74-af66e15f45ca
# ╠═70c2c378-3e01-11eb-34c7-d17201765158
# ╠═8fe5aaf2-3def-11eb-2708-d395cd38ffde
# ╠═0a9f638a-3dfc-11eb-3c1a-d3d4f9eab152
# ╠═40bb7a5e-3e05-11eb-046e-d9cef0009c51
# ╠═7f768d70-3e04-11eb-0a0c-91bb5677385a
# ╠═9d343d2c-3def-11eb-155e-6ff7501e4fb6
# ╠═fe31a3d8-3e00-11eb-3b37-cdc824855e51
# ╠═e665342a-3e09-11eb-2e93-fb0813ae24af
# ╠═c0ee620a-3e09-11eb-2ded-7d1ffe2c0243


import DataFrames
import GLM
import StatsModels

function get_unique_symbol_name(
        v::AbstractVector{<:Symbol},
        new_symbol::Symbol = :x;
        append_string::String = "1"
        )::Symbol
    old_strings_set::Set{String} = Set(strip.(string.(v)))
    new_string::String = strip.(string.(new_symbol))
    _append_string::String = strip(append_string)
    while new_string in old_strings_set
        new_string = string(new_string, _append_string)
    end
    new_symbol::Symbol = Symbol(new_string)
    return new_symbol
end

function columns_are_linearly_independent(
        df::DataFrames.DataFrame,
        columns::AbstractVector{<:Symbol} = Symbol[],
        )::Bool
    df_copy = deepcopy(df)
    if length(columns) == 0
        columns_copy = deepcopy(DataFrames.names(df_copy))
    else
        columns_copy = deepcopy(columns)
    end
    temporary_dependent_variable::Symbol = get_unique_symbol_name(
        columns_copy,
        :y,
        )
    formula = generate_formula(
        temporary_dependent_variable,
        columns_copy;
        intercept = false,
        interactions = 1,
        )
    df_copy[temporary_dependent_variable] = randn(size(df_copy, 1))
    try_lm_result::Bool = try
        lm = GLM.lm(formula, df_copy, false)
        true
    catch ex
        @debug("caught exception: ", exception=ex,)
        false
    end
    cols_result = length(linearly_dependent_columns(df, columns)) == 0
    final_result::Bool = try_lm_result && cols_result
    return final_result
end

function linearly_dependent_columns(
        df::DataFrames.DataFrame,
        columns::AbstractVector{<:Symbol} = Symbol[])::Vector{Symbol}
    df_copy = deepcopy(df)
    if length(columns) == 0
        columns_copy = deepcopy(DataFrames.names(df_copy))
    else
        columns_copy = deepcopy(columns)
    end
    temporary_dependent_variable::Symbol = get_unique_symbol_name(
        columns_copy,
        :y,
        )
    formula = generate_formula(
        temporary_dependent_variable,
        columns_copy;
        intercept = false,
        interactions = 1,
        )
    df_copy[temporary_dependent_variable] = randn(size(df_copy, 1))
    lm = GLM.lm(formula, df_copy, true)
    coeftable = StatsModels.coeftable(lm)
    result_strings::Vector{String} = sort(
        unique(strip.(coeftable.rownms[isnan.(coeftable.cols[4])]))
        )
    result_symbols::Vector{Symbol} = Symbol.(result_strings)
    return result_symbols
end


# # [Custom objective terms](@id custom_terms)
#
# ProximalAlgorithms relies on the first-order primitives defined in [ProximalCore](https://github.com/JuliaFirstOrder/ProximalCore.jl).
# While a rich library of function types, implementing such primitives, is provided by [ProximalOperators](https://github.com/JuliaFirstOrder/ProximalOperators.jl),
# one may need to formulate problems using custom objective terms.
# When that is the case, one only needs to implement the right first-order primitive,
# ``\nabla f`` or ``\operatorname{prox}_{\gamma f}`` or both, for algorithms to be able
# to work with ``f``.
#
# Defining the proximal mapping for a custom function type requires adding a method for [`ProximalCore.prox!`](@ref).
#
# To compute gradients, ProximalAlgorithms provides a fallback definition for [`ProximalCore.gradient!`](@ref),
# relying on [Zygote](https://github.com/FluxML/Zygote.jl) to use automatic differentiation.
# Therefore, you can provide any (differentiable) Julia function wherever gradients need to be taken,
# and everything will work out of the box.
#
# If however one would like to provide their own gradient implementation (e.g. for efficiency reasons),
# they can simply implement a method for [`ProximalCore.gradient!`](@ref).
#
# ```@docs
# ProximalCore.prox
# ProximalCore.prox!
# ProximalCore.gradient
# ProximalCore.gradient!
# ```
#
# ## Example: constrained Rosenbrock
#
# Let's try to minimize the celebrated Rosenbrock function, but constrained to the unit norm ball. The cost function is

rosenbrock2D(x) = 100 * (x[2] - x[1]^2)^2 + (1 - x[1])^2

# To enforce the constraint, we define the indicator of the unit ball, together with its proximal mapping:
# this is simply projection onto the unit norm ball, so it is sufficient to normalize any given point that lies
# outside of the set.

using LinearAlgebra
using ProximalCore

struct IndUnitBall end

(::IndUnitBall)(x) = norm(x) > 1 ? eltype(x)(Inf) : eltype(x)(0)

function ProximalCore.prox!(y, ::IndUnitBall, x, gamma)
    if norm(x) > 1
        y .= x ./ norm(x)
    else
        y .= x
    end
    return zero(eltype(x))
end

# We can now minimize the function, for which we will use `PANOC`, which is a Newton-type method:

using ProximalAlgorithms

panoc = ProximalAlgorithms.PANOC()
solution, iterations = panoc(x0=-ones(2), f=rosenbrock2D, g=IndUnitBall())

# Plotting the solution against the cost function contour and constraint, gives an idea of its correctness.

using Plots

contour(-2:0.1:2, -2:0.1:2, (x,y) -> rosenbrock2D([x, y]), fill=true, framestyle=:none, background=nothing)
plot!(Shape(cos.(0:0.01:2*pi), sin.(0:0.01:2*pi)), opacity=.5, label="feasible set")
scatter!([solution[1]], [solution[2]], color=:red, markershape=:star5, label="computed solution")

# ## Example: counting operations
#
# It is often interesting to measure how many operations (gradient- or prox-evaluation) an algorithm is taking.
# In fact, in algorithms involving backtracking or some other line-search logic, the iteration count may not be entirely
# representative of the amount of operations are being performed; or maybe some specific implementations require
# additional operations to be performed when checking stopping conditions. All of this makes it difficult to quantify the exact
# iteration complexity.
#
# We can achieve this by wrapping functions in a dedicated `Counting` type:

mutable struct Counting{T}
    f::T
    gradient_count::Int
    prox_count::Int
end

Counting(f::T) where T = Counting{T}(f, 0, 0)

# Now we only need to intercept any call to `gradient!` and `prox!` and increase counters there:

function ProximalCore.gradient!(y, f::Counting, x)
    f.gradient_count += 1
    return ProximalCore.gradient!(y, f.f, x)
end

function ProximalCore.prox!(y, f::Counting, x, gamma)
    f.prox_count += 1
    return ProximalCore.prox!(y, f.f, x, gamma)
end

# We can run again the previous example, this time wrapping the objective terms within `Counting`:

f = Counting(rosenbrock2D)
g = Counting(IndUnitBall())

solution, iterations = panoc(x0=-ones(2), f=f, g=g)

# and check how many operations where actually performed:

println(f.gradient_count)
println(g.prox_count)


function toeplitz(c,r)
    @assert c[1] == r[1]
    nc = length(c)
    nr = length(r)
    A  = zeros(eltype(c), nc, nr)
    A[:,1] = c
    A[1,:] = r
    for i in 2:nr
        A[2:end,i] = A[1:end-1,i-1]
    end
    A
end

obslength(y::AbstractMatrix) = size(y,1)
obslength(y::AbstractVector) = length(y[1])

Base.oftype(x::Vector{<:Vector}, y::Vector{<:Vector}) = y
Base.oftype(x::AbstractVector, y::Vector{<:Vector}) = getindex.(y, 1)
Base.oftype(x::Matrix, y::Vector{<:Vector}) = reduce(hcat,y)
Base.oftype(x::Matrix, y::Matrix) = y
Base.oftype(x::Vector{<:Vector}, y::Matrix) = [y[:,i] for i in 1:size(y,2)]
function Base.oftype(x::Number, y::AbstractArray)
    length(y) == 1 || throw(ArgumentError("Can only autoconvert one-element arrays"))
    y[]
end

rms(x::AbstractVector) = sqrt(mean(abs2,x))
sse(x::AbstractVector) = x⋅x
mse(x::AbstractVector) = sse(x)/length(x)

rms(x::AbstractMatrix) = sqrt.(mean(abs2.(x),dims=2))[:]
sse(x::AbstractMatrix) = sum(abs2,x,dims=2)[:]
modelfit(y,yh) = 100 * (1 .-rms(y.-yh)./rms(y.-mean(y,dims=2)))
modelfit(y::T,yh::T) where T <: AbstractVector = 100 * (1 .-rms(y.-yh)./rms(y.-mean(y)))
aic(x::AbstractVector,d) = log(sse(x)) .+ 2d/size(x,2)
const nrmse = modelfit


"""
`ind = nearestbelow(array, nearto)` return the indices for the elements in `array` that are nearest below `nearto`.
The nearest elements are `array[ind]`.
# Example
julia> nearestbelow([1,2,3,4,5,6],[1.9, 3.3])
2-element Vector{Float64}:
 1
 3
"""
function nearestbelow(array::Vector{<:Real}, nearto::Vector{<:Real})
	ind = Vector{Int64}(undef,length(nearto));
	for (i,target) in enumerate(nearto)
		ind[i] = findlast(array .<= target);
	end
	return ind
end


## Domains

show(io::IO,d::Segment) = print(io,"the segment [$(leftendpoint(d)),$(rightendpoint(d))]")
function show(io::IO,d::Line)
    if d.center == angle(d) == 0 && d.α == d.β == -1.
        print(io,"ℝ")
    elseif  d.α == d.β == -1.
        print(io,"Line($(d.center),$(angle(d)))")
    else
        print(io,"Line($(d.center),$(angle(d)),$(d.α),$(d.β))")
    end
end

function show(io::IO,d::Ray)
    if d.orientation && angle(d)==0
        print(io,"【$(d.center),∞❫")
    elseif  d.orientation && angle(d)==1.0π
        print(io,"【$(d.center),-∞❫")
    elseif  d.orientation
        print(io,"【$(d.center),exp($(angle(d))im)∞❫")
    elseif !d.orientation  && angle(d)==0
        print(io,"❪∞,$(d.center)】")
    elseif !d.orientation && angle(d)==1.0π
        print(io,"❪-∞,$(d.center)】")
    else
        print(io,"❪exp($(angle(d))im)∞,$(d.center)】")
    end
end

show(io::IO,d::PeriodicSegment) = print(io,"【$(leftendpoint(d)),$(rightendpoint(d))❫")
show(io::IO,d::Circle) =
    print(io,(d.radius==1 ? "" : string(d.radius))*
                    (d.orientation ? "🕒" : "🕞")*
                    (d.center==0 ? "" : "+$(d.center)"))
show(io::IO,d::Point) = print(io,"Point($(d.x))")

## Spaces

show(io::IO,::ConstantSpace{AnyDomain}) = print(io,"ConstantSpace")
show(io::IO,S::ConstantSpace) = print(io,"ConstantSpace($(domain(S)))")
show(io::IO,f::Fun{ConstantSpace{AnyDomain}}) =
    print(io,"$(convert(Number,f)) anywhere")

show(io::IO,f::Fun{ConstantSpace{DD,RR}}) where {DD,RR} =
    print(io,"$(convert(Number,f)) on $(domain(f))")

for typ in ("Chebyshev","Fourier","Laurent","Taylor","SinSpace","CosSpace")
    TYP = Meta.parse(typ)
    @eval function show(io::IO,S::$TYP{D,R}) where {D,R}
        print(io,$typ*"(")
        show(io,domain(S))
        print(io,")")
    end
end

function show(io::IO, S::Ultraspherical)
    print(io,"Ultraspherical($(order(S)),")
    show(io,domain(S))
    print(io,")")
end

function show(io::IO,S::Jacobi)
    S.a == S.b == 0 ? print(io,"Legendre(") : print(io,"Jacobi($(S.b),$(S.a),")
    show(io,domain(S))
    print(io,")")
end

show(io::IO, S::Chebyshev{<:ChebyshevInterval}) = print(io, "Chebyshev()")
show(io::IO, S::Ultraspherical{<:Any,<:ChebyshevInterval}) = print(io, "Ultraspherical($(order(S)))")
show(io::IO, S::Jacobi{<:ChebyshevInterval}) =
    S.a == S.b == 0 ? print(io,"Legendre()") : print(io,"Jacobi($(S.b),$(S.a))")

function show(io::IO,s::JacobiWeight)
    d=domain(s)
    #TODO: Get shift and weights right
    if s.α==s.β
        print(io,"(1-x^2)^$(s.α)[")
    elseif s.β==0
        print(io,"(1-x)^$(s.α)[")
    elseif s.α==0
        print(io,"(1+x)^$(s.β)[")
    else
        print(io,"(1+x)^$(s.β)*(1-x)^$(s.α)[")
    end

    show(io,s.space)
    print(io,"]")
end

function show(io::IO,s::LogWeight)
    d=domain(s)
    #TODO: Get shift and weights right
    if s.α==s.β
        print(io,"log((1-x^2)^$(s.α))[")
    else
        print(io,"log((1+x)^$(s.α)*(1-x)^$(s.β))[")
    end

    show(io,s.space)
    print(io,"]")
end

function show(io::IO,s::QuotientSpace)
    show(io,s.space)
    print(io," /\n")
    show(io,s.bcs;header=false)
end

function show(io::IO,ss::SumSpace)
    s = components(ss)
    show(io,s[1])
    for sp in s[2:end]
        print(io,"⊕")
        show(io,sp)
    end
end

function show(io::IO,ss::PiecewiseSpace)
    s = components(ss)
    show(io,s[1])
    for sp in s[2:end]
        print(io,"⨄")
        show(io,sp)
    end
end

summary(ss::ArraySpace) = string(Base.dims2string(length.(axes(ss))), " ArraySpace")
function show(io::IO,ss::ArraySpace;header::Bool=true)
    header && print(io,summary(ss)*":\n")
    show(io, ss.spaces)
end

function show(io::IO,s::TensorSpace)
    d = length(s.spaces)
    for i=1:d-1
        show(io,s.spaces[i])
        print(io,"⊗")
    end
    show(io,s.spaces[d])
end

function show(io::IO,s::SubSpace)
    print(io,s.space)
    print(io,"|")
    show(io,s.indexes)
end

## Fun

show(io::IO, ::MIME"text/plain", f::Fun) = show(io, f)

function show(io::IO, f::Fun)
    print(io,"Fun(")
    show(io,f.space)
    print(io,",")
    show(io,f.coefficients)
    print(io,")")
end

## MultivariateFun

function show(io::IO,L::LowRankFun)
    print(io,"LowRankFun on ",space(L)," of rank ",rank(L),".")
end

function show(io::IO,P::ProductFun)
    print(io,"ProductFun on ",space(P),".")
end

## Operator

summary(B::Operator) =
    string(typeof(B).name.name)*" : "*string(domainspace(B))*" → "*string(rangespace(B))

struct PrintShow
    str
end
Base.show(io::IO,N::PrintShow) = print(io,N.str)

function show(io::IO,B::Operator;header::Bool=true)
    header && println(io,summary(B))
    dsp=domainspace(B)

    if !isambiguous(domainspace(B)) && (eltype(B) <: Number)
        if isbanded(B) && isinf(size(B,1)) && isinf(size(B,2))
            BM=B[1:10,1:10]

            M=Matrix{Any}(undef,11,11)
            fill!(M,PrintShow("⋅"))
            for j = 1:size(BM,2),k = colrange(BM,j)
                M[k,j]=BM[k,j]
            end

            for k=max(1,11-bandwidth(B,2)):11
                M[k,end]=PrintShow("⋱")
            end
            for j=max(1,11-bandwidth(B,1)):10
                M[end,j]=PrintShow("⋱")
            end

            print_array(io, M)
        elseif isinf(size(B,1)) && isinf(size(B,2))
            BM=B[1:10,1:10]

            M=Matrix{Any}(undef,11,11)
            for k=1:10,j=1:10
                M[k,j]=BM[k,j]
            end

            M[1,end]=PrintShow("⋯")
            M[end,1]=PrintShow("⋮")

            for k=2:11
                M[k,end]=M[end,k]=PrintShow("⋱")
            end

            print_array(io, M)
        elseif isinf(size(B,1))
            BM=B[1:10,1:size(B,2)]

            M=Matrix{Any}(undef,11,size(B,2))
            for k=1:10,j=1:size(B,2)
                M[k,j]=BM[k,j]
            end
            for k=1:size(B,2)
                M[end,k]=PrintShow("⋮")
            end

            print_array(io, M)
        elseif isinf(size(B,2))
            BM=B[1:size(B,1),1:10]

            M=Matrix{Any}(undef,size(B,1),11)
            for k=1:size(B,1),j=1:10
                M[k,j]=BM[k,j]
            end
            for k=1:size(B,1)
                M[k,end]=PrintShow("⋯")
            end

            print_array(io, M)
        else
            print_array(io, AbstractMatrix(B)[1:size(B,1),1:size(B,2)])
        end
    end
end


################################
###    Summary statistics    ###
################################
"""
	poisson_fixation(observed_values,λds, λdn)

Divergence sampling from Poisson distribution. The expected neutral and selected fixations are subset through their relative expected rates ([`fix_neut`](@ref), [`fix_neg_b`](@ref), [`fix_pos_sim`](@ref)). Empirical values are used are used to simulate the locus *L* along a branch of time *T* from which the expected *Ds* and *Dn* raw count estimated given the mutation rate (``\\mu``). Random number generation is used to subset samples arbitrarily given the success rate ``\\lambda`` in the distribution.

```math
\\mathbb{E}[D_N] = X \\in Poisson\\left(\\lambda = D \\times \\left[\\frac{\\mathbb{E}[D_+] + \\mathbb{E}[D_-]}{\\mathbb{E}[D_+] + \\mathbb{E}[D_-] + \\mathbb{E}[D_0]}\\right]\\right)
```
```math
\\mathbb{E}[D_S] = X \\in Poisson\\left(\\lambda = D \\times \\left[\\frac{\\mathbb{E}[D_0]}{\\mathbb{E}[D_+] + \\mathbb{E}[D_-] + \\mathbb{E}[D_0]}\\right]\\right)
```
# Arguments
 - `observed_values::Array`: Array containing the total observed divergence.
 - ` λds::Float64`: expected neutral fixations rate.
 - ` λdn::Float64`: expected selected fixations rate.
# Returns
 - `Array{Int64,1}` containing the expected count of neutral and selected fixations.

"""
function poisson_fixation(;observed_values::Vector{Int64}, λds::Vector{Float64}, λdn::Vector{Float64},λweak::Vector{Float64},λstrong::Vector{Float64})

	ds = @. λds / (λds + λdn) * observed_values
	dn = @. λdn / (λds + λdn) * observed_values
	dweak = @. λweak / (λds + λdn) * observed_values
	dstrong = @. λstrong / (λds + λdn) * observed_values

	sampled_ds     = pois_rand.(ds)
	sampled_dn     = pois_rand.(dn)
	sampled_weak   = pois_rand.(dweak)
	sampled_strong = pois_rand.(dstrong)

	alphas = @. [sampled_weak/sampled_dn sampled_strong/sampled_dn (sampled_weak+sampled_strong)/sampled_dn]

	out = alphas,sampled_dn, sampled_ds
	return out
end

"""
	poisson_polymorphism(observed_values,λps,λpn)

Polymorphism sampling from Poisson distributions. The total expected neutral and selected polimorphism are subset through the relative expected rates at the frequency spectrum ([`fix_neut`](@ref), [`sfs_neut`](@ref),). Empirical sfs are used to simulate the locus *L* along a branch of time *T* from which the expected *Ps* and *Pn* raw count are estimated given the mutation rate (``\\mu``). Random number generation is used to subset samples arbitrarily from the whole sfs given each frequency success rate ``\\lambda`` in the distribution.

The success rate managing the Poisson distribution by the observed count each frequency.  We considered both sampling variance and process variance is affecting the number of variable alleles we sample from SFS. This variance arises from the random mutation-fixation process along the branch. To incorporate this variance we do one sample per frequency-bin and use the total sampled variation and the SFS for the summary statistics.

```math
\\mathbb{E}[P_N] = \\sum_{x=0}^{x=1} X \\in Poisson\\left(\\lambda = SFS_{(x)} \\times \\left[\\frac{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}]}{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}] + \\mathbb{E}[P_{0(x)}]}\\right]\\right)
```

```math
\\mathbb{E}[P_S] = \\sum_{x=0}^{x=1} X \\in Poisson\\left(\\lambda = SFS_{(x)} \\times \\left[\\frac{\\mathbb{E}[P_{0(x)}]}{\\mathbb{E}[P_{+(x)}] + \\mathbb{E}[P_{-(x)}] + \\mathbb{E}[P_{0(x)}]}\\right]\\right)
```

# Arguments
 - `observed_values::Array{Int64,1}`: Array containing the total observed divergence.
 - ` λps::Array{Float64,1} `: expected neutral site frequency spectrum rate.
 - ` λpn::Array{Float64,1} `: expected selected site frequency spectrum rate.
# Returns
 - `Array{Int64,2}` containing the expected total count of neutral and selected polymorphism.

"""
function poisson_polymorphism(;observed_values::Vector{Float64}, λps::Matrix{Float64}, λpn::Matrix{Float64})

	# Neutral λ;
	λ1 = @. λps / (λps + λpn) * observed_values
	# Selected λ;
	λ2 = @. λpn / (λps + λpn) * observed_values

	# Replace negative and 0 rates values. Strange behaviour depending on θ and Γ values
	# Relative rates output NaN values due to 0 divisons.

	replace!(λ1,NaN=>1)
	replace!(λ2,NaN=>1)

	sampled_ps = pois_rand.(λ1)
	sampled_pn = pois_rand.(λ2)

	return (sampled_pn, sampled_ps)
end

"""
	sampled_alpha(observed_values,λds, λdn)

Ouput the expected values from the Poisson sampling process. Please check [`poisson_fixation`](@ref) and [`poisson_polymorphism`](@ref) to understand the samplingn process. α(x) is estimated through the expected values of Dn, Ds, Pn and Ps.

# Arguments
 - `param::parameters`: Array containing the total observed divergence.
 - `d::Array`: observed divergence.
 - `afs::Array`: observed polymorphism.
 - ` λdiv::Array{Float64,2}`: expected fixations rate.
 - ` λdiv::Array{Float64,2}`: expected site frequency spectrum rates.
# Returns
α_summaries,exp_dn,exp_ds,exp_pn,exp_ps,ssAlpha
 - `Array{Int64,2}` containing α(x) values.
 - `Array{Int64,1}` expected non-synonymous divergence.
 - `Array{Int64,1}` expected synonymous divergence.
 - `Array{Int64,1}` expected non-synonymous polymorphism.
 - `Array{Int64,1}` expected synonymous polymorphism.
 - `Array{Int64,1}` expected synonymous polymorphism.
 - `Array{Int64,1}` expected synonymous polymorphism.
 - `Array{Int64,2}` containing α(x) binned values.

	sampled_from_rates(gammaL,gammaH,ppos_l,ppos_h,observedData,nopos)

"""
function sampled_from_rates(models::SubArray,fs::Vector{Float64},d::Vector{Int64},neut::SubArray,sel::SubArray,dsdn::SubArray,output::String)

	ds             = dsdn[:,1]
	dn             = dsdn[:,2]
	dweak          = dsdn[:,3]
	dstrong        = dsdn[:,4]
	gn             = abs.(models[:,4])
	sh             = round.(models[:,end-1],digits=5)

	## Outputs
	alphas, exp_dn, exp_ds = poisson_fixation(observed_values=d,λds=ds,λdn=dn,λweak=dweak,λstrong=dstrong);

	exp_pn, exp_ps         = poisson_polymorphism(observed_values=fs,λps=permutedims(neut),λpn=permutedims(sel));

	## Alpha from expected values. Used as summary statistics
	α_summaries = @. round(1 - ((exp_ds/exp_dn) * (exp_pn/exp_ps)'),digits=5);

	expected_values = hcat(round.(alphas,digits=5),gn,sh,α_summaries);

	expected_values = filter_expected(expected_values);

	w(x,name) = CSV.write(name,DataFrame(x,:auto),delim='\t',header=false);
	w(expected_values,output)
end

"""
	summary_statistics(param::parameters,rates::JLD2.JLDFile,analysis_folder::String,summstat_size::Int64,replicas::Int64,bootstrap::Bool)

Estimate summary statistics using observed data and analytical rates. *analysis_folder* will check for the SFS and divergence file and will be used to output summary statistics

# Arguments
 - `param::parameters` : Mutable structure containing the models
 - `h5_file::String` : HDF5 containing solved models, fixation and polymorphic rates
 - `sfs::Vector{Matrix{Float64}}`: SFS matrix.
 - `divergence::Vector{Matrix{Float64}}` : divergence matrix.
 - `summstat_size::Int64` : Number of summary statistics
 - `replicas::Int64` : Number of bootstrap replicas
 - `bootstrap::Bool` : Boolean to perform or not bootstrapping
# Output
 - Summary statistics to ABC inference
"""
function summary_statistics(;param::parameters,h5_file::String,sfs::Vector,divergence::Vector,analysis_folder::String,summstat_size::Int64,bootstrap::Union{Bool,Int64}=false,threads::Bool=false)

	## Opening files
	if length(sfs) > 1 & bootstrap != false
		throw(ArgumentError("You have more than one SFS and divergence file. Please be sure you have on set of files to bootstrap manually your data."))
	end

	assertion_params(param);

	α,sfs,divergence = data_to_poisson(sfs,divergence,param.dac,bootstrap);

	if any(0 .∈ sfs) | any(0 .∈ divergence)
		throw(ArgumentError("Your SFS contains 0 values at the selected DACs or the divergence is 0. Please consider to bin the SFS and re-estimate the rates using the selected bin as sample the new sample size."))
	end

	# Open rates
	string_cutoff = "cutoff=[" * string(param.cutoff[1]) * "," * string(param.cutoff[end]) * "]";
	h         = jldopen(h5_file);
	tmp       = h[string(param.N)* "/" * string(param.n) * "/" * string_cutoff]

	# Subset index
	idx    = sample(1:size(tmp["models"],1),summstat_size,replace=false);
	models = view(Array(tmp["models"]),idx,:);
	dsdn   = view(Array(tmp["dsdn"]),idx,:);

	# Filtering polymorphic rate by dac
	n    = hcat(map(x -> view(tmp["neut"][x],:),param.dac)...);
	s    = hcat(map(x -> view(tmp["sel"][x],:),param.dac)...);

	neut = view(n,idx,:);
	sel  = view(s,idx,:);

	# Making summaries
	summ_output = analysis_folder .* "/summstat_" .* string.(1:size(sfs,1)) .* ".txt"

	if threads
		ThreadsX.map((x,y,z) -> sampled_from_rates(models,x,y,neut,sel,dsdn,z),sfs,divergence,summ_output);
	else
		progress_pmap((x,y,z) -> sampled_from_rates(models,x,y,neut,sel,dsdn,z),sfs,divergence,summ_output);
	end

	w(x,name) = CSV.write(name,DataFrame(x,:auto),delim='\t',header=false);
	α_output = analysis_folder * "/alphas_" .* string.(1:size(sfs,1)) .* ".txt";

	w.(α, α_output);

	return(summ_output);
end

function filter_expected(x::Matrix{Float64})

	replace!(x, -Inf=>NaN)
	x = x[vec(.!any(isnan.(x),dims=2)),:]
	x = x[(x[:,3] .< 1),:]
	x = x[(x[:,1] .> 0),:]

	return(x)
end

function pol_correction!(sfs_all::Vector{Matrix{Float64}},sfs_in::Vector{Matrix{Float64}};column::Vector{Int}=[2])

	pn_all = map(x->sum(x[:,column]),sfs_all)

	sfs_pn = map(x -> sum(x[:,column]),sfs_in)

	ratio_all_sub = map( (x,y) -> x/y,pn_all,sfs_pn)

	map( (x,y) -> y[:,column] .= x * y[:,column],ratio_all_sub,sfs_in)
end


using FiniteVolumes, StaticArrays, OrdinaryDiffEq, Plots

mesh = PeriodicCartesianMesh(50, 50)
u0 = [SVector(1.0 + exp(-100*((x - 0.4)^2 + (y - 0.3)^2)), 0.0, 0.0) for (x, y) in cell_centers(mesh)]

function f(v, n)
    h, ux, uy = v[1], v[2]/v[1], v[3]/v[1]
    un = ux*n[1] + uy*n[2]
    return SVector(h*un, h*ux*un + h^2*n[1]*9.81/2, h*uy*un + h^2*n[2]*9.81/2)
end

dudt(u, t, p) = -FiniteVolumes.div(f, mesh, u, Upwind())

problem = ODEProblem(dudt, u0, (0.0, 0.2))
sol = solve(problem, Euler(), dt=0.001, saveat=0.01)

plot(
    plot(mesh, sol(0.0), 1),
    plot(mesh, sol(0.06), 1),
    plot(mesh, sol(0.13), 1),
    plot(mesh, sol(0.20), 1),
)


"""
$(SIGNATURES)

Circular path for a swept mesh.
Specify the radius by evaluating it ahead of time.
For example, define `my_circle(θ) = f_swept_circle(0.1, θ)`, and then pass `my_circle` to [`create_swept_mesh`](@ref).
"""
function f_swept_circle(r::Float64, θ::Float64)
    n̂_1 = SVector{3,Float64}( cos(θ), sin(θ), 0.0)
    n̂_2 = SVector{3,Float64}(-sin(θ), cos(θ), 0.0)
    return r * n̂_1, n̂_1, n̂_2
end

"""
$(SIGNATURES)

Straight path for a swept mesh.
Pass this to [`create_swept_mesh`](@ref) to make a basic swept mesh.
"""
function f_swept_triv(θ::Float64)
	n̂_1 = SVector(0.0, 0.0, -1.0)
	n̂_2 = SVector(0.0, 1.0,  0.0)
	return n̂_2 * θ, n̂_1, n̂_2
end

function add_rot_sym_segment!(eM, fun_gen, θ, ϕ, rad, is_open::NTuple{2,Bool})
    # x̂ is in the "radial" direction
    # ŷ is along the path

    p1, x̂1, ŷ1 = fun_gen(θ[1])
    p2, x̂2, ŷ2 = fun_gen(θ[2])
    p3 = (p1 + p2) * 0.5
    p4 = p1 + AngleAxis(ϕ[1], ŷ1...) * x̂1 * rad[1]
    p6 = p1 + AngleAxis(ϕ[2], ŷ1...) * x̂1 * rad[1]
    p5 = p2 + AngleAxis(ϕ[1], ŷ2...) * x̂2 * rad[2]
    p7 = p2 + AngleAxis(ϕ[2], ŷ2...) * x̂2 * rad[2]
    n_offset = length(eM.point)
    i1 = SVector{4,Int64}(1,3,4,6) .+ n_offset
    i2 = SVector{4,Int64}(3,2,5,7) .+ n_offset
    i3 = SVector{4,Int64}(3,4,6,7) .+ n_offset
    i4 = SVector{4,Int64}(4,3,5,7) .+ n_offset
    i5 = SVector{3,Int64}(4,6,7) .+ n_offset
    i6 = SVector{3,Int64}(4,7,5) .+ n_offset
    append!(eM.point, [p1, p2, p3, p4, p5, p6, p7])
    append!(eM.tri, [i5, i6])
    append!(eM.tet, [i1, i2, i3, i4])
    ϵ = zeros(7)
    ϵ[1:3] .= 1.0
    if is_open[1]
        ϵ[1] = 0.0
        push!(eM.tri, SVector{3,Int64}(1,6,4) .+ n_offset)
    end
    if is_open[2]
        ϵ[2] = 0.0
        push!(eM.tri, SVector{3,Int64}(2,5,7) .+ n_offset)
    end
    append!(eM.ϵ, ϵ)
    return nothing
end

"""
$(SIGNATURES)

Creates a mesh by sweeping a 3D path for the given inputs:
* `fun_gen`: function that takes a single input arc length and outputs 1.) the position on path, 2.) a single direction normal to path and 3.) direction along the path.
* `lr`: arc length locations of nodes on path
* `rad`: thickness of swept mesh
* `n_side`: number of sides of swept path
* `is_open`: if the path starts and ends at the same point (e.g. ring) set this to false
* `rot_half`: if the sides of the path are "off" set this to false
"""
function create_swept_mesh(fun_gen::Function, lr::Union{LinRange,Vector{Float64}}, rad::Union{Float64,Vector{Float64}}, n_side::Int64=4,
		is_open::Bool=true; rot_half::Bool=true)

	l_lr = length(lr)
	if isa(rad, Float64)
		rad = zeros(l_lr) .+ rad
	else
		(l_lr == length(rad)) || error("the length of lr and length of rad must be the same")
	end
    eM = eMesh{Tri,Tet}()
    Δ_ϕ = 2 * pi / n_side
    rad = rad ./ cos(Δ_ϕ / 2)
	n_θ = l_lr - 1
    for k_θ = 1:n_θ
        θ = (lr[k_θ], lr[k_θ + 1])
		rad_k = (rad[k_θ], rad[k_θ + 1])
        for k_ϕ = 1:n_side
            ϕ_0 = Δ_ϕ * (k_ϕ  - 0.5 * rot_half)
            ϕ_1 = ϕ_0 + Δ_ϕ
            ϕ = (ϕ_0, ϕ_1)
            if is_open
                is_open_ = (k_θ==1, k_θ==n_θ)
            else
                is_open_ = (false, false)
            end
            add_rot_sym_segment!(eM, fun_gen, θ, ϕ, rad_k, is_open_)
        end
    end
	remove_degenerate!(eM)
    mesh_repair!(eM)
    return eM
end

function f_swept_helix(θ::Float64, coil_sep::Float64)
    delta_z = 1 / (2 * pi) * coil_sep
    r = SVector{3,Float64}(cos(θ), sin(θ), θ * delta_z)
    dir_1 = normalize(SVector{3,Float64}(-sin(θ), cos(θ), delta_z))
    dir_2 = SVector{3,Float64}( cos(θ), sin(θ), 0.0)
    return r, dir_1, dir_2
end

f_swept_circle(θ::Float64) = f_swept_helix(θ, 0.0)


# Benchmarks to compare the performance of
#
#   - devec macro (in Devectorize)
#   - hand-coded devectorized for-loop
#   - vectorized Julia code
#

using Devectorize

# benchmark helpers

repeat = 20

macro my_bench(Name)
    quote
        # create task
        println("bench: ", $(string(Name)))
        task = ($Name)()

        vec_eval(task, a, b, c)
        t1 = @elapsed for i = 1 : repeat
            vec_eval(task, a, b, c)
        end
        @printf "    vec_eval   : %8.4f sec  |  gain = %7.4f\n" t1 t1 / t1

        devec_eval(task, a, b, c)
        t2 = @elapsed for i = 1 : repeat
            devec_eval(task, a, b, c)
        end
        @printf "    devec_eval : %8.4f sec  |  gain = %7.4f\n" t2 t1 / t2

        hand_loop(task, a, b, c)
        t3 = @elapsed for i = 1 : repeat
            hand_loop(task, a, b, c)
        end
        @printf "    hand_loop  : %8.4f sec  |  gain = %7.4f\n" t3 t1 / t3

        println()
    end
end

# prepare data

m = 1000
n = 1000
a = zeros(m, n)
b = zeros(m, n)
c = zeros(m, n)

# simple element-wise

type simple_ewise end
formula(::simple_ewise) = "(a - b).^2 + c"

function vec_eval{T<:Real}(::simple_ewise, a::Array{T}, b::Array{T}, c::Array{T})
    r = sqr(a - b) + c   # Note: Devectorize defines a sqr on arrays
end

function devec_eval{T<:Real}(::simple_ewise, a::Array{T}, b::Array{T}, c::Array{T})
    @devec r = sqr(a - b) + c
end

function hand_loop{T<:Real}(::simple_ewise, a::Array{T}, b::Array{T}, c::Array{T})
    r = similar(a)
    for i = 1 : length(a)
        v = (a[i] - b[i])
        r[i] = v * v + c[i]
    end
end

# complex element-wise

type complex_ewise end
formula(::complex_ewise) = "log(exp((a - b).^2) + exp(a + b)) - c .* log(c)"

function vec_eval{T<:Real}(::complex_ewise, a::Array{T}, b::Array{T}, c::Array{T})
    r = log(exp((a - b).^2) + exp(a + b)) - c .* log(c)
end

function devec_eval{T<:Real}(::complex_ewise, a::Array{T}, b::Array{T}, c::Array{T})
    @devec r = log(exp((a - b).^2) + exp(a + b)) - c .* log(c)
end

function hand_loop{T<:Real}(::complex_ewise, a::Array{T}, b::Array{T}, c::Array{T})
    r = similar(a)
    for i = 1 : length(a)
        r[i] = log(exp((a[i] - b[i]).^2) + exp(a[i] + b[i])) - c[i] .* log(c[i])
    end
end

# shift dot

type shift_dot end
formula(::shift_dot) = "sum( (a - mean(a)) .* (b - mean(b)) )"

function vec_eval{T<:Real}(::shift_dot, a::Array{T}, b::Array{T}, c::Array{T})
    r = sum( (a - mean(a)) .* (b - mean(b)) )
end

function devec_eval{T<:Real}(::shift_dot, a::Array{T}, b::Array{T}, c::Array{T})
    @devec r = sum( (a - mean(a)) .* (b - mean(b)) )
end

function hand_loop{T<:Real}(::shift_dot, a::Array{T}, b::Array{T}, c::Array{T})
    n = length(a)

    # calculate mean(a)
    sa = 0.
    for i = 1 : n
        sa += a[i]
    end
    ma = sa / n

    # calculate mean(b)
    sb = 0.
    for i = 1 : n
        sb += b[i]
    end
    mb = sb / n

    # calculate shift dot
    r = 0.
    for i = 1 : n
        r += (a[i] - ma) * (b[i] - mb)
    end
end

# column wise sum

type colwise_sum end
formula(::colwise_sum) = "sum(a, 1)"

function vec_eval{T<:Real}(::colwise_sum, a::Array{T}, b::Array{T}, c::Array{T})
    r = sum(a, 1)
end

function devec_eval{T<:Real}(::colwise_sum, a::Array{T}, b::Array{T}, c::Array{T})
    @devec r = sum(a, 1)
end

function hand_loop{T<:Real}(::colwise_sum, a::Array{T}, b::Array{T}, c::Array{T})
    m, n = size(a)
    r = zeros(1, n)
    for j = 1 : n
        s = 0.
        for i = 1 : m
            s += a[i, j]
        end
        r[j] = s
    end
end

# row-wise sum

type rowwise_sum end
formula(::rowwise_sum) = "sum(a, 2)"

function vec_eval{T<:Real}(::rowwise_sum, a::Array{T}, b::Array{T}, c::Array{T})
    r = sum(a, 2)
end

function devec_eval{T<:Real}(::rowwise_sum, a::Array{T}, b::Array{T}, c::Array{T})
    @devec r = sum(a, 2)
end

function hand_loop{T<:Real}(::rowwise_sum, a::Array{T}, b::Array{T}, c::Array{T})
    m, n = size(a)
    r = zeros(1, n)
    for j = 1 : n
        for i = 1 : m
            r[i] += a[i,j]
        end
    end
end

# col-wise euclidean

type colwise_eucdist end
formula(::colwise_eucdist) = "sqrt(sum((a - b).^2, 1))"

function vec_eval{T<:Real}(::colwise_eucdist, a::Array{T}, b::Array{T}, c::Array{T})
    r = sqrt(sum(sqr(a - b), 1))
end

function devec_eval{T<:Real}(::colwise_eucdist, a::Array{T}, b::Array{T}, c::Array{T})
    @devec r = sqrt(sum(sqr(a - b), 1))
end

function hand_loop{T<:Real}(::colwise_eucdist, a::Array{T}, b::Array{T}, c::Array{T})
    m, n = size(a)
    r = zeros(1, n)
    for j = 1 : n
        s = 0.
        for i = 1 : m
            v = a[i, j] - b[i, j]
            s += v * v
        end
        r[j] = sqrt(s)
    end
end

#  Benchmarks

@my_bench simple_ewise
@my_bench complex_ewise
@my_bench shift_dot
@my_bench colwise_sum
@my_bench rowwise_sum
@my_bench colwise_eucdist


using CollegeStratCollege, Test

mdl = CollegeStratCollege;

function study_time_grid_test(perCourse)
	@testset "Study Time grid" begin
        stg = mdl.make_test_study_time_grid(perCourse);
        @test mdl.validate_stg(stg);

        @test stime_grid_per_course(stg) == perCourse;
        n = n_study_times(stg);

        nTried = 3;
        sTimeGridV = study_time_grid(stg, nTried);
        @test all(sTimeGridV .> 0.0);
        @test length(sTimeGridV) == n;
        @test max_study_time(stg, n) ≈ maximum(sTimeGridV);
        sTimeGrid2V = study_time_from_grid(stg, stg.studyTimeV, nTried);
        @test isapprox(sTimeGridV, sTimeGrid2V);
	end
end

@testset "StudyTimeGrid" begin
    for perCourse in (true, false)
        study_time_grid_test(perCourse);
    end
end

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


"""
`find_all_combinations(vectors)`

Find all possible combinations among the vectors of values.

**Arguments**
- `vectors`: Vector of vectors representing different value sets.

**Returns**
- `combinations`: Matrix containing all combinations of values from the
    individual vectors. Each row represents one combination.

**Example**
```
julia> vectors = [[1, 2], [3, 4, 5], [6]];
julia> combinations = find_all_combinations(vectors);
julia> print(combinations);
[1 3 6; 1 4 6; 1 5 6; 2 3 6; 2 4 6; 2 5 6]
```
"""
function find_all_combinations(vectors)

    # Number of input vectors
    n = length(vectors)

    # The first column of the input matrix is the last input vector
    combinations = hcat(vectors[n])

    # Iterate the remaining vectors
    for i = (n-1):-1:1

        # Get the current vector and its length
        x = vectors[i]
        nx = length(x)

        # Get the current size of the combinations matrix
        nc = size(combinations, 1)

        # Repeat the current combinations for each element in x
        combinations = repeat(combinations, nx)

        # Add an element from x to each repetition of combinations
        x = vcat([fill(xi, nc) for xi in x]...)
        combinations = hcat(x, combinations)

    end

    return combinations

end


@testset "1637.widest-vertical-area-between-two-points-containing-no-points.jl" begin
    @test max_width_of_vertical_area([[8,7],[9,9],[7,4],[9,7]]) == 1
    @test max_width_of_vertical_area([[3,1],[9,0],[1,0],[1,4],[5,3],[8,8]]) == 3
end


@test begin
    l = Segment(Point(0, 0), Point(1, 1))
    L = Segment(Point(2, 0), Point(3, 1))
    !collides(l, L)
end

@test begin
    l = Segment(Point(0, 0), Point(1, 1))
    L = Segment(Point(1, 0), Point(2, 1))
    !collides(l, L)
end

@test begin
    l = Segment(Point(0, 0), Point(1, 1))
    L = Segment(Point(0.75, 0), Point(1, 0.25))
    !collides(l, L)
end

@test begin
    l = Segment(Point(0, 0), Point(1, 0))
    L = Segment(Point(0, 1), Point(1, 1))
    !collides(l, L)
end

@test begin
    l = Segment(Point(0, 0), Point(1, 0))
    L = Segment(Point(2, 0), Point(3, 0))
    !collides(l, L)
end

@test collides(Segment(Point(-1, 0), Point(0, 0)), Segment(Point(0, 0), Point(1, 0)))

@test begin
    l = Segment(Point(0, 0), Point(1, 0))
    L = Segment(Point(1, 0), Point(3, 0))
    collides(l, L)
end

@test begin
    l = Segment(Point(0, 0), Point(1, 1))
    L = Segment(Point(0, 1), Point(1, 0))
    collides(l, L)
end

@test begin
    l = Segment(Point(0, 0), Point(1, 1))
    collides(l, l)
end

# Box vs segment
@test begin
    b = Box(Point(0, 0), Point(1, 1))
    l = Segment(Point(2, 2), Point(3, 3))
    !collides(l, b)
end

@test begin
    b = Box(Point(0, 0), Point(1, 1))
    l = Segment(Point(1, 1), Point(3, 3))
    collides(l, b)
end

@test begin
    b = Box(Point(0, 0), Point(1, 1))
    l = Segment(Point(1, 1), Point(0, 1))
    collides(l, b)
end

@test begin
    b = Box(Point(0, 0), Point(1, 1))
    l = Segment(Point(1, 1), Point(0, 0))
    collides(l, b)
end

@test begin
    b = Box(Point(0, 0), Point(1, 1))
    l = Segment(Point(0.25, 0.25), Point(0.75, 0.75))
    collides(l, b)
end

# -------------------------------------------------------------------
# # Both in
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0.25, 0.25, 0.25), Point(0.75, 0.75, 0.75))
    collides(l, b)
end

# One in other out
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0.5, 0.5, 0.5), Point(3, 3, 3))
    collides(l, b)
end

# Both out but cross
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0.5, 0.5, -1), Point(0.5, 0.5, 3))
    collides(l, b)
end

# Both out but cross
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0.25, 0.25, -1), Point(0.75, 0.75, 3))
    collides(l, b)
end

# Both out but cross
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0.123, 0.1, -1), Point(0.567, 0.456, 3))
    collides(l, b)
end

# Point collision
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(1, 1, 1), Point(3, 3, 3))
    collides(l, b)
end

# Both out
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(2, 2, 2), Point(3, 3, 3))
    !collides(l, b)
end

# Both out
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0.5, 0.5, 0), Point(0.5, 0.5, -2))
    collides(l, b)
end

# Entire segment collision
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0, 0, 0), Point(1, 0, 0))
    collides(l, b)
end

# Intersect in plane
@test begin
    b = Box(Point(0, 0, 0), Point(1, 1, 1))
    l = Segment(Point(0, 0, 0), Point(1, 1, 0))
    collides(l, b)
end


using DrWatson
@quickactivate "ExercisesRepo"
include(srcdir("style.jl"))

# %% sensitive dependence
using DynamicalSystems, PyPlot, Random
Random.seed!(5)
close("all")
fig = figure(figsize=(2figx/3, figx/2))
ax = gca()
u0 = [10,10.0,10]
lo = Systems.lorenz([10,10.0,10])

for i in 1:3
    u = u0 .+ i*1e-3
    tr = trajectory(lo, 15, u)
    plot(0:0.01:15, tr[:, 1],  c = COLORS[[2,3,1][i]])
end

xlabel("\$t\$",labelpad = -10)
yticks(-15:10:15)
xticks(0:5:15)
ylabel("\$x\$",labelpad = -15)
tight_layout()
subplots_adjust(bottom = 0.15, left = 0.12)
fsave(joinpath(figdir, "sensitive"), fig)

# %% stretching and folding for logistic
lo = Systems.logistic()
X = trajectory(lo, 1000, Ttr = 100)
# close("all")
fig = figure()
# ax = subplot2grid((1,3), (0,0))
ax = subplot(121)
ax.scatter(X[1:end-1], X[2:end], color = COLORS[1], s = 5, label = "\$x_{i+2}\$")
ax.set_xlabel("\$x_{i}\$")
ax.set_yticks([0, 0.5, 1])
ax.scatter(X[1:end-2], X[3:end], color = COLORS[2], s = 5, label = "\$x_{i+2}\$")
ax.legend(markerscale = 5)

using3D()
# ax2 = subplot2grid((1,3), (0,1), projection = "3d", colspan = 2)
ax2 = subplot(122, projection = "3d")
ax2.scatter(X[1:end-2], X[2:end-1], X[3:end], color = COLORS[1], s = 5)

ax2.set_xticklabels([])
ax2.set_yticklabels([])
ax2.set_zticklabels([])
ax2.set_xlabel("\$x_{i}\$", labelpad = -10)
ax2.set_ylabel("\$x_{i+1}\$", labelpad = -10)
ax2.set_zlabel("\$x_{i+2}\$", labelpad = -10)

# fold arrow
s = (0.15, 0.5, 1.2)
e = (0.85, 0.5, 1.2)
ax2.quiver3D(e..., (s .- e)..., color = COLORS[3])
ax2.text3D(0.5, 0.1, 1.2, "Fold", color = COLORS[3])

# stretch arrow
s = (0.75, 0.35, 0.9)
e = (0.95, 0.15, 0.1)
ax2.quiver3D(e..., (s .- e)..., color = COLORS[5])

s = (0.55, 0.45, 0.1)
e = (0.75, 0.35, 0.9)
ax2.quiver3D(e..., (s .- e)..., color = COLORS[5])
ax2.text3D(1.1, 0.45, 0.1, "Stretch", color = COLORS[5])

tight_layout()
fig.subplots_adjust(top = 0.95, wspace = 0.1)
add_identifiers!(fig)
savefig(joinpath(@__DIR__, "stretching.png"))

# %% Folding in henon map
a = 1.4; b = 0.3
he = Systems.henon(;a=a, b=b)
close("all")
fig = figure()

x0, y0 = columns(trajectory(he, 5000, Ttr = 100))
x1 = x0
y1 = @. 1 - a*x0^2 + y0

x2 = b .* x1
y2 = y1

y3 = x2
x3 = y2

ax1 = subplot(141)
ax1.scatter(x0, y0, color = COLORS[1], s = 1, label = "initial")
ax1.text(0.1, 0.05, "initial", transform=ax1.transAxes)

αprev = 0.75
ax2 = subplot(142)
ax2.scatter(x0, y0, color = COLORS[1], s = 0.25, alpha = αprev)
ax2.scatter(x1, y1, color = COLORS[2], s = 1, label = "stretch/bend")
ax2.text(0.1, 0.05, "bend", transform=ax2.transAxes)
ax2.arrow(0,-0.19, 0, 0.8, head_width=0.1, head_length=0.1, linewidth=4, color=COLORS[6], length_includes_head=true)
ax2.arrow(-1.5, 0.0,  0, -0.5, head_width=0.1, head_length=0.1, linewidth=4, color=COLORS[6], length_includes_head=true)
ax2.arrow(1.5, 0.0,  0, -0.5, head_width=0.1, head_length=0.1, linewidth=4, color=COLORS[6], length_includes_head=true)

ax3 = subplot(143)
ax3.scatter(x1, y1, color = COLORS[2], s = 0.25, alpha = αprev)
ax3.scatter(x2, y2, color = COLORS[3], s = 1)
ax3.text(0.1, 0.05, "squeeze", transform=ax3.transAxes)
ax3.arrow(-1.5, 0.0,  1, 0, head_width=0.1, head_length=0.1, linewidth=4, color=COLORS[6], length_includes_head=true)
ax3.arrow(1.5, 0.0,  -1, 0, head_width=0.1, head_length=0.1, linewidth=4, color=COLORS[6], length_includes_head=true)

ax4 = subplot(144)
ax4.scatter(x2, y2, color = COLORS[3], s = 0.25, alpha = αprev)
ax4.scatter(x3, y3, color = COLORS[4], s = 1, label = "reflect")
ax4.text(0.1, 0.05, "reflect", transform=ax4.transAxes)

ax4.plot([-2,2], [-2,2], ls = "dashed", lw = 2, color = COLORS[6])

# ticks and stuff
ax1.set_ylabel("\$y\$")
for ax in (ax1, ax2, ax3, ax4)
    ax.set_xlim(-2, 2)
    ax.set_ylim(-1.5, 1.5)
    ax.set_yticks(-1.5:0.5:1.5)
    ax.set_xticks(-1.5:1:1.5)
    ax.grid("on")
    ax.set_yticklabels([])
    ax.set_xticklabels([])
    ax.set_xlabel("\$x\$", labelpad = -10)
    # ax.legend(markerscale = 4, scatterpoints=1)
end
tight_layout()
subplots_adjust(left = 0.05, bottom = 0.1, wspace = 0.1)
add_identifiers!(fig; xloc = 0.01)
# savefig(joinpath(figdir, "stretchinghenon"))

# %% definition of lyapunov
using LinearAlgebra
he = Systems.henon()
X₁ = trajectory(he, 50)
u₂ = get_state(he) .+ 1e-6
X₂ = trajectory(he, 50, u₂)
δ  = norm.(X₂.data .- X₁.data)
λ = lyapunov(he, 10000)

close("all")
figure(figsize=(figx/2,figx/4))
ax = gca()
ax.plot(0:50, log.(δ), c = COLORS[1], label ="\$\\ln(\\delta(n)))\$")
ax.set_yticks(-12:4:0)
ax.set_xlabel("\$n\$", labelpad=-10)
# Lyapunov
ax.plot([0, 25] .+ 5, λ .* [0, 25] .- 13, color = COLORS[2])
ax.text(20, -9, "\$\\lambda\$=$(round(λ;digits=2))", color = COLORS[2])
ax.legend()
xticks(0:15:50)
tight_layout()
subplots_adjust(bottom = 0.2, left = 0.12)
savefig(joinpath(figdir, "lyapunov"))

# %% gali map
using DynamicalSystems, PyPlot
figure()
hh = Systems.henonheiles()
ics = Systems.henonheiles_ics(0.13, 10)
cmap = matplotlib.cm.get_cmap("viridis")
for ic in ics
    psos = poincaresos(hh, (1, 0.0), 10000; u0 = ic)
    GALI, t = gali(hh, 1000, 4; u0 = ic)
    v = clamp(t[end]/1000, 0, 1)
    scatter(psos[:, 2], psos[:, 4], color = cmap(v), s = 2)
end

cb = colorbar()
xticks()
xlabel("\$y\$", labelpad = -10)
yticks(-0.5:0.3:0.5)
ylabel("\$p_y\$", labelpad = -10)
cb.set_ticks([0, 1])
cb.set_label("regularity")
tight_layout()
subplots_adjust(bottom = 0.15, top = 0.95, left = 0.1, wspace = 0.01)
savefig(joinpath(figdir, "gali.png"))


module LiftProps

export fitliftslope,findclmax,findclmaxlinear,findclmin

"""
    findclmax(aoa,cl)
Determines clmax as greatest lift coefficient. Returns aoaclmax, clmax
"""
function findclmax(aoa,cl)
  indclmax = argmax(cl)
  clmax = cl[indclmax]
  aoaclmax = aoa[indclmax]
  return aoaclmax,clmax
end #findclmax

"""
    findclmax(aoa,cl,npnts::Integer)
Determines clmax as when all angles greater than the clmax angle for `npnts`
data points are less than clmax. Assumes points are ordered by angle of attack.
Returns aoaclmax, clmax
"""
function findclmax(aoa,cl,npnts::Integer)
  # default to basic method if conditions not satisfied
  aoaclmax,clmax = findclmax(aoa,cl)
  # start at zero
  idxzero = argmin(abs.(aoa))
  for i = idxzero:length(cl)-npnts
    # check if all points in range are less than current point
    if count(cl[(i+1):(i+npnts)] .> cl[i]) == 0
      clmax = cl[i]
      aoaclmax = aoa[i]
      break
    end
  end
  return aoaclmax,clmax
end #findclmax

"""
    findclmax(aoa,cl,range::AbstractFloat)
Determines clmax as when all angles greater than the clmax angles for `range`
radians are less than clmax. Assumes points are ordered by angle of attack.
Returns aoaclmax, clmax
"""
function findclmax(aoa,cl,range::AbstractFloat)
  # default to basic method if conditions not satisfied
  aoaclmax,clmax = findclmax(aoa,cl)
  # start at zero
  idxzero = argmin(abs.(aoa))
  for i = idxzero:length(cl)
    # find index of point approximately `range` greater
    idx = argmin(abs.(aoa[i].+range.-aoa))
    if idx == i
      break
    else
      # check if all points in range are less than current point
      if count(cl[(i+1):idx] .> cl[i]) == 0
        clmax = cl[i]
        aoaclmax = aoa[i]
        break
      end
    end
  end
  return aoaclmax,clmax
end #findclmax

"""
    findclmin(aoa,cl)
Determines clmin as smallest lift coefficient.  Returns aoaclmin, clmin
"""
function findclmin(aoa,cl)
  indclmin = argmin(cl)
  clmin = cl[indclmin]
  aoaclmin = aoa[indclmin]
  return aoaclmin,clmin
end #findclmax

"""
    findclmin(aoa,cl,npnts::Integer)
Determines clmin as when all angles less than the clmin angle for `npnts` data
points are greater than clmin. Assumes points are ordered by angle of attack.
Returns aoaclmin, clmin
"""
function findclmin(aoa,cl,npnts::Integer)
  # default to basic method if conditions not satisfied
  aoaclmin,clmin = findclmin(aoa,cl)
  # start at zero
  idxzero = argmin(abs.(aoa))
  for i = idxzero:-1:(npnts+1)
    # check if all points in range are less than current point
    if count(cl[(i-1):-1:(i-npnts)] .< cl[i]) == 0
      clmin = cl[i]
      aoaclmin = aoa[i]
      break
    end
  end
  return aoaclmin,clmin
end #findclmax

"""
    findclmin(aoa,cl,range::AbstractFloat)
Determines clmin as when all points less than the clmin angle for `range`
radians are greater than clmin. Returns aoaclmin, clmin
"""
function findclmin(aoa,cl,range::AbstractFloat)
  # default to basic method if conditions not satisfied
  aoaclmin,clmin = findclmin(aoa,cl)
  # start at zero
  idxzero = argmin(abs.(aoa))
  for i = idxzero:-1:1
    # find index of point approximately `range` less
    idx = argmin(abs.(aoa[i].-range.-aoa))
    if idx == i
      break
    else
      # check if all points in range are greater than current point
      if count(cl[(i-1):-1:idx] .< cl[i]) == 0
        clmin = cl[i]
        aoaclmin = aoa[i]
        break
      end
    end
  end
  return aoaclmin,clmin
end #findclmax

"""
    fitliftslope(aoa,cl,tol::Real=0.05,allnegfit::Bool=false)
Returns the lift slope and zero lift angle of attack as well as the data used
to compute these parameters.  The tolerance `tol` (in terms of cl) for the fit
and whether to include all data points below zero angle of attack `allnegfit`
may be specified. Returns liftslope,zeroliftangle,aoafit,clfit
"""
function fitliftslope(aoa,cl,tol::Real=0.05,allnegfit::Bool=false,center=0.0)

  if length(aoa) != length(cl)
    error("aoa and cl must have the same length")
  end

  # split into positive and negative angles of attack
  idxcenter = argmin(abs.(aoa .- center))
  idxneg = findall(aoa .< aoa[idxcenter])
  sort!(idxneg,rev = true)
  idxpos = findall(aoa .> aoa[idxcenter])
  sort!(idxpos)

  # initialize arrays of data to fit
  aoafit = [aoa[idxcenter]]
  clfit = [cl[idxcenter]]
  # create indexes for increasing and decreasing aoa
  ipos = 1
  ineg = 1
  # initialize flag to stop adding data > 0
  upperflag = false
  if ipos > length(idxpos)
    upperflag = true
  end
  # initialize flag to stop adding data < 0
  lowerflag = false
  if ineg > length(idxneg)
    lowerflag = true
  end
  # compute lift slope and zero lift angle of attack
  success = false
  liftslope = 2*pi
  zeroliftangle = 0.0
  for i = 1:(length(idxpos)+length(idxneg))
    # add positive angle of attack point to fit if closest to zero aoa
    if upperflag == false
      if lowerflag == true || abs(aoa[idxpos[ipos]]) <= abs(aoa[idxneg[ineg]])
        # add point
        push!(aoafit,aoa[idxpos[ipos]])
        push!(clfit,cl[idxpos[ipos]])
        # order points
        order = sortperm(aoafit)
        aoafit = aoafit[order]
        clfit = clfit[order]
        # create least squares fit
        lsqsol = hcat(aoafit, -ones(length(aoafit)))\clfit
        # if liftslope not defined from previous iteration, assign it here
        if ipos == 1
            liftslope = lsqsol[1]
        end
        # get slope and error of added point
        yerror = lsqsol[1]*(aoa[idxpos[ipos:end]].-lsqsol[2]/lsqsol[1]).-cl[idxpos[ipos:end]]
        # slope = (clfit[end] - clfit[end-1])/(aoafit[end] - aoafit[end-1])
        # check if on linear portion of lift curve
        if any(yerror .< tol) #slope > 0.5*liftslope ||
          liftslope = lsqsol[1]
          zeroliftangle = lsqsol[2]/lsqsol[1]
          success = true
        else # if not on linear portion remove from fit
          aoafit = aoafit[1:(end-1)]
          clfit = clfit[1:(end-1)]
          upperflag = true
        end
        #  increment counter and check for more points greater than aoa=0
        ipos = ipos+1
        if ipos > length(idxpos)
          upperflag = true
        end
      end
    end
    # add negative angle of attack point to fit if closest to zero
    if lowerflag == false
      if upperflag == true || abs(aoa[idxneg[ineg]]) <= abs(aoa[idxpos[ipos]])
        # add point
        pushfirst!(aoafit,aoa[idxneg[ineg]])
        pushfirst!(clfit,cl[idxneg[ineg]])
        # create least squares fit
        lsqsol = hcat(aoafit, -ones(length(aoafit)))\clfit
        # get error of added point
        yerror = lsqsol[1]*(aoa[idxneg[ineg:end]].-lsqsol[2]/lsqsol[1]).-cl[idxneg[ineg:end]]
        # check if on linear portion of lift curve
        if any(abs.(yerror) .< tol) || allnegfit
          liftslope = lsqsol[1]
          zeroliftangle = lsqsol[2]/lsqsol[1]
          success = true
        else # if not on linear portion remove from fit
          aoafit = aoafit[2:end]
          clfit = clfit[2:end]
          lowerflag = true
        end
        #  increment counter and check for more points less than aoa=0
        ineg = ineg+1
        if ineg > length(idxneg)
          lowerflag = true
        end
      end
    end
    # break if no more points to add
    if upperflag && lowerflag
      break
    end
  end
  if success == false
    liftslope,zeroliftangle,aoafit,clfit = fitliftslope(aoa,cl,tol,allnegfit,aoa[idxcenter]+1.0*pi/180)
  end

  return liftslope,zeroliftangle,aoafit,clfit
end

"""
    findclmaxlinear(aoa,cl,liftslope::Real,zeroliftangle::Real;tol::Real=0.1,
    interpolate::Bool=true)
Returns the angle of attack and cl of the maximum linear lift coefficient. `tol`
is used to determine the liftslope line tolerance. Returns aoaclmaxlinear,clmaxlinear
"""
function findclmaxlinear(aoa,cl,liftslope::Real,zeroliftangle::Real;
  tol::Real=0.1,interpolate::Bool=true)

  # find where difference in cl is greater than tol
  indclmaxlinear = 0
  idxzero = argmin(abs.(aoa))
  for i = idxzero:length(aoa)
    predictedcl = liftslope*(aoa[i]-zeroliftangle)
    indclmaxlinear = i
    if abs(cl[i]-predictedcl) > tol
      break
    end
  end
  if interpolate
    m = (cl[indclmaxlinear]-cl[indclmaxlinear-1])/(aoa[indclmaxlinear]-aoa[indclmaxlinear-1])
    aoaclmaxlinear = (m*aoa[indclmaxlinear-1]-liftslope*zeroliftangle-tol-cl[indclmaxlinear-1])/(m-liftslope)
    clmaxlinear = liftslope*(aoaclmaxlinear-zeroliftangle)-tol
  else
    aoaclmaxlinear = aoa[indclmaxlinear]
    clmaxlinear = cl[indclmaxlinear]
  end
  return aoaclmaxlinear,clmaxlinear
end #findclmaxlinear

end #module


module CPortGraphDynam

using Catlab.WiringDiagrams
using Catlab.WiringDiagrams.CPortGraphs
using  Catlab.CategoricalAlgebra
using  Catlab.CategoricalAlgebra.FinSets
using Catlab.Graphs
import Catlab.Graphs: Graph
using Catlab.Graphics.GraphvizGraphs: to_graphviz

import Catlab.WiringDiagrams: oapply

import ..UWDDynam: nstates, nports, eval_dynamics, euler_approx
import ..DWDDynam: AbstractMachine, ContinuousMachine, DiscreteMachine,
ninputs, noutputs
import Catlab.CategoricalAlgebra.CSets: migrate!

using Base.Iterators
import Base: show, eltype

ContinuousMachine{T}(nports::Int, nstates::Int, d::Function, r::Function) where T =
    ContinuousMachine{T}(nports, nstates, nports, d, r)
DiscreteMachine{T}(nports::Int, nstates::Int, d::Function, r::Function) where T =
    DiscreteMachine{T}(nports, nstates, nports, d, r)

migrate!(g::Graph, pg::OpenCPortGraph) = migrate!(g, migrate!(CPortGraph(), pg))

draw(g::Graph) = to_graphviz(g, prog="neato", edge_labels=true, node_labels=true)
draw(pg::OpenCPortGraph) = draw(migrate!(Graph(), pg))

concat(xs::Vector) = (collect ∘ Iterators.flatten)(xs)

nports(d::OpenCPortGraph, b::Int) = incident(d, b, :box) |> length
nports(d::OpenCPortGraph, b) = map(length, incident(d, b, :box))
nports(d::OpenCPortGraph, b::Colon) = map(length, incident(d, :, :box))

colimitmap!(f::Function, output, C::Colimit, input) = begin
    for (i,x) in enumerate(input)
        y = f(i, x)
        I = legs(C)[i](1:length(y))
        # length(I) == length(y) || error("colimitmap! attempting to fill $(length(I)) slots with $(length(y)) values")
        output[I] .= y
    end
    return output
end

@inline fillreadouts!(y, d, xs, Ports, statefun) = colimitmap!(y, Ports, xs) do i,x
    return x.readout(statefun(i))
end

@inline fillstates!(y, d, xs, States, statefun, inputfun, p, t) = colimitmap!(y, States, xs) do i, x
    return x.dynamics(statefun(i), inputfun(i), p, t)
end

fillreadins!(readins, d, readouts) = begin
    for b in parts(d, :Box)
        ports = incident(d, b, :box)
        for p in ports
            ws = incident(d, p, :tgt)
            qs = d[ws, :src]
            readins[p] += sum(readouts[qs])
        end
    end
    return readins
end

"""    oapply(d::OpenCPortGraph, ms::Vector)

Implements the operad algebras for directed composition of dynamical systems given a
composition pattern (implemented by an open circular port graph `d`)
and primitive systems (implemented by a collection of
machines `ms`).

Each box of the composition pattern `d` is filled by a machine with the
appropriate type signature. Returns the composite machine.
"""
function oapply(d::OpenCPortGraph, xs::Vector{Machine}) where {T, Machine<:AbstractMachine{T}}
    x -> FinSet(x.nstates)
    S = coproduct((FinSet∘nstates).(xs))
    Inputs = coproduct((FinSet∘ninputs).(xs))
    Ports = coproduct([FinSet.(nports(d, b)) for b in parts(d, :Box)])
    state(u::Vector, b::Int) = u[legs(S)[b](1:xs[b].nstates)]
    readouts = zeros(T, length(apex(Ports)))
    readins  = zeros(T, length(apex(Ports)))
    ϕ = zeros(T, length(apex(S)))
    v = (u::AbstractVector, x::AbstractVector, p, t::Real) -> begin
        # length(p) == length(d[:, :con]) || error("Expected $(length(d[:, :con])) parameters, have $(length(p))")
        statefun(b) = state(u,b)
        inputfun(b) = readins[incident(d, b, :box)]
        fillreadouts!(readouts, d, xs, Ports, statefun)
        # communicate readouts to the ports at the other end of the wires, external connections directly fill ports
        readins .= 0
        fillreadins!(readins, d, readouts)
        readins[d[:, :con]] .+= x
        fillstates!(ϕ, d, xs, S, statefun, inputfun, p, t)
        return ϕ
    end
    function readout(u::AbstractVector)
        statefun(b) = state(u,b)
        fillreadouts!(readouts, d, xs, Ports, statefun)
        return readouts[d[:, :con]]
    end
    return Machine( nparts(d, :OuterPort), apex(S).set, v, readout)
end

"""    oapply(d::OpenCPortGraph, m::AbstractMachine)

A version of `oapply` where each box of `d` is filled with the machine `m`.
"""
function oapply(d::OpenCPortGraph, x::AbstractMachine)
    oapply(d, collect(repeated(x, nparts(d, :Box))))
end

"""    barbell(n::Int)

Constructs an open CPG with two boxes each with `n` ports. The ``i``th ports on each box are connected.
"""
barbell(k::Int) = begin
  g = OpenCPortGraph()
  add_parts!(g, :Box, 2)
  add_parts!(g, :Port, 2k; box=[fill(1,k); fill(2,k)])
  add_parts!(g, :Wire, k; src=1:k, tgt=k+1:2k)
  add_parts!(g, :Wire, k; tgt=1:k, src=k+1:2k)
  return g
end

meshpath(n::Int) = begin
    gt = @acset OpenCPortGraph begin
        Box = 1
        Port = 3
        Wire = 0
        OuterPort = 2
        box= ones(Int, 3)
        con= [3,2]
    end
    gm = @acset OpenCPortGraph begin
        Box = 1
        Port = 4
        Wire = 0
        OuterPort = 2
        box= ones(Int, 4)
        con= [4,2]
    end
    subs = [gt]
    for i in 2:n-1
        push!(subs, gm)
    end
    push!(subs, gt)
    X = coproduct(subs)
    for i in 1:n-1
        xi = subs[i]
        xj = subs[i+1]
        p = legs(X)[i][:Port](nparts(xi, :Port)-1)
        q = legs(X)[i+1][:Port](1)
        add_parts!(apex(X), :Wire, 2, src=[p,q], tgt=[q,p])
    end
    c₁ = apex(X)[1:2:nparts(apex(X),:OuterPort) ,:con]
    c₂ = apex(X)[2:2:nparts(apex(X),:OuterPort) ,:con]
    apex(X)[:,:con] = vcat(c₁,c₂)
    return X
end

function gridpath(n::Int, width::Int)
    node = @acset OpenCPortGraph begin
        Box = 1
        Port = 0
        Wire = 0
        box = 1
    end
    add_parts!(node, :Port, 2width, box=1)
    X = coproduct(collect(repeated(node, n)))
    L = legs(X)
    A = apex(X)
    for i in 1:n-1
        for j in 1:width
            s = L[i][:Port](j)
            t = L[i+1][:Port](j+width)
            add_part!(A, :Wire, src=s, tgt=t)
            add_part!(A, :Wire, src=t, tgt=s)
        end
    end
    upstream = L[1][:Port](width+1:2width)
    add_parts!(A, :OuterPort, width, con=upstream)
    downstream = L[end][:Port](1:width)
    add_parts!(A, :OuterPort, width, con=downstream)
    return X
end

grid(n::Int, m::Int) = ocompose(apex(gridpath(n,m)), collect(repeated(apex(meshpath(m)), n)))

end


using GNSSTools
using SatelliteToolbox
using ProgressMeter
using Makie, FileIO, Colors, GeometryBasics, GLMakie
using AbstractPlotting, AbstractPlotting.MakieLayout

GM = 3.986004418e14  # m³s⁻²
eop = get_eop();
user_lla = (40.01, -105.2437, 1655);

###########ONE-WEB############
a = Rₑ+1200*1000;  # meters
plane_num = 18
sat_per_plane = 36;
incl = 87.9;  # degrees
ΔΩ = 360/plane_num/2;
###########IRIDIUM############
a = Rₑ+780*1000;  # meters
plane_num = 6;
sat_per_plane = 11;
incl = 86;
ΔΩ = 360/plane_num/2;
###########STARLINK###########
# a = Rₑ+1150*1000;  # meters
# plane_num = 32;
# sat_per_plane = 50;
# incl = 53;
# ΔΩ = 360/plane_num;
a_starlink = Rₑ+1150*1000;  # meters
#############GPS##############
# a = Rₑ+20000*1000;
# plane_num = 6;
# sat_per_plane = 4;
# incl = 56;
# ΔΩ = 360/plane_num;
##############################
orbital_period = 2π*sqrt(a^3 / GM)  # seconds
orbital_period_starlink = 2π*sqrt(a_starlink^3 / GM)  # seconds
# orb_period_fraction = 2
# orb_period_fraction = (60*60*24) / orbital_period
orb_period_fraction = 1
frames_per_second = 25  # frames/second
total_orbit_time = orb_period_fraction*orbital_period
total_orbit_time_starlink = orb_period_fraction*orbital_period_starlink
total_vid_time = 10  # seconds
# Δt = total_orbit_time/total_vid_time/frames_per_second  # seconds
Δt = total_orbit_time_starlink/total_vid_time/frames_per_second  # seconds
t_range = Array(0:Δt:total_orbit_time);
freq = 9650e6;
constellation = define_constellation(a, plane_num,
                                     sat_per_plane, incl,
                                     t_range; obs_lla=user_lla,
                                     print_steps=true,
                                     eop=eop,
                                     ΔΩ=ΔΩ,
                                     show_plot=false);
# Loop through each time step and produce a constellation plot to be combined
# into a video later
obs_ecef = GeodetictoECEF(deg2rad(user_lla[1]), deg2rad(user_lla[2]), user_lla[3])
N = length(constellation.t_range)
sat_num = length(constellation.satellites)
# Earth
u = Array(range(0, 2π, length=200)) .+ (180*pi/180)
v = Array(range(0, 1π, length=200))
u, v = meshgrid(u, v)
E_x = Rₑ.*cos.(u).*sin.(v)
E_y = Rₑ.*sin.(u).*sin.(v)
E_z = Rₑ.*cos.(v)

# earth = load(download("https://svs.gsfc.nasa.gov/vis/a000000/a002900/a002915/bluemarble-2048.png"));
earth = load("bluemarble-2048.png");

# function make_plot(i, scene=Scene(resolution=(1000, 1000), color=:black);
#                    sat2user_vectors=true)
#     pt_size = 9e8
#     strokewidth = 4
#     linewidth = 3
#     ϕ = deg2rad(270)
#     θ = deg2rad(65)
#     scaler = 3
#     camera_pos = (scaler*constellation.a).*[cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)]
#     light = Vec{3,Float32}[[1.0, 1.0, 1.0], [0.1, 0.1, 0.1], [0.9, 0.9, 0.9], [0.0, -20.0, 0.0]]
#     to_eci = rECEFtoECI(ITRF(), TEME(), constellation.epoch+constellation.t_range[i], eop)
#     # Transform Earth and CO from ECEF to ECI frames
#     E_xt = zeros(size(E_x))
#     E_yt = zeros(size(E_y))
#     E_zt = zeros(size(E_z))
#     for j in 1:size(E_x)[1]
#         for k in 1:size(E_x)[2]
#             x, y, z = to_eci*[E_x[j,k], E_y[j,k], E_z[j,k]]
#             E_xt[j,k] = x
#             E_yt[j,k] = y
#             E_zt[j,k] = z
#         end
#     end
#     surface!(E_xt, E_yt, E_zt, color=earth, shading=false, show_axis=false)
#     obs_eci = to_eci*[obs_ecef[1], obs_ecef[2], obs_ecef[3]]
#     for j in 1:sat_num
#         r_ecef = constellation.satellites[j].r_ecef[i,:]
#         r_eci = kepler_to_sv(constellation.satellites[j].orbit[i]).r
#         sat_range, az, el = calcelevation(r_ecef, user_lla)
#         scatter!([r_eci[1]], [r_eci[2]], [r_eci[3]],
#                 strokewidth=strokewidth, strokecolor=:grey)
#         if el > 10
#             lines!([obs_eci[1], r_eci[1]],
#                    [obs_eci[2], r_eci[2]],
#                    [obs_eci[3], r_eci[3]], color=:red, linewidth=linewidth)
#         end
#     end
#     for j in 1:sat_num
#         if (j%sat_per_plane) == (sat_per_plane-1)
#             r_ecis = zeros(size(constellation.satellites[j].r_ecef))
#             for k in 1:size(r_ecis)[1]
#                 # r_ecis[k,:] = to_eci*constellation.satellites[j].r_ecef[k,:]
#                 r_ecis[k,:] =  kepler_to_sv(constellation.satellites[j].orbit[k]).r
#             end
#             lines!(r_ecis[:,1], r_ecis[:,2], r_ecis[:,3], color=:grey, linewidth=linewidth)
#         end
#     end
#     update_cam!(scene, camera_pos, 0)
#     return scene
# end

function rotate_earth(i)
    to_eci = rECEFtoECI(ITRF(), TEME(), constellation.epoch+constellation.t_range[i], eop)
    # Transform Earth and CO from ECEF to ECI frames
    E_xt = zeros(size(E_x))
    E_yt = zeros(size(E_y))
    E_zt = zeros(size(E_z))
    for j in 1:size(E_x)[1]
        for k in 1:size(E_x)[2]
            x, y, z = to_eci*[E_x[j,k], E_y[j,k], E_z[j,k]]
            E_xt[j,k] = x
            E_yt[j,k] = y
            E_zt[j,k] = z
        end
    end
    return (E_xt, E_yt, E_zt)
end

function get_obj_positions(i)
    to_eci = rECEFtoECI(ITRF(), TEME(), constellation.epoch+constellation.t_range[i], eop)
    obs_eci = to_eci*obs_ecef
    sat_pos = zeros(length(constellation.satellites), 3)
    planes = []
    visible = []
    visible_sat = []
    for j in 1:length(constellation.satellites)
        r_ecef = constellation.satellites[j].r_ecef[i,:]
        r_eci = kepler_to_sv(constellation.satellites[j].orbit[i]).r
        sat_pos[j,:] = r_eci
        sat_range, az, el = calcelevation(r_ecef, user_lla)
        if el > 10
            push!(visible, [[obs_eci[1], r_eci[1]], [obs_eci[2], r_eci[2]], [obs_eci[3], r_eci[3]]])
            push!(visible_sat, r_eci)
        end
        if (constellation.satellites[j].id%sat_per_plane) == 1
            r_ecis = zeros(size(constellation.satellites[j].r_ecef))
            for k in 1:size(r_ecis)[1]
                r_ecis[k,:] =  kepler_to_sv(constellation.satellites[j].orbit[k]).r
            end
            push!(planes, r_ecis)
        end
    end
    visible_sats = zeros(length(sat_pos),3)
    for j in 1:size(sat_pos)[1]
        if j > length(visible_sat)
            visible_sats[j,:] = [NaN, NaN, NaN]
        else
            visible_sats[j,:] = visible_sat[j]
        end
    end
    return (sat_pos, planes, visible, visible_sats)
end

function initialize_figure(i; show_vectors=true, visible_color=missing)
    ϕ = deg2rad(180)
    θ = deg2rad(70)
    scaler = 3
    camera_pos = (scaler*constellation.a).*[cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)]
    light_ϕ = deg2rad(270)
    light_θ = deg2rad(90)
    light_scaler = 149597870700/Rₑ
    light_camera_pos = (light_scaler*constellation.a).*[cos(light_ϕ)*sin(light_θ), sin(light_ϕ)*sin(light_θ), cos(light_θ)]
    lightpos = Vec3f0(light_camera_pos...)
    pt_size = 9e8
    strokewidth = 3
    linewidth = 2
    linestyle = nothing
    pt_color = :lightgrey
    line_color = :lightgrey
    vector_color = :lightblue
    background_color = :white
    if ismissing(visible_color)
        visible_color = pt_color
    end
    E_xt, E_yt, E_zt = rotate_earth(i)
    sat_pos, planes, visible, visible_sats = get_obj_positions(i)
    scene = Scene(resolution=(1000, 1000), backgroundcolor=background_color)
    surf = surface!(scene, E_xt, E_yt, E_zt, shading=true, color=earth,
                    show_axis=false, lightposition=lightpos,
                    colormap=(:white, :white),
                    ambient = Vec3f0(0.5, 0.5, 0.5),
                    diffuse = Vec3f0(0.4, 0.4, 0.2),
                    specular = Vec3f0(0.5, 0.5, 0.5),
                    shininess = 0.5f0)
    scatter_sat = scatter!(scene, sat_pos[:,1], sat_pos[:,2],
    sat_pos[:,3],
                           strokewidth=strokewidth, strokecolor=pt_color)
    scatter_visible_sat = scatter!(scene,
                                   visible_sats[:,1],
                                   visible_sats[:,2],
                                   visible_sats[:,3],
                                   strokewidth=strokewidth,
                                   strokecolor=visible_color)
    for j in 1:length(planes)
        line_planes = lines!(scene, planes[j][:,1], planes[j][:,2],
                             planes[j][:,3],
                             color=line_color, linewidth=linewidth,
                             linestyle=linestyle)
    end
    if show_vectors
        val = [NaN, NaN]
        for j in 1:ceil(Int, length(constellation.satellites)/2)
            line_visible = lines!(scene, val, val,  val,
                                  color=vector_color, linewidth=linewidth)
        end
        for j in 1:length(visible)
            scene.plots[j+length(planes)+3][1] = visible[j][1]
            scene.plots[j+length(planes)+3][2] = visible[j][2]
            scene.plots[j+length(planes)+3][3] = visible[j][3]
        end
    end
    update_cam!(scene, camera_pos, 0)
    return scene
end

function modify_figure!(scene, i; show_vectors=true)
    ϕ = deg2rad(180)
    θ = deg2rad(70)
    scaler = 3
    camera_pos = (scaler*constellation.a).*[cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)]
    E_xt, E_yt, E_zt = rotate_earth(i)
    sat_pos, planes, visible, visible_sats = get_obj_positions(i)
    scene.plots[1][1] = E_xt
    scene.plots[1][2] = E_yt
    scene.plots[1][3] = E_zt
    scene.plots[2][1] = sat_pos[:,1]
    scene.plots[2][2] = sat_pos[:,2]
    scene.plots[2][3] = sat_pos[:,3]
    scene.plots[3][1] = visible_sats[:,1]
    scene.plots[3][2] = visible_sats[:,2]
    scene.plots[3][3] = visible_sats[:,3]
    for j in 1:length(planes)
        scene.plots[j+3][1] = planes[j][:,1]
        scene.plots[j+3][2] = planes[j][:,2]
        scene.plots[j+3][3] = planes[j][:,3]
    end
    if show_vectors
        val = [NaN, NaN]
        for j in 1:(length(scene.plots)-3-length(planes))
            scene.plots[j+length(planes)+3][1] = val
            scene.plots[j+length(planes)+3][2] = val
            scene.plots[j+length(planes)+3][3] = val

        end
        for j in 1:length(visible)
            scene.plots[j+length(planes)+3][1] = visible[j][1]
            scene.plots[j+length(planes)+3][2] = visible[j][2]
            scene.plots[j+length(planes)+3][3] = visible[j][3]
        end
    end
    update_cam!(scene, camera_pos, 0)
    return scene
end

function make_plots(N; show_vectors=true, visible_color=missing)
    scene = initialize_figure(1; show_vectors=show_vectors,
                              visible_color=visible_color)
    p = Progress(N, 1, "Generating $N frames...")
    compression = 10
    record(scene, "gps_constellation_black_compressed.mp4", 1:N; sleep=false,
           compression=compression, framerate=frames_per_second) do i
        modify_figure!(scene, i; show_vectors=show_vectors)
        next!(p)
    end
end

function zoom_earth(N=11*frames_per_second, start_i=9*frames_per_second,
                    end_i=10*frames_per_second;
                    show_vectors=true, visible_color=missing)
    a1 = Rₑ+20000*1000
    a2 = Rₑ+1150*1000;
    ϕ = deg2rad(180)
    θ = deg2rad(70)
    scaler = 3
    camera_pos = (scaler*a1).*[cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)]
    light_ϕ = deg2rad(270)
    light_θ = deg2rad(90)
    light_scaler = 149597870700/Rₑ
    light_camera_pos = (light_scaler*a1).*[cos(light_ϕ)*sin(light_θ), sin(light_ϕ)*sin(light_θ), cos(light_θ)]
    lightpos = Vec3f0(light_camera_pos...)
    E_xt, E_yt, E_zt = rotate_earth(1)
    scene = Scene(resolution=(1000, 1000), backgroundcolor=:black)
    surf = surface!(scene, E_xt, E_yt, E_zt, shading=true, color=earth,
                    show_axis=false, lightposition=lightpos,
                    colormap=(:white, :white),
                    ambient = Vec3f0(0.5, 0.5, 0.5),
                    diffuse = Vec3f0(0.4, 0.4, 0.2),
                    specular = Vec3f0(0.5, 0.5, 0.5),
                    shininess = 0.5f0)
    as = fill(a1, N)
    Δas = (a2 - a1)
    for i in start_i:end_i
        as[i] = a1 - Δas*(cos((i - start_i) * π / (end_i - start_i))/2 - 0.5)
    end
    as[end_i+1:end] .= a2
    p = Progress(N, 1, "Generating $N frames...")
    compression = 0
    record(scene, "zoom_earth.mp4", 1:N; sleep=false,
           compression=compression, framerate=frames_per_second) do i
        E_xt, E_yt, E_zt = rotate_earth(i)
        scene.plots[1][1] = E_xt
        scene.plots[1][2] = E_yt
        scene.plots[1][3] = E_zt
        camera_pos = (scaler*as[i]).*[cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)]
        update_cam!(scene, camera_pos, 0)
        next!(p)
    end
end

function make_hist_plots()
    # GPS
    a = Rₑ+20000*1000;
    plane_num = 6;
    sat_per_plane = 4;
    incl = 56;
    ΔΩ = 360/plane_num;
    orbital_period = 2π*sqrt(a^3 / GM)  # seconds
    # orb_period_fraction = 2
    # orb_period_fraction = (60*60*24) / orbital_period
    orb_period_fraction = 10
    frames_per_second = 1  # frames/second
    total_orbit_time = orb_period_fraction*orbital_period
    total_vid_time = total_orbit_time
    # Δt = total_orbit_time/total_vid_time/frames_per_second  # seconds
    Δt = total_orbit_time/total_vid_time/frames_per_second  # seconds
    t_range = Array(0:Δt:total_orbit_time);
    # freq = 9650e6;
    freq = L1_freq
    doppler, doppler_rate, doppler_bounds,
    doppler_rate_bounds = doppler_distribution(a, plane_num,
                                               sat_per_plane, incl,
                                               t_range, user_lla,
                                               freq, show_plot=false,
                                               print_steps=true,
                                               eop=eop,
                                               ΔΩ=ΔΩ);
    # Make plot
    scene, layout = layoutscene(resolution=(840, 360))
    ax1 = layout[1, 1] = LAxis(scene, xlabel="Doppler Frequency (kHz)")
    hideydecorations!(ax1, label=true, ticklabels=true, ticks=true, grid=true)
    hidexdecorations!(ax1, label=false, ticklabels=false, ticks=false, grid=true)
    barplot!(ax1, doppler.edges[1][2:end]./1000, doppler.weights, color=:black,
            strokewidth=1, strokecolor=:black, show_axis=true)
    save("gps_hist_doppler.png", scene)
    scene, layout = layoutscene(resolution=(840, 360))
    ax1 = layout[1, 1] = LAxis(scene, xlabel = "Doppler Frequency Rate (Hz/s)")
    hideydecorations!(ax1, label=true, ticklabels=true, ticks=true, grid=true)
    hidexdecorations!(ax1, label=false, ticklabels=false, ticks=false, grid=true)
    barplot!(ax1, doppler_rate.edges[1][2:end], doppler_rate.weights, color=:black,
            strokewidth=1, strokecolor=:black, show_axis=true)
    save("gps_hist_doppler_rate.png", scene)
    # Starlink
    a = Rₑ+1150*1000;  # meters
    plane_num = 32;
    sat_per_plane = 50;
    incl = 53;
    ΔΩ = 360/plane_num;
    orbital_period = 2π*sqrt(a^3 / GM)  # seconds
    # orb_period_fraction = 2
    # orb_period_fraction = (60*60*24) / orbital_period
    orb_period_fraction = 10
    frames_per_second = 1  # frames/second
    total_orbit_time = orb_period_fraction*orbital_period
    total_vid_time = total_orbit_time
    # Δt = total_orbit_time/total_vid_time/frames_per_second  # seconds
    Δt = total_orbit_time/total_vid_time/frames_per_second  # seconds
    t_range = Array(0:Δt:total_orbit_time);
    # freq = 9650e6;
    freq = L1_freq
    doppler, doppler_rate, doppler_bounds,
    doppler_rate_bounds = doppler_distribution(a, plane_num,
                                               sat_per_plane, incl,
                                               t_range, user_lla,
                                               freq, show_plot=false,
                                               print_steps=true,
                                               eop=eop,
                                               ΔΩ=ΔΩ);
    # Make plot
    scene, layout = layoutscene(resolution=(840, 360))
    ax1 = layout[1, 1] = LAxis(scene, xlabel="Doppler Frequency (kHz)")
    hideydecorations!(ax1, label=true, ticklabels=true, ticks=true, grid=true)
    hidexdecorations!(ax1, label=false, ticklabels=false, ticks=false, grid=true)
    barplot!(ax1, doppler.edges[1][2:end]./1000, doppler.weights, color=:black,
            strokewidth=1, strokecolor=:black, show_axis=true)
    save("starlink_hist_doppler.png", scene)
    scene, layout = layoutscene(resolution=(840, 360))
    ax1 = layout[1, 1] = LAxis(scene, xlabel = "Doppler Frequency Rate (Hz/s)")
    hideydecorations!(ax1, label=true, ticklabels=true, ticks=true, grid=true)
    hidexdecorations!(ax1, label=false, ticklabels=false, ticks=false, grid=true)
    barplot!(ax1, doppler_rate.edges[1][2:end], doppler_rate.weights, color=:black,
            strokewidth=1, strokecolor=:black, show_axis=true)
    ax1.xticks = -200:50:0
    save("starlink_hist_doppler_rate.png", scene)
end
# run(`ffmpeg -r $frames_per_second -f image2 -i constellation_figures/constellation_%0$(pad_num)d.png -y constellation.mp4`)
run(`ffmpeg -i LabVideo_Sergei.mp4 -filter_complex "[0:v] palettegen" -y palette.png`)
run(`ffmpeg -i LabVideo_Sergei.mp4 -i palette.png -filter_complex "[0:v][1:v] paletteuse" -y LabVideo_Sergei.gif`)

# L1 C/A signal undergoing extreme doppler
using PyPlot
pygui(false)
t_length = 2
sig_freq = 9650e6;
sigtype = define_l1ca_code_type(t_length; sig_freq=sig_freq);
doppler_t = Array(0:0.001:t_length+0.1);
f_d = 150*1e3
fd_rate = -1000
fdd_rate = 175
fddd_rate = 0
doppler_curve = f_d .+ fd_rate .* doppler_t .+ fdd_rate .* doppler_t.^2 .+ fddd_rate .* doppler_t.^3;
signal = definesignal(sigtype, 5e6, t_length);
generatesignal!(signal; doppler_curve=doppler_curve, doppler_t=doppler_t);
acqresults, trackresults = process(signal, sigtype, 1; fd_center=f_d, show_plot=false);
# Iridium constellation @ X-band
t_range = Array(0:1:12*60*60);
a = Rₑ+780*1000;
plane_num = 6;
sat_per_plane = 11;
incl = 86;
freq = 9650e6;
ΔΩ = 30;
doppler, doppler_rate, doppler_bounds,
doppler_rate_bounds = doppler_distribution(a, plane_num,
                                           sat_per_plane, incl,
                                           t_range, user_lla,
                                           freq,
                                           print_steps=true,
                                           eop=eop,
                                           ΔΩ=ΔΩ;
                                           show_plot=false);
# Make animated plots
figsize = (17,8)
frames_per_second = 25
total_length = 10
N = floor(Int, frames_per_second*t_length)
pad_num = length(string(N))
ΔN = floor(Int, length(trackresults.t)/N)
p = Progress(N, 1, "Generating $N frames...")
for i in 0:N
    idx = i*ΔN
    fig = figure(figsize=figsize)
    ax1 = subplot(221)
    ax1.plot(trackresults.t[1:idx], trackresults.n0[1:idx].%1023, "k")
    xlim([trackresults.t[1], trackresults.t[end]])
    ylim([minimum(trackresults.n0.%1023), maximum(trackresults.n0.%1023)])
    title("Code Phase", fontsize=20)
    xticks([])
    yticks([])
    ax2 = subplot(222)
    ax2.plot(trackresults.t[1:idx], trackresults.dphi_meas[1:idx], "k")
    xlim([trackresults.t[1], trackresults.t[end]])
    ylim([minimum(trackresults.dphi_meas), maximum(trackresults.dphi_meas)])
    title("Phase Residuals", fontsize=20)
    xticks([])
    yticks([])
    ax3 = subplot(223)
    ax3.plot(trackresults.t[1:idx], trackresults.fds[1:idx], "k")
    xlim([trackresults.t[1], trackresults.t[end]])
    ylim([minimum(trackresults.fds), maximum(trackresults.fds)])
    title("Doppler", fontsize=20)
    xticks([])
    yticks([])
    ax4 = subplot(224)
    ax4.plot(trackresults.t[1:idx], real.(trackresults.ZP[1:idx]), "k")
    xlim([trackresults.t[1], trackresults.t[end]])
    ylim([minimum(real.(trackresults.ZP)), maximum(real.(trackresults.ZP))])
    # ax4.plot(trackresults.t, trackresults.data_bits, "k", linestyle="-")
    title("Navigation Data", fontsize=20)
    xticks([])
    yticks([])
    subplots_adjust(top=0.96, bottom=0.004, right=0.998, left=0.001,
                    hspace=0.2, wspace=0.1)
    # margins(0,0,0)
    # gca().xaxis.set_major_locator(PyPlot.plt.NullLocator())
    # gca().yaxis.set_major_locator(PyPlot.plt.NullLocator())
    if i == N
        num = string(0, pad=pad_num)
    else
        num = string(i+1, pad=pad_num)
    end
    savefig("data_figures/data_$num.png")
    close(fig)
    next!(p)
end
run(`ffmpeg -r $frames_per_second -f image2 -i data_figures/data_%0$(pad_num)d.png -y data.mp4`)
run(`ffmpeg -i data.mp4 -filter_complex "[0:v] palettegen" -y palette1.png`)
run(`ffmpeg -i data.mp4 -i palette1.png -filter_complex "[0:v][1:v] paletteuse" -y data.gif`)


@define_scenario SimpleScenario = begin
    q₁::Float64
    q₂::Float64
    d₁::Float64
    d₂::Float64
end

ξ₁ = SimpleScenario(24.0, 28.0, 500.0, 100.0, probability = 0.4)
ξ₂ = SimpleScenario(28.0, 32.0, 300.0, 300.0, probability = 0.6)

simple_sp = StochasticProgram([ξ₁, ξ₂], Deterministic(), GLPK.Optimizer)

@first_stage simple_sp = begin
    @decision(model, x₁ >= 40)
    @decision(model, x₂ >= 20)
    @objective(model, Min, 100*x₁ + 150*x₂)
    @constraint(model, x₁ + x₂ <= 120)
end

@second_stage simple_sp = begin
    @known x₁ x₂
    @uncertain q₁ q₂ d₁ d₂ from SimpleScenario
    @recourse(model, 0 <= y₁ <= d₁)
    @recourse(model, 0 <= y₂ <= d₂)
    @objective(model, Max, q₁*y₁ + q₂*y₂)
    @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
    @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
end