## Assumptions
## http://docs.sympy.org/0.7.2/modules/assumptions/index.html

"""
refine: http://docs.sympy.org/dev/modules/assumptions/refine.html
"""
refine(ex, assumpts...) = sympy.refine(ex, assumpts...)
export refine

"""

`ask`. Returns true, false or nothing

Example:

```
ask(Q.integer(x*y), And(Q.integer(x), Q.integer(y)))
## really slow isprime:
filter(x -> ask(Q.prime(x)), [1:1000])
```

"""
ask(x::Sym, args...) = sympy.ask(x, args...)
ask(x::Bool, args...) = x
ask(x::Nothing, args...) = x
export ask

## should we support & and | for (sym,sym) pairs? Not sure
## dependso on what happens to | (x==0) and ex == x usage for ==
## for now, we can combine terms logically with And, Or, Not...
## these are in logic module

## We make a module Q to hold the assumptions
## this follows this page http://docs.sympy.org/0.7.5/_modules/sympy/assumptions/ask.html
"""

Documentation for the `Q` module.

SymPy allows for
[assumptions](http://docs.sympy.org/latest/modules/sympy/assumptions/)
on variables. These may be placed on free sympols at construction.

For example, the following creates a real value variable `x` and a postive, real variable `y`:

```
@vars x real=true
@vars y real=true positive=true
```

The `Q` module exposes a means to *q*uery the assumptions on a
variable. For example,

```
ask(Q.positive(y))  # true
ask(Q.negative(y))  # false
ask(Q.positive(x))  # `nothing`
ask(Q.positive(x^2)) # `nothing` -- might be 0
ask(Q.positive(1 + x^2) # true  -- must be postive now.
```

The ask function uses tri-state logic, returning one of 3 values:
`true`; `false`; or `nothing`, when the query is indeterminate.

The construction of predicates is done through `Q` methods. These can
be combined logically. For example, this will be `true`:

```
ask(Q.positive(y) & Q.negative(-x^2 - 1))
```

The above use `&` as an infix operation for the binary operator
`And`. Values can also be combined with `Or`, `Not`, `Xor`, `Nand`,
`Nor`, `Implies`, `Equivalent`, and `satisfiable`.

Note: As `SymPy.jl` converts symbolic matrices into Julia's `Array`
type and not as matrices within Python, the predicate functions from SymPy for
matrices are not used, though a replacement is given.
"""
module Q
import SymPy
import PyCall
import LinearAlgebra: det, norm

##http://docs.sympy.org/dev/_modules/sympy/assumptions/ask.html#ask
Q_predicates = (:antihermitian,
                :bounded, :finite, # bounded deprecated
                :commutative,
                :complex,
                :composite,
                :even,
                :extended_real,
                :hermitian,
                :imaginary,
                :infinitesimal,
                :infinity, :infinite, # :infinity deprecated
                :integer,
                :irrational,
                :rational,
                :algebraic,
                :transcendental,
                :negative,
                :nonzero, :zero,
                :positive,
                :prime,
                :real,
                :odd,
                :is_true,
                :nonpositive,
                :nonnegative
#                :symmetric,
#                :invertible,
#                :singular,
#                :orthogonal,
#                :unitary,
#                :normal,
#                :positive_definite,
#                :upper_triangular,
#                :lower_triangular,
#                :diagonal,
#                :triangular,
#                :unit_triangular,
#                :fullrank,
#                :square,
#                :real_elements,
#                :complex_elements,
#                :integer_elements
)

for meth in Q_predicates
   nm = string(meth)
      @eval begin
            ($meth)(x) = PyCall.pycall(SymPy.sympy.Q.$nm, SymPy.Sym, x)::SymPy.Sym
   end
end

symmetric(M::Array{T,2}) where {T <: SymPy.Sym} = SymPy.issymmetric(M)
function invertible(M::Array{T,2}) where {T <: SymPy.Sym}
    d = det(M)
    pos = SymPy.ask(positive(d))
    if pos == nothing
        return nothing
    elseif pos == true
        return true
    end
    neg = SymPy.ask(negative(d))
    if neg == nothing
        return nothing
    end
    z = SymPy.ask(zero(d))
    if z == nothing
        return nothing
    elseif z == true
        return false
    end

    return true
end

function singular(M::Array{T,2}) where {T <: SymPy.Sym}
    !invertible(M)
end

function orthogonal(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = SymPy.simplify.(SymPy.simplify.(M*transpose(M)) .== one(T))
    no_nothing = 0
    for val in vals
        a = SymPy.ask(zero(val))
        if a == nothing
            no_nothing += 1
        elseif a == false
            return false
        end
    end

    no_nothing > 0 && return nothing
    return true
end

function unitary(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = SymPy.simplify.(SymPy.simplify.(M*ctranspose(M)) .== one(T))
    no_nothing = 0
    for val in vals
        a = SymPy.ask(zero(val))
        if a == nothing
            no_nothing += 1
        elseif a == false
            return false
        end
    end

    no_nothing > 0 && return nothing
    return true
end

function normal(M::Array{T,2}) where {T <: SymPy.Sym}
    lhs = ctranspose(M) * M
    rhs = M * ctranspose(M)
    vals = zero.(SymPy.simplify.(lhs - rhs))
    no_nothing = 0
    for val in vals
        a = SymPy.ask(val)
        if a == nothing
            no_nothing += 1
        elseif a == false
            return false
        end
    end

    no_nothing > 0 && return nothing
    return true
end

# Use [Sylvester's](https://en.wikipedia.org/wiki/Sylvester%27s_criterion) Criteria
function positive_definite(M::Array{T,2}) where {T <: SymPy.Sym}
    !SymPy.ask(square(M)) && return false
    !SymPy.ask(symmetric(M))  && return false
    m, n = size(M)
    no_false = 0
    no_nothing = 0
    for i in 1:m
        a = SymPy.ask(Q.positive(det(M[1:i, 1:i])))
        if a == nothing no_nothing += 1 end
        if a == false no_false += 1 end
    end
    if no_false > 0
        return false
    elseif no_nothing > 0
        return nothing
    else
        return true
    end
end

upper_triangular(M::Array{T,2}) where {T <: SymPy.Sym} = SymPy.istriu(M)
lower_triangular(M::Array{T,2}) where {T <: SymPy.Sym} = SymPy.istril(M)
diagonal(M::Array{T,2}) where {T <: SymPy.Sym} = upper_triangular(M) && lower_triangular(M)
triangular(M::Array{T,2}) where {T <: SymPy.Sym} = upper_triangular(M) || lower_triangular(M)

## This is likely not the best way as it is a bit fidgety due
## to the call to rref.
function full_rank(M::Array{T,2}) where {T <: SymPy.Sym}
    m,n = size(M)
    m <= n || return full_rank(transpose(M))

    rr, p = SymPy.rref(M)
    lr = rr[end, :] # is this zero?
    no_nothing = 0
    no_nonzero = 0
    for val in lr
        a = SymPy.ask(nonzero(val))
        if a == nothing
            no_nothing += 1
        end
        if a == true
            no_nonzero += 1
        end
    end
    if no_nothing > 0
        return nothing
    elseif no_nonzero == 0
        return false
    else
        return true
    end

end

function square(M::Array{T,2}) where {T <: SymPy.Sym}
    m,n = SymPy.size(M)
    m == n
end

function real_elements(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = real.(M)
    for val in vals
        a = SymPy.ask(real(val))
        (a == nothing || a == false) && return false
    end
    return true
end

function complex_elements(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = real.(M)
    for val in vals
        a = SymPy.ask(SymPy.sympy."Q".complex(val))
        (a == nothing || a == false) && return false
    end
    return true

end

function integer_elements(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = real.(M)
    for val in vals
        a = SymPy.ask(integer(val))
        (a == nothing || a == false) && return false
    end
    return true
end

end

## export
export Q


struct MLS{AT} <: IRMeasurement
    sig::AT
    reps::Int
    gain::Float64
    prepad::Int
end

MLS(sig::AT, reps, gain, prepad) where AT = MLS{AT}(sig, reps, gain, prepad)

prepadding(m::MLS) = m.prepad

"""
    mls(L, reps=2; options...)

Generate a maximum-length sequence (MLS), also known as Schroeder's method.
Actual length is `L-1`. The sequence will be repeated `reps` times, which should
be at least 2, as only the last `reps-1` repetitions are used.

## Options:
$optiondoc_gain
$optiondoc_prepad
"""
function mls(len, reps=2; gain=mls_gain, prepad=default_prepad, samplerate=1)
    N = Int(log2(nextpow(2, durnorm(len, samplerate))))
    1 <= N <= length(mlspolys) || throw(ArgumentError("N must be positive and <= $(length(mlspolys))"))
    poly = mlspolys[N]
    # where we'll keep the state used to compute the next value. The
    # rightmost bit is the most recent value. Nothing super special about this
    # initialization value, as long as there's at least 1 nonzero bit in the
    # filter span
    register = 0x0000000000000001
    # bit mask for the filter taps
    taps = reduce(|, (0x0000000000000001 << (N-1-shift) for shift in poly))
    out = zeros(2^N-1)
    # initialize our output array to match the register above - with only the
    # most recent sample as 1.
    out[1:N-1] .= -1
    out[N] = 1
    for i in N+1:length(out)
        nextval = UInt64(count_ones(register & taps)) & 0x0000000000000001
        register = (register << 1) | nextval
        out[i] = float(nextval) * 2 - 1
    end

    MLS(out, reps, gain, durnorm(prepad, samplerate))
end

stimulus(m::MLS) = [zeros(m.prepad); repeat(m.sig, m.reps)*m.gain]

"""
Perform circular cross-correlation, assuming `u` and `v` are equal length, but
allowing for the possibility that their length is not amenable to FFT
processing.

Uses algorithm from:

Daigle, J. N., & Xiang, N. (2006). A specialized fast cross-correlation for
acoustical measurements using coded sequences. The Journal of the Acoustical
Society of America, 119(1), 330–335. https://doi.org/10.1121/1.2141236

"""
function circxcorr(u, v)
    L = length(u)
    @assert L == length(v)

    if L == nextfastfft(L)
        _circxcorr(u,v)
    else
        N = nextfastfft(2L)
        u = DSP._zeropad(u, N)
        v = DSP._zeropad(v, N)
        xc = _circxcorr(u,v)
        xc[1:L] + xc[end-L+1:end]
    end
end

"""
Perform circular cross-correlation, assuming `u` and `v` are vectors of equal
length and that length is amenable to FFT processing.
"""
_circxcorr(u, v) = irfft(rfft(u) .* conj.(rfft(v)), length(u))

function _analyze(m::MLS, response::AbstractArray)
    if m.reps < 2
        throw(ArgumentError("Can't perform MLS analysis with only one repetition"))
    end
    L = length(m.sig)
    avgresp = zeros(L, Base.tail(size(response))...)
    for rep in 1:m.reps-1
        avgresp .+= timeslice(response, (1:L) .+ (m.prepad + rep*L))
    end
    avgresp ./= (m.reps-1)
    mapslices(avgresp, dims=1) do v
        # compensate for the amplitude drop
        circxcorr(v, m.sig) ./ L / m.gain
    end
end

# these polynomial coefficients represent the tap indices of the filter, where
# zero is the oldest value (delay of N)
# they are from:
# Stahnke, Wayne 1973. Primitive Binary Polynomials.
# Mathematics of Computation 27(124): 977–980.
# the paper has higher orders, but this seems like more than enough for
# reasonable audio purposes. Feel free to submit a PR with more if they are
# useful to you.

const mlspolys = [
    [0],            # 01
    [1,  0],        # 02
    [1,  0],        # 03
    [1,  0],        # 04
    [2,  0],        # 05
    [1,  0],        # 06
    [1,  0],        # 07
    [6,  5,  1, 0], # 08
    [4,  0],        # 09
    [3,  0],        # 10
    [2,  0],        # 11
    [7,  4,  3, 0], # 12
    [4,  3,  1, 0], # 13
    [12, 11, 1, 0], # 14
    [1,  0],        # 15
    [5,  3,  2, 0], # 16
    [3,  0],        # 17
    [7,  0],        # 18
    [6,  5,  1, 0], # 19
    [3,  0],        # 20
    [2,  0],        # 21
    [1,  0],        # 22
    [5,  0],        # 23
    [4,  3,  1, 0], # 24
    [3,  0],        # 25
    [8,  7,  1, 0], # 26
    [8,  7,  1, 0], # 27
    [3,  0],        # 28
    [2,  0],        # 29
    [16, 15, 1, 0]  # 30
]


n = rand(1:10)
@test matrixdepot("cauchy", n) == matrixdepot("cauchy", Float64[1:n;])
x = ones(Float64, n)
y = x + 2
A = zeros(Float64, n, n)
fill!(A, 0.25)
@test_approx_eq matrixdepot("cauchy", x, y) A
println("'cauchy' passed test...")


# `topological_sort`
let
    # simple graph
    G = DiGraph(4)
    add_edge!(G,1,3)
    add_edge!(G,1,4)
    add_edge!(G,2,3)
    add_edge!(G,3,4)
    ordering = topological_sort(G)
    @test ordering == [1,2,3,4]
end
let
    # disjoint graph
    G = DiGraph(4)
    add_edge!(G,1,2)
    add_edge!(G,3,4)
    ordering = topological_sort(G)
    @test length(ordering) == nv(G)
    @test findfirst(ordering .== 1) < findfirst(ordering .== 2)
    @test findfirst(ordering .== 3) < findfirst(ordering .== 4)
end
let
    # cyclic graph (should throw error)
    G = DiGraph(3)
    add_edge!(G,1,2)
    add_edge!(G,2,3)
    add_edge!(G,3,1)
    @test_throws AssertionError topological_sort(G)
end
# `find_index_in_sorted_array` and `insert_to_sorted_array`
let
    array = [1.0,2.0,3.0,4.0,5.0]
    array = insert_to_sorted_array!(array, 3.5)
    @test(array[4] == 3.5)
    array = insert_to_sorted_array!(array, 0.5)
    @test(array[1] == 0.5)
    array = insert_to_sorted_array!(array, 5.5)
    @test(array[end] == 5.5)
end
let
    v = [0,1,2,4,5,6]
    L = length(v)
    x = 3
    @test find_index_in_sorted_array(v,x) == 4
    insert_to_sorted_array!(v,x)
    @test length(v) == L + 1
    v[4] == x
end
let
    v = [0,1,1,1,1,2]
    L = length(v)
    x = 1
    @test find_index_in_sorted_array(v,x) == 2
end
let
    v = Vector{Int}()
    insert_to_sorted_array!(v,1)
    @test length(v) == 1
end


function rate_plot(sepp::SEPP; step = nothing, from_idx = nothing, to_idx = nothing)

    ts = sepp.data
    isnothing(ts) && error("No data in the model, can't plot")

    times = ts.times
    isnothing(from_idx) && (from_idx = 1)
    isnothing(to_idx) && (to_idx = length(times))
    starttime = times[from_idx]
    endtime = times[to_idx]

    ((starttime isa TimeType) && !(step isa Union{TimeType, Nothing})) && (@warn "If times are TimeType step must be TimeType, ignoring step value."; step=nothing)

    if isnothing(step)
        if starttime isa TimeType
            anytimes = DateTime(starttime):Dates.Hour(1):DateTime(endtime)
        else
            anytimes = (starttime:oneunit(starttime-endtime):endtime)
        end
    else
        anytimes = (starttime:step:endtime)
    end

    μ = sepp.μ
    ϕ = sepp.ϕ
    γ = sepp.γ
    δ = sepp isa SEMPPExpKern ? sepp.δ : 0

    lamb = μ .+ ϕ .* volfunc(anytimes, ts, γ, δ)

    points = fill(0, (to_idx-from_idx + 1,))

    rate_layer = layer(x = anytimes, y = lamb, color = [color("black")], Geom.line)

    points_layer = layer(x = times[from_idx:to_idx], y = points, alpha = [0.7], Geom.point)

    plt = plot(rate_layer, points_layer, Guide.xlabel("time"), Guide.ylabel("intensity"), Theme(grid_line_width=0mm, highlight_width = 0mm))

    return plt
end

function marks_plot(mts::MarkedTimeSeries; from_idx::Union{Nothing, Integer} = nothing, to_idx::Union{Nothing, Integer} = nothing)
    marks = mts.marks
    times = mts.times

    isnothing(from_idx) && (from_idx = 1)
    isnothing(to_idx) && (to_idx = length(times))

    plt = plot(x=times[from_idx:to_idx], y=marks[from_idx:to_idx], Geom.hair, Geom.point, Guide.xlabel("time"), Guide.ylabel("marks"))

    return plt
end

function marks_plot(sempp::SEMPPExpKern)
    mts = sempp.data
    isnothing(mts) && error("No data in the model, can't plot")

    return marks_plot(mts)
end

function marked_rate_plot(sempp::SEMPPExpKern; step = nothing, from_idx = nothing, to_idx = nothing)

    mts = sempp.data
    isnothing(mts) && error("No data in the model, can't plot")

    rate_plt = rate_plot(sempp, step = step, from_idx = from_idx, to_idx = to_idx)

    marks_plt = marks_plot(mts, from_idx = from_idx, to_idx = to_idx)

    plt = vstack(rate_plt, marks_plt)

    return plt
end


include("../src/components/co2cycle_component.jl")

@testset "co2cycle" begin

    Precision = 1.0e-11
    T = length(MimiDICE2010.model_years)
    f = readxlsx(joinpath(@__DIR__, "..", "data", "DICE2010_082710d.xlsx"))

    m = Model()

    set_dimension!(m, :time, MimiDICE2010.model_years)

    add_comp!(m, co2cycle, :co2cycle)

# Set the parameters that would normally be internal connection from their Excel values
    set_param!(m, :co2cycle, :E, read_params(f, "B109:BI109", T))

# Load the rest of the external parameters
    p = dice2010_excel_parameters(joinpath(@__DIR__, "..", "data", "DICE2010_082710d.xlsx"))
    set_param!(m, :co2cycle, :mat0, p[:mat0])
    set_param!(m, :co2cycle, :mat1, p[:mat1])
    set_param!(m, :co2cycle, :mu0, p[:mu0])
    set_param!(m, :co2cycle, :ml0, p[:ml0])
    set_param!(m, :co2cycle, :b12, p[:b12])
    set_param!(m, :co2cycle, :b23, p[:b23])
    set_param!(m, :co2cycle, :b11, p[:b11])
    set_param!(m, :co2cycle, :b21, p[:b21])
    set_param!(m, :co2cycle, :b22, p[:b22])
    set_param!(m, :co2cycle, :b32, p[:b32])
    set_param!(m, :co2cycle, :b33, p[:b33])

# Run the one-component model
    run(m)

# Extract the generated variables
    MAT     = m[:co2cycle, :MAT]
    MAT_final   = m[:co2cycle, :MAT_final]
    ML      = m[:co2cycle, :ML]
    MU      = m[:co2cycle, :MU]

# Extract the true values
    True_MAT    = read_params(f, "B112:BI112", T)
    True_MAT_final  = read_params(f, "BJ112")
    True_ML     = read_params(f, "B115:BI115", T)
    True_MU     = read_params(f, "B114:BI114", T)

# Test that the values are the same
    @test maximum(abs, MAT .- True_MAT) ≈ 0. atol = Precision
    @test abs(MAT_final - True_MAT_final) ≈ 0. atol = Precision
    @test maximum(abs, ML .- True_ML) ≈ 0. atol = Precision
    @test maximum(abs, MU .- True_MU) ≈ 0. atol = Precision

end


"Abstract vector space type."
abstract type AbstractVectorSpace end

counter = 0

"""
    VectorSpace(dim::Union{Int, Nothing})

A vector space with optional dimension. Attribute id distinguishes between
different vector spaces of the same dimension. User should not need to set id.
"""
struct VectorSpace <: AbstractVectorSpace
    dim::Union{Int, Nothing}
    id::Int
    function VectorSpace(dim::Union{Int, Nothing})
        global counter
        counter += 1
        new(dim, counter)
    end
end

VectorSpace() = VectorSpace(nothing)

"""
    DualVectorSpace(vectorspace::AbstractVectorSpace)

Dual of a given vector space.
"""
struct DualVectorSpace <: AbstractVectorSpace
    vectorspace::AbstractVectorSpace
end

"""
    RnSpace(dim::Int)

Vector space corresponding to R^n for some n.
"""
struct RnSpace <: AbstractVectorSpace
    dim::Int
end

"""
    dual(V::AbstractVectorSpace)

Returns the dual of V.
"""
dual(V::VectorSpace) = DualVectorSpace(V)
dual(Vstar::DualVectorSpace) = Vstar.vectorspace
dual(R::RnSpace) = R
"""
    dim(V::AbstractVectorSpace)

Returns the dimension V.
"""
dim(V::VectorSpace) = V.dim
dim(Vstar::DualVectorSpace) = Vstar.vectorspace.dim
dim(R::RnSpace) = R.dim

"""
    Base.adjoint(V::AbstractVectorSpace)

Returns the [`dual(V)`](@ref).
"""
Base.adjoint(V::AbstractVectorSpace) = dual(V)

Base.show(io::IO, Vstar::DualVectorSpace) = print(io, Vstar', "*")

Base.show(io::IO, R::RnSpace) = print(io,  "R^",dim(R))

function Base.show(io::IO, V::VectorSpace)
    if dim(V) === nothing
        print(io, "V", "_", V.id)
    else
        print(io, "V", dim(V), "_", V.id)
    end
end


module LearnedJacobi

function step_iteration!(unext, uprev, rhs, stepsize, parameters)
    wL, wC, wR, wF = parameters
    numpts = length(unext)
    for i = 2:numpts-1
        unext[i] =
            wL * uprev[i-1] +
            wC * uprev[i] +
            wR * uprev[i+1] +
            wF * stepsize^2 * rhs[i]
    end
end

function run_parametric_iterations(
    bcleft,
    bcright,
    rhs,
    stepsize,
    numiter,
    parameters,
)
    T = eltype(parameters)
    numpts = length(rhs)

    u0 = zeros(T, numpts)
    u0[1] = bcleft
    u0[numpts] = bcright

    uprev = copy(u0)
    unext = copy(u0)

    for iter = 1:numiter
        step_iteration!(unext, uprev, rhs, stepsize, parameters)
        uprev .= unext
    end
    return unext
end

function error_for_parameters(
    bcleft,
    bcright,
    rhs,
    stepsize,
    numiter,
    uexact,
    parameters,
)

    itersol = run_parametric_iterations(
        bcleft,
        bcright,
        rhs,
        stepsize,
        numiter,
        parameters,
    )
    err = sum((uexact - itersol) .^ 2)
    return err
end

function run_parametric_iterations_all_rhs(
    bcleft,
    bcright,
    stepsize,
    numiter,
    parameters,
)
    T = eltype(parameters)

    solution = zeros(T, numpts, numpts - 2)
    solution[1, :] .= bcleft
    solution[end, :] .= bcright

    rhs = diagm(ones(numpts))[:, 2:numpts-1]
    rows, cols = size(rhs)

    for col = 1:cols
        u0 = solution[:, col]
        solution[:, col] .= run_parametric_iterations(
            u0,
            rhs[:, col],
            stepsize,
            numiter,
            parameters,
        )
    end
    return solution
end

function error_for_parameters_all_rhs(
    bcleft,
    bcright,
    stepsize,
    numiter,
    uexact,
    parameters,
)

    itersol = run_parametric_iterations_all_rhs(
        bcleft,
        bcright,
        stepsize,
        numiter,
        parameters,
    )
    err = sum((itersol - uexact) .^ 2)
    return err
end

# end module
end
# end module


type InterpolatedHazardRateCurve{DC <: DayCount, P <: Interpolation, B <: BusinessCalendar} <: InterpolatedDefaultProbabilityCurve{P}
  settlementDays::Int
  referenceDate::Date
  dc::DC
  interp::P
  cal::B
  dates::Vector{Date}
  times::Vector{Float64}
  data::Vector{Float64}
end

function InterpolatedHazardRateCurve(dates::Vector{Date},
                                    hazardRates::Vector{Float64},
                                    dc::DayCount,
                                    interpolator::Interpolation)
  ihc = InterpolatedHazardRateCurve(0, dates[1], dc, interpolator, NullCalendar(), dates, zeros(length(dates)), hazardRates)

  initialize!(ihc)

  return ihc
end

function initialize!(ihc::InterpolatedHazardRateCurve)
  length(ihc.dates) == length(ihc.data) || error("dates / data mismatch")

  ihc.times[1] = 0.0
  for i = 2:length(ihc.dates)
    ihc.times[i] = year_fraction(ihc.dc, ihc.dates[1], ihc.dates[i])
  end

  # initialize interpolator
  Math.initialize!(ihc.interp, ihc.times, ihc.data)

  Math.update!(ihc.interp)

  return ihc
end

function hazard_rate_impl(ihc::InterpolatedHazardRateCurve, t::Float64)
  if t <= ihc.times[end]
    return Math.value(ihc.interp, t)
  end

  # flat hazard rate extrap
  return ihc.data[end]
end

function survival_probability_impl(ihc::InterpolatedHazardRateCurve, t::Float64)
  if t == 0.0
    return 1.0
  end

  if t <= ihc.times[end]
    integral = get_primitive(ihc.interp, t, true)
  else
    # flat hazard rate extrapolation
    integral = get_primitive(ihc.interp, ihc.times[end], true) + ihc.data[end] * (t - ihc.times[end])
  end

  return exp(-integral)
end


using PkgSRA
using Plots, Random, OrdinaryDiffEq, Statistics
using Distributions, Turing
using StatsPlots, DSP
using Flux
using BSON: @save
pyplot()
Random.seed!(11)

#####
##### Define the controlled dynamical system
#####
# Load helper functions
# include("../../utils/sindy_statistics_utils.jl")
include("../../utils/sindy_turing_utils.jl")
include("../scripts/paper_settings.jl")
# Load example problem
include(EXAMPLE_FOLDERNAME*"example_lorenz.jl")

# Define the multivariate forcing function
num_ctr = 50;
    U_starts = rand(3, num_ctr) .* tspan[2]
    U_widths = [0.1, 0.1, 0.0];
    amplitudes = [150.0, 100.0, 0.0]
my_U_func_time2(t, u) = U_func_time_multivariate(t, u,
                        U_widths, U_starts,
                        F_dim=[1, 2, 3],
                        amplitudes=amplitudes)

#####
##### Produce data
#####
sol = solve_lorenz_system(my_U_func_time2)
dat = Array(sol)
# plot(dat[1,:],dat[2,:],dat[3,:],label="Data")
# Derivatives
numerical_grad = numerical_derivative(dat, ts)

## Also get baseline true/ideal cases
# Uncontrolled
dat_raw = Array(solve_lorenz_system())
numerical_grad_raw = numerical_derivative(dat_raw, ts)
true_grad_raw = zeros(size(dat))
for i in 1:size(dat,2)
    true_grad_raw[:,i] = lorenz_system(dat_raw[:,i], Float64.(p), [0])
end

# Controlled; true variables
true_grad = zeros(size(dat))
for i in 1:size(dat,2)
    true_grad[:,i] = lorenz_system(true_grad[:,i], dat[:,i], p, [0])
end
let d = dat
    plot(d[1,:],d[2,:],d[3,:],label="Data")
end

# True control signal
U_true = zeros(size(dat))
for (i, t) in enumerate(ts)
    U_true[:,i] = my_U_func_time2(t, dat[:,i])
end
plot(U_true')

#####
##### Calculate distribution of residuals
#####
# Turing.setadbackend(:forward_diff)
# CHEATING: known structure
#   Calculate posterior distribution of parameters
# chain, train_ind = generate_chain(dat,
#                         numerical_grad, lorenz_grad_residual)
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);
Random.seed!(13)

# Upgrade: Use AIC
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_unctr,best_criterion,all_criteria,all_models) =
    sindyc_ensemble(dat, numerical_grad, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,20),
                    use_clustering_minimization=true)
print_equations(sindy_unctr)
sindy_grad1 = sindy_unctr(dat, 0)
# scatter(sum.(val_list), all_criteria)

# Generate a full trajectory from SINDy model
condition(u,t,integrator) = any(abs.(u).>2e2)
cb = DiscreteCallback(condition, terminate!)
prob_unctr = ODEProblem(sindy_unctr, u0, tspan, [0], callback=cb)
sindy_dat_unctr = Array(solve(prob_unctr, Tsit5(), saveat=ts));

plot(dat[1,:],dat[2,:],dat[3,:],label="Data")
    plot!(sindy_dat_unctr[1,:],sindy_dat_unctr[2,:],sindy_dat_unctr[3,:],label="Turing model with enforced zeros")
    title!("Integrated SINDy model")

# grad_unctr = sindy_unctr(dat)
# let d = grad_unctr, d0 = numerical_grad
#     plot(d[1,:],d[2,:],d[3,:],label="Model")
#     plot!(d0[1,:],d0[2,:],d0[3,:],label="Data")
#     title!("Gradients (L2 SINDy model)")
# end
turing_unctr = convert_sindy_to_turing_enforce_zeros(sindy_unctr;
                                dat_noise_prior=Normal(0.0, 20.0),
                                coef_noise_std=2.0)
chain = generate_chain(dat, numerical_grad, turing_unctr,
                            iterations=200,
                            num_training_pts=100, start_ind=101)[1]
turing_unctr_sample = sindy_from_chain(sindy_unctr, chain,
                                        enforced_zeros=true)
# plot(chain["all_coef[3]"])
# Generate a single full sample trajectory (FROM TURING)
# prob_unctr = ODEProblem(sindy_unctr, u0, tspan, [0], callback=cb)
prob_unctr = ODEProblem(turing_unctr_sample, u0, tspan, [0], callback=cb)
sol_unctr = solve(prob_unctr, Tsit5(), saveat=ts);
dat_unctr = Array(sol_unctr)

# plot(dat[1,:],dat[2,:],dat[3,:],label="Data")
plot(dat_raw[1,:],dat_raw[2,:],dat_raw[3,:],label="Data")
    plot!(dat_unctr[1,:],dat_unctr[2,:],dat_unctr[3,:],label="Turing model with enforced zeros")
    title!("Integrated Turing model")

# Get the full posterior distribution
sample_ind = 1:length(ts)
sample_gradients3d, sample_noise = sample_sindy_posterior_grad(chain,
            dat, sample_ind, sindy_unctr)
# let y = sample_gradients3d[1,:,:], y2 = sample_gradients3d[2,:,:], n=numerical_grad
#     plot3d(y[:,1], y[:,2], y[:,3], label="Naive Sample 1")
#     plot3d!(y2[:,1], y2[:,2], y2[:,3], label="Naive Sample 2")
#     plot3d!(n[1,:], n[2,:], n[3,:], label="Data")
#     title!("Sampled Turing Gradients")
# end

# Calculate the residuals
sample_gradients = transpose(drop_all_1dims(mean(sample_gradients3d, dims=1)))
    sample_gradients_std = transpose(drop_all_1dims(std(sample_gradients3d, dims=1)))
    sample_trajectory_noise = mean(sample_noise)
dat_grad = numerical_grad[:, sample_ind]

residual = dat_grad .- sample_gradients
ctr_guess = process_residual(residual, sample_trajectory_noise)

# Visualization
ind = 101:1000
# ind = sample_ind
i = 2
    plot(dat_grad[i,ind], label="Data gradient")
    plot!(sample_gradients[i,ind], label="Sample gradient")
i = 2
    plot(residual[i,ind], ribbon=sample_trajectory_noise)
    plot!(U_true[i,ind], label="True")
    title!("Residual and true control")
i = 1
    plot(sample_gradients_std[i,ind])
    plot!(U_true[i,ind], label="True")
    title!("Residual and true control")
i = 2
    plot(residual[i,ind],ribbon=
            sample_trajectory_noise.+sample_gradients_std[i,ind])
    plot!(U_true[i,ind], label="True")
    title!("Residual and true control (with std)")
# i = 1
#     plot(ctr_guess[i,ind], label="Control guess")
#     plot!(U_true[i,ind], label="True", show=true)

#####
##### INSTEAD: subsample the data based on the above residual
#####

# Get "confident" indices
accepted_ind = subsample_using_residual(residual,
            sample_trajectory_noise, min_length=4)

num_pts = 400
start_ind = 101
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1]
dat_sub = dat[:,subsample_ind]
grad_sub = numerical_grad[:,subsample_ind]

# Any control signal in the subset?
U_sub = U_true[:,subsample_ind]
plot(U_sub')
    # plot!(residual[:,subsample_ind]', ribbon=sample_trajectory_noise)
    title!("Control signals in the subsampled dataset")

# SINDY SETUP
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(sindy_sub,best_criterion,all_criteria,all_models,best_index) =
    sindyc_ensemble(dat_sub, grad_sub, sindy_library, val_list,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms",
                    selection_dist=Normal(0.0,sample_trajectory_noise),
                    use_clustering_minimization=true)
println("Best index is $best_index")
print_equations(sindy_sub)
scatter(sum.(val_list), all_criteria)
    title!("AIC for various sparsities")
    xlabel!("Number of nonzero terms")

# Generate a single trajectory
prob_ctr = ODEProblem(sindy_sub, u0, tspan, [0], callback=cb)
sindy_dat_ctr = Array(solve(prob_ctr, Tsit5(), saveat=ts));
sindy_grad_ctr = sindy_sub(dat, 0)

plot(dat[1,:],dat[2,:],dat[3,:],label="Data")
    plot!(sindy_dat_ctr[1,:],sindy_dat_ctr[2,:],sindy_dat_ctr[3,:],label="Turing model with enforced zeros")
    title!("Integrated SINDy model (after control)")

i = 1
    plot(dat_grad[i,ind], label="Data Gradient")
    plot!(sindy_grad1[i,ind], label="Old SINDy gradient")
    plot!(sindy_grad_ctr[i,ind], label="New SINDy Gradient")

# Just get the residual from the SINDy model, not Turing
residual_ctr = dat_grad .- sindy_grad_ctr

i = 2
    plot(residual[i,ind], ribbon=sample_trajectory_noise, label="Original residual")
    plot!(residual_ctr[i,ind], label="Residual2")
    plot!(U_true[i,ind], label="True")
    title!("Iteration 2: Residual and true control for i=$i")

# TURING ANALYSIS
turing_sub = convert_sindy_to_turing_enforce_zeros(sindy_sub;
                                dat_noise_prior=Normal(0.0, 1.0),
                                coef_noise_std=0.1)
chain_sub = generate_chain(dat, numerical_grad, turing_sub,
                            train_ind=subsample_ind,
                            iterations=300)[1]
turing_sub_sample = sindy_from_chain(sindy_sub, chain_sub,
                                        enforced_zeros=true)
println("Turing sample equations")
print_equations(turing_sub_sample)
println("SINDy equations")
print_equations(sindy_sub)
# Generate test trajectories from the CORRECTED posterior
sample_gradients3d_ctr, sample_noise_ctr =
            sample_sindy_posterior_grad(chain_sub,
                    dat, sample_ind, sindy_sub)

density(chain_sub)
# Calculate the residuals
sample_gradients_ctr = transpose(drop_all_1dims(mean(sample_gradients3d_ctr, dims=1)))
    sample_trajectory_noise_ctr = mean(sample_noise_ctr)
# dat_grad2 = numerical_grad[:, sample_ind]

residual2 = dat_grad .- sample_gradients_ctr
ctr_guess2 = process_residual(residual2, sample_trajectory_noise_ctr)

# Generate a single full sample trajectory (FROM TURING)
prob_ctr = ODEProblem(turing_sub_sample, u0, tspan, [0], callback=cb)
turing_dat_ctr = Array(solve(prob_ctr, Tsit5(), saveat=ts));

# Plot!
plot(dat[1,:],dat[2,:],dat[3,:],label="Data")
    plot!(dat_ctr[1,:],dat_ctr[2,:],dat_ctr[3,:],label="Turing model with enforced zeros")
    title!("Integrated Turing model (after control)")
i = 2
    plot(residual[i,ind], ribbon=sample_trajectory_noise, label="Original residual")
    plot!(residual2[i,ind], ribbon=sample_trajectory_noise_ctr, label="Residual2")
    plot!(U_true[i,ind], label="True")
    title!("Iteration 2: Residual and true control for i=$i")
i = 2
    plot(dat_grad[i,ind], label="Data gradient")
    plot!(sample_gradients_ctr[i,ind], label="Sample gradient2")
    plot!(sample_gradients[i,ind], label="Sample gradient1")
    title!("Gradient predictions")

println("SINDy of subset")
print_equations(sindy_sub)
println("Turing-SINDy of subset")
print_equations(turing_sub_sample)
println("True")
print_equations(core_dyn_true)

#####
##### Save everything; plotted in a different script
#####
this_dat_name = DAT_FOLDERNAME*"TMP_dat_flowchart_inkscape_"

# Intermediate data
fname = this_dat_name*"intermediate.bson";
@save fname chain chain_sub sol sample_gradients sample_gradients_ctr sample_noise sample_noise_ctr sample_ind numerical_grad

# Uncontrolled data
fname = this_dat_name*"uncontrolled_lorenz.bson"
@save fname dat_raw numerical_grad_raw

# Original ODE and controller variables
fname = this_dat_name*"ode_vars.bson";
@save fname dat dat_grad true_grad U_true

# SINDy variables 1: before control
fname = this_dat_name*"naive_vars_sindy.bson";
unctr_nz_terms = get_nonzero_terms(sindy_unctr)
unctr_nz_names = get_nonzero_term_names(sindy_unctr)
# NOTE: saving the SINDy model directly currently breaks BSON
@save fname sindy_dat_unctr unctr_nz_terms unctr_nz_names #sindy_unctr

# Bayesian variables 1: before control
fname = this_dat_name*"naive_vars_bayes.bson";
@save fname dat_unctr sample_trajectory_noise residual accepted_ind ctr_guess

# SINDY variables 2: after control
fname = this_dat_name*"ctr_vars_sindy.bson";
sub_nz_terms = get_nonzero_terms(sindy_sub)
sub_nz_names = get_nonzero_term_names(sindy_sub)
@save fname sindy_dat_ctr sindy_grad_ctr sub_nz_terms sub_nz_names #sindy_sub

# Bayesian variables 2: after control
fname = this_dat_name*"ctr_vars_turing.bson";
@save fname dat_ctr sample_trajectory_noise_ctr sample_gradients_ctr turing_dat_ctr ctr_guess2

# TODO: NN variables


#! /usr/bin/julia

# Rosetta Code, Euler's sum of powers conjecture

const lim = 250
const pwr = 5
const p = [i^pwr for i in 1:lim]

x = zeros(Int, pwr-1)
y = 0

for a in combinations(1:lim, pwr-1)
    b = searchsorted(p, sum(p[a]))
    0 < length(b) || continue
    x = a
    y = b[1]
    break
end

if y == 0
    println("No solution found for power = ", pwr, " and limit = ", lim, ".")
else
    s = [@sprintf("%d^%d", i, pwr) for i in x]
    s = join(s, " + ")
    println("A solution is ", s, " = ", @sprintf("%d^%d", y, pwr), ".")
end


DATA_SIZE = 360
index_timetype = Date(2000, 1,1) + Day.(0:(DATA_SIZE - 1))
vec1 = randn(DATA_SIZE)
vec2 = randn(DATA_SIZE)
vec3 = randn(DATA_SIZE)
df = DataFrame(Index = index_timetype, vec1 = vec1,vec2 = vec2, vec3 = vec3)
ts = TSx.TS(df, 1)

functions = [mean, median, sum, minimum, maximum, std]
cols = [2, 3, :vec1, :vec3]          # col 1 is index
windowsize = [1, 5, 100, DATA_SIZE]

@test typeof(TSx.rollapply(mean, ts, 2, 5)) == TSx.TS

for fun in functions
    res = @test TSx.rollapply(fun, ts, 1, 5).coredata[!, 2] == RollingFunctions.rolling(fun, df[!, 2], 5)
end

for fun in functions
    colname = "vec1_rolling_$(fun)"
    res = @test TSx.rollapply(fun, ts, :vec1, 5).coredata[!, colname] == RollingFunctions.rolling(fun, df[!, :vec1], 5)

end

for fun in functions
    colname = "vec3_rolling_$(fun)"
    res = @test TSx.rollapply(fun, ts, :vec3, 100).coredata[!, colname] == RollingFunctions.rolling(fun, df[!, :vec3], 100)
end

for fun in functions
    colname = "vec3_rolling_$(fun)"
    res = @test TSx.rollapply(fun, ts, :vec3, DATA_SIZE).coredata[!, colname] == RollingFunctions.rolling(fun, df[!, :vec3], DATA_SIZE)
end

for fun in functions
    res = @test_throws ErrorException TSx.rollapply(fun, ts, 2, 10000)
end


"""
Overview: We here collect all functions and data structures required to manage
solid materials. The package has an internal materials database that is
updateed by the user to insert his own materials.
"""

"""
> struct material
- ρ : density of the material
- Dᵢⱼₖₗ : fourth order elasticity tensor in Voigt notation
"""
struct material
  #
  ρ::Float64
  Dᵢⱼₖₗ::Matrix{Float64}
  #
end

"""
> load_material( name::String )
It loads the data structure from the database
- name : name of the material to load
"""
function MORFE_load_material(name::String)
  pth = joinpath(get_materials_path(),name*".jld2")
  data = load(pth)
  return data["material"]
end

"""
> add_material( name::String, ρ::Float64, E::Float64, ν::Float64 )
It adds an isotropic material to the internal database of the package.
- name = name of the material
- ρ = density of the material
- E = Young's modulus
- ν = Poisson's ratio
**return**  *nothing*
"""
function MORFE_add_material(name::String, ρ::Float64, E::Float64, ν::Float64)
  #
  pth = get_materials_path()
  #
  Dᵢⱼₖₗ = zeros(Float64,(dim*(dim-1),dim*(dim-1)))
  #
  λ = E*ν/((1.0+ν)*(1.0-2.0*ν))
  μ = E/(1+ν)
  #
  Dᵢⱼₖₗ[1,1] = λ+μ
  Dᵢⱼₖₗ[1,2] = λ
  Dᵢⱼₖₗ[2,1] = λ
  Dᵢⱼₖₗ[2,2] = λ+μ
  Dᵢⱼₖₗ[1,3] = λ
  Dᵢⱼₖₗ[3,1] = λ
  Dᵢⱼₖₗ[2,3] = λ
  Dᵢⱼₖₗ[3,2] = λ
  Dᵢⱼₖₗ[3,3] = λ+μ
  Dᵢⱼₖₗ[4,4] = μ/2.0
  Dᵢⱼₖₗ[5,5] = μ/2.0
  Dᵢⱼₖₗ[6,6] = μ/2.0
  #
  mat = material(ρ, Dᵢⱼₖₗ)
  mname = joinpath(pth,name*".jld2")
  save(mname, Dict("material"=>mat))
  #
  return nothing
end

"""
> add_material( name::String, ρ::Float64, Dᵢⱼₖₗ::Matrix{Float64} )
It adds an anisotropic material to the internal database of the package.
- name = name of the material
- ρ = density of the material
- Dᵢⱼₖₗ = fourth order elasticity tensor
**return**  *nothing*
"""
function MORFE_add_material(name::String, ρ::Float64, Dᵢⱼₖₗ::Matrix{Float64})
  #
  pth = get_materials_path()
  #
  mat = material(ρ, Dᵢⱼₖₗ)
  mname = joinpath(pth,name*".jld2")
  save(mname, Dict("material"=>mat))
  #
  return nothing
end

"""
>  get_materials_path()
It returns the path associated to the material database
**return**  pth::String
"""
function get_materials_path()
  tail = length(joinpath("src","MORFEInvariantManifold.jl"))
  pth = chop(pathof(MORFEInvariantManifold), head = 0, tail = tail)*"materials"
  return pth
end

"""
> list_materials()
It lists all materials in the database.
**return**  *nothing*
"""
function MORFE_list_materials()
  pth = get_materials_path()
  tail = length(".jld2")
  files = readdir(pth)
  for f in files
    println(chop(f,head=0,tail=tail))
  end
  return nothing
end

"""
> delete_materials()
It delets a material in the database.
- name = name of the material to delete
**return**  *nothing*
"""
function MORFE_delete_material(name::String)
  pth = joinpath(get_materials_path(),name)*".jld2"
  try
    rm(pth)
  catch
    println("Material not present")
  end
  return nothing
end


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

using JuliaFEM, Test

# Example of 2d linear elasticity + volume load + surface load

# Dictionary containing node coordinates
X = Dict(1 => [0.0, 0.0], 2 => [1.0, 0.0], 3 => [1.0, 1.0], 4 => [0.0, 1.0])

# Create new problem of type `Elasticity`
block = Problem(Elasticity, "block", 2)
block.properties.formulation = :plane_stress
block.properties.finite_strain = false
block.properties.geometric_stiffness = false

# Add volume element
element = Element(Quad4, (1, 2, 3, 4))
update!(element, "geometry", X)
update!(element, "youngs modulus", 288.0)
update!(element, "poissons ratio", 1/3)
update!(element, "displacement load 2", 576.0)
add_element!(block, element)

# Add boundary element for tractoin force
traction_element = Element(Seg2, (3, 4))
update!(traction_element, "geometry", X)
update!(traction_element, "displacement traction force 2", 288.0)
add_element!(block, traction_element)

# Define boundary conditions
bc = Problem(Dirichlet, "symmetry boundary conditions", 2, "displacement")
bc_elements = [Element(Seg2, (1, 2)), Element(Seg2, (4, 1))]
update!(bc_elements, "geometry", X)
update!(bc_elements[1], "displacement 2", 0.0)
update!(bc_elements[2], "displacement 1", 0.0)
add_elements!(bc, bc_elements)

# Create analysis, add problems to it and run analysis
analysis = Analysis(Linear)
add_problems!(analysis, block, bc)
run!(analysis)

# Analytical solution is known:
f = 288.0
g = 576.0
E = 288.0
nu = 1/3
u3 = block("displacement", 0.0)[3]
u3_expected = f/E*[-nu, 1] + g/(2*E)*[-nu, 1]
@test isapprox(u3, u3_expected)


function conv_bn_block(input, output, kernel_size)
    return Chain(Conv((kernel_size...,), input => output), BatchNorm(output), x -> relu.(x))
end

function fc_bn_block(input, output)
    return Chain(Dense(input, output), BatchNorm(output), x -> relu.(x))
end


using StaticArrays
using OrdinaryDiffEq
using Test

@inbounds @inline function ż(z, p, t)
    A, B, D = p
    p₀, p₂ = z[1:2]
    q₀, q₂ = z[3:4]

    return SVector{4}(
        -A * q₀ - 3 * B / √2 * (q₂^2 - q₀^2) - D * q₀ * (q₀^2 + q₂^2),
        -q₂ * (A + 3 * √2 * B * q₀ + D * (q₀^2 + q₂^2)),
        A * p₀,
        A * p₂
  )
end

condition(u, t, integrator) = u
affect!(integrator) = nothing
cbf(idx) = ContinuousCallback(condition,
    affect!, nothing, save_positions=(false, true), idxs=idx)
z0 = SVector{4}(7.1989885061904335, -0.165912283356219, 0., -3.63534900748947)

tspan=(0.,300.)
prob=ODEProblem(ż, z0, tspan, (A=1, B=0.55, D=0.4), callback=cbf(3))
sol=solve(prob, Vern9(), abstol=1e-14, reltol=1e-14,
    save_everystep=false, save_start=false, save_end=false, maxiters=1e6)

@test length(sol) > 100

prob=ODEProblem(ż, z0, (0,400.), (A=1, B=0.55, D=0.4), callback=cbf(3))
sol=solve(prob, Vern9(), abstol=1e-14, reltol=1e-14, save_everystep=false, save_start=false, save_end=false, maxiters=2e4)

@test length(sol) > 100

prob=ODEProblem(ż, z0, (0,5000.), (A=1, B=0.55, D=0.4), callback=cbf(3))
sol=solve(prob, Vern9(), abstol=1e-14, reltol=1e-14, save_everystep=false, save_start=false, save_end=false, maxiters=1e6)

@test length(sol) > 1500

f = function (du,u,p,t)
  du[1] = u[2]
  du[2] = -p[1]
end
function condition(u,t,integrator) # Event when event_f(u,t) == 0
  u[1]
end
function affect!(integrator)
  @test integrator.u[1] >= 0
  integrator.u[2] = -integrator.u[2]
end
cb2 = ContinuousCallback(condition,affect!)
tspan = (0.0,10000.0)
u0 = [50.0,0.0]
p = 9.8
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,Tsit5(),callback=cb2)
@test minimum(sol') > -40

function vcondition!(out,u,t,integrator)
  out[1] = u[1]
  out[2] = u[2]
end

function vaffect!(integrator, event_idx)
  @test integrator.u[1] >= 0.0
  if event_idx == 1
    integrator.u[2] = -integrator.u[2]
  else
    integrator.p = 0.0
  end
end

u0 = [50.0,0.0]
tspan = (0.0,15.0)
p = 9.8
prob = ODEProblem(f,u0,tspan,p)
Vcb = VectorContinuousCallback(vcondition!,vaffect!, 2 , save_positions=(true,true))
sol = solve(prob,Tsit5(), callback=Vcb)

f = function (du,u,p,t)
  du[1] = u[2]
  du[2] = -9.81
end

function condition(u,t,integrator) # Event when event_f(u,t) == 0
  u[1]    # Event when height crosses from positive to negative
end

function affect_neg(integrator)
  @test integrator.u[1] >= 0.0
  integrator.u[2] = -0.8*integrator.u[2]
end

cb = ContinuousCallback(condition,nothing,affect_neg! = affect_neg)

u0 = [1.0,0.0]
tspan = (0.0, 3.0)
prob = ODEProblem(f,u0,tspan)
sol = solve(prob,Tsit5(),saveat=0.01,callback=cb)

f! = function (du,u,p,t)
  du[1] = u[2]
  du[2] = -p
end

function condition!(out,u,t,integrator)
  out[1] = u[1]
  out[2] = u[2]
end

function affect!(integrator, event_idx)
  if event_idx == 1
    integrator.u[2] = -integrator.u[2]
  else
    integrator.p = 0.0
  end
end

u0 = [50.0,0.0]
tspan = (0.0,15.0)

begin
    p = 9.8
    prob = ODEProblem(f!,u0,tspan,p)
    Vcb = VectorContinuousCallback(condition!,affect!, 2 , save_positions=(true,true))
    sol = solve(prob,Tsit5(), callback=Vcb)
end


import NMF.nndsvd

function init_nndsvd!(glrm::GLRM; scale::Bool=true, zeroh::Bool=false,
                      variant::Symbol=:std, max_iters::Int=0)
    # NNDSVD initialization:
    #    Boutsidis C, Gallopoulos E (2007). SVD based initialization: A head
    #    start for nonnegative matrix factorization. Pattern Recognition
    m,n = size(glrm.A)

    # only initialize based on observed entries
    A_init = zeros(m,n)
    for i = 1:n
        A_init[glrm.observed_examples[i],i] = glrm.A[glrm.observed_examples[i],i]
    end

    # scale all columns by the Loss.scale parameter
    if scale
        for i = 1:n
            A_init[:,i] .*= glrm.losses[i].scale
        end
    end

    # run the first nndsvd initialization
    W,H = nndsvd(A_init, glrm.k, zeroh=zeroh, variant=variant)
    glrm.X = W'
    glrm.Y = H

    # If max_iters>0 do a soft impute for the missing entries of A.
    #   Iterate: Estimate missing entries of A with W*H
    #            Update (W,H) nndsvd estimate based on new A
    for iter = 1:max_iters
        # Update missing entries of A_init
        for j = 1:n
            for i = setdiff(1:m,glrm.observed_examples[j])
                A_init[i,j] = dot(glrm.X[:,i],glrm.Y[:,j])
            end
        end
        # Re-estimate W and H
        W,H = nndsvd(A_init, glrm.k, zeroh=zeroh, variant=variant)
        glrm.X = W'
        glrm.Y = H
    end
end


using WignerSymbols

function compute_all(max_angular)
    for l1 = 0:max_angular
        for l2 = 0:max_angular
            for l3 = 0:max_angular
                for m1 = -l1:l1
                    for m2 = -l2:l2
                        for m3 = -l3:l3
                            wigner3j(Float64, l1, l2, l3, m1, m2, m3)
                        end
                    end
                end
            end
        end
    end
end

for max_angular in [4, 8, 12, 16, 20]
    println("max_angular = $max_angular")
    # warmup & compile
    compute_all(max_angular)
    GC.gc()

    empty!(WignerSymbols.Wigner3j)
    @time compute_all(max_angular)
end


using PtFEM

data = Dict(
  # Frame(nels, nn, ndim, nst, nip, finite_element(nod, nodof))
  :struc_el => Frame(4, 5, 3, 1, 1, Line(2, 2)),
  :properties => [5.0e5;],
  :x_coords => [0.0, 1.25, 3.5, 4.0, 2.0],
  :y_coords => [0.0, 3.0, 2.0, 1.0, 1.5],
  :z_coords => [0.0, 0.0, 0.0, 0.0, 3.0],
  :g_num => [1 2 3 4;5 5 5 5],
  :support => [
    (1, [0 0 0]),
    (2, [0 0 0]),
    (3, [0 0 0]),
    (4, [0 0 0])
    ],
  :loaded_nodes => [
    (5, [20.0 -20.0 30.0])],
  :fixed_freedoms => [
    (5, 2, -0.0005)],
  :penalty => 1e19
)

data |> display
println()

@time fem, dis_df, fm_df = p42(data)
println()

display(dis_df)
println()
display(fm_df)


import Images, AxisArrays
using AxisArrays: Axis
using Colors

# Create a cone in 3d that changes color over time
sz = [201, 301, 31]
center = [(s+1)>>1 for s in sz]  # ceil(Int, sz/2)
C3 = Bool[(i-center[1])^2+(j-center[2])^2 <= k^2 for i = 1:sz[1], j = 1:sz[2], k = sz[3]:-1:1]
cmap1 = round.(UInt32, range(0, stop=255, length=60))
cmap = Array{UInt32}(undef, length(cmap1))
for i = 1:length(cmap)
    cmap[i] = cmap1[i]<<16 + cmap1[end-i+1]<<8 + cmap1[i]
end
C4 = Array{UInt32}(undef, sz..., length(cmap))
for i = 1:length(cmap)
    C4[:,:,:,i] = C3*cmap[i]
end
img = AxisArray(reinterpret(RGB24, C4), Axis{:x}(1:size(C4,1)), Axis{:y}(1:size(C4,2)), Axis{:z}(range(1,step=3,length=size(C4,3))), Axis{:time}(1:size(C4,4)))

nothing


using Fast12thRoot
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

@test fast12throot(0.0) == 0.0
#@test fast12throot(-0.0) == 0.0
@test fast12throot(1.0) == 1.0
@test isnan(fast12throot(NaN)) == true
@test isinf(fast12throot(Inf)) == true

# test smallest and largest normal and subnormal doubles
const vsmallestsubnorm = 2.0^(-1074)
const vlargestsubnorm = 2.0^(-1022)*(1.0-2.0^-52)
const vsmallestnorm = 2.0^(-1023)*(2.0 - 2.0^-52)
const vlargestnorm = 2.0^(1022)*(2.0-2.0^-52)

@test abs(big(vsmallestsubnorm)^(one(BigFloat)/big(12))/fast12throot(vsmallestsubnorm)-1.0)<0.5eps(Float64)
@test abs(big(vlargestsubnorm)^(one(BigFloat)/big(12))/fast12throot(vlargestsubnorm)-1.0)<0.5eps(Float64)
@test abs(big(vsmallestnorm)^(one(BigFloat)/big(12))/fast12throot(vsmallestnorm)-1.0)<0.5eps(Float64)
@test abs(big(vlargestnorm)^(one(BigFloat)/big(12))/fast12throot(vlargestnorm)-1.0)<0.5eps(Float64)

for ii=1:100
    @test fast12throot(Float64(ii)^(12.0)) == Float64(ii)
end


using PtFEM

data = Dict(
  # Beam(ndim, nst, nxe, nip, direction, finite_element(nod, nodof), axisymmetric)
  :struc_el => Beam(2, 1, 4, 1, :x, Line(2, 1), false),
  :properties => [1.0;],
  :x_coords => [0.0, 0.25, 0.5, 0.75, 1.0],
  :support => [
    (1, [0 1]),
    (5, [0 0])
    ],
  :limit => 100,
  :tol => 0.00001
)

data |> display
println()

@time m = p46(data)
println()

println("\nThe buckling load = $(m[1])")
m[4][1] = 0.0
println("\nThe buckling mode (iterations=$(m[2])):\n")
for i in 1:m[3]
  println("$i    $(m[4][m[5][:, i]+1])")
end


#####################################################################
##                      Data Type Declarations                     ##
#####################################################################

type StateSpace <: LTISystem
    A::Matrix{Float64}
    B::Matrix{Float64}
    C::Matrix{Float64}
    D::Matrix{Float64}
    Ts::Float64
    nx::Int
    nu::Int
    ny::Int
    statenames::Vector{String}
    inputnames::Vector{String}
    outputnames::Vector{String}

    function StateSpace(A::Matrix{Float64}, B::Matrix{Float64},
            C::Matrix{Float64}, D::Matrix{Float64}, Ts::Float64,
            statenames::Vector{String}, inputnames::Vector{String},
            outputnames::Vector{String})
        nx = size(A, 1)
        nu = size(B, 2)
        ny = size(C, 1)

        if size(A, 2) != nx && nx != 0
            error("A must be square")
        elseif size(B, 1) != nx
            error("B must have the same row size as A")
        elseif size(C, 2) != nx
            error("C must have the same column size as A")
        elseif nu != size(D, 2)
            error("D must have the same column size as B")
        elseif ny != size(D, 1)
            error("D must have the same row size as C")
        end
        # Validate names of state, input, and output
        if size(statenames, 1) != nx
            error("Must have same number of statenames as states")
        elseif size(inputnames, 1) != nu
            error("Must have same number of inputnames as inputs")
        elseif size(outputnames, 1) != ny
            error("Must have same number of outputnames as outputs")
        end
        # Validate sampling time
        if Ts < 0 && Ts != -1
            error("Ts must be either a positive number, 0
                   (continuous system), or -1 (unspecified)")
        end
        new(A, B, C, D, Ts, nx, nu, ny, statenames, inputnames, outputnames)
    end
end

function StateSpace(A::Array, B::Array, C::Array, D::Array, Ts::Real,
        statenames::Vector{String}, inputnames::Vector{String},
        outputnames::Vector{String})
    return StateSpace(float64mat(A), float64mat(B), float64mat(C),
            float64mat(D), map(Float64, Ts), statenames, inputnames, outputnames)
end

#####################################################################
##                      Constructor Functions                      ##
#####################################################################

@doc """`ss(A,B,C,D[, Ts, statenames=..., inputnames=..., outputnames=...]) -> sys`

Create a state-space model.
This is a continuous-time model if Ts is omitted or set to 0.
Otherwise, this is a discrete-time model with sampling period Ts.
Set Ts=-1 for a discrete-time model with unspecified sampling period.

State, input and output names: each can be either a vector of strings (one string per dimension),
or a single string (e.g., "x"). In the latter case, an index is automatically appended to identify
the coordinates for each dimension (e.g. "x1", "x2", ...).

`sys = ss(D[, Ts, ...])` specifies a static gain matrix D.""" ->
function ss(A::Array, B::Array, C::Array, D::Array, Ts::Real=0; kwargs...)
    # Check the kwargs for metadata
    nu = size(B, 2)
    ny, nx = size(C, 1, 2)
    kvs = Dict(kwargs)
    statenames = validate_names(kvs, :statenames, nx)
    inputnames = validate_names(kvs, :inputnames, nu)
    outputnames = validate_names(kvs, :outputnames, ny)
    return StateSpace(A, B, C, D, Ts, statenames, inputnames, outputnames)
end

# Function for accepting scalars
function ss(A::Union{Real,Array}, B::Union{Real,Array}, C::Union{Real,Array}, D::Union{Real,Array}, args...; kwargs...)
    A = float64mat(A)
    B = float64mat(B)
    C = float64mat(C)
    if D == 0
        D = zeros(size(C,1),size(B,2))
    else
        D = float64mat(D)
    end
    ss(A, B, C, D, args..., kwargs...)
end

# Function for creation of static gain
function ss(D::Array, Ts::Real=0; kwargs...)
    ny, nu = size(D, 1, 2)
    A = zeros(0, 0)
    B = zeros(0, nu)
    C = zeros(ny, 0)
    return ss(A, B, C, D, Ts, kwargs...)
end
ss(d::Real, Ts::Real=0; kwargs...) = ss([d], Ts, kwargs...)

# ss(sys) converts to StateSpace
ss(sys::LTISystem) = convert(StateSpace, sys)

# Create a random statespace system
function rss(nx::Int, nu::Int=1, ny::Int=1, feedthrough::Bool=true)
    Q = randn(nx, nx)
    A = Q*diagm(-100*abs(randn(nx)))*Q'
    B = randn(nx, nu)
    C = randn(ny, nx)
    if feedthrough
        D = randn(ny, nu)
    else
        D = zeros(ny, nu)
    end
    return ss(A, B, C, D)
end

#####################################################################
##                         Math Operators                          ##
#####################################################################

## EQUALITY ##
function ==(s1::StateSpace, s2::StateSpace)
    fields = [:Ts, :nx, :ny, :nu, :A, :B, :C, :D, :inputnames, :outputnames,
            :statenames]
    for field in fields
        if getfield(s1, field) != getfield(s2, field)
            return false
        end
    end
    return true
end

## Approximate ##
function isapprox(s1::StateSpace, s2::StateSpace)
    fieldsApprox = [:Ts, :nx, :ny, :nu, :A, :B, :C, :D]
    fieldsEqual = [:inputnames, :outputnames, :statenames]
    for field in fieldsApprox
        if !(getfield(s1, field) ≈ getfield(s2, field))
            return false
        end
    end
    for field in fieldsEqual
        if getfield(s1, field) != getfield(s2, field)
            return false
        end
    end
    return true
end

## ADDITION ##
function +(s1::StateSpace, s2::StateSpace)
    #Ensure systems have same dimensions
    if size(s1) != size(s2)
        error("Systems have different shapes.")
    elseif s1.Ts != s2.Ts
        error("Sampling time mismatch")
    end

    A = [s1.A zeros(s1.nx, s2.nx);
         zeros(s2.nx, s1.nx) s2.A]
    B = [s1.B ; s2.B]
    C = [s1.C s2.C;]
    D = [s1.D + s2.D;]

    # Naming strategy: If only one sys is named, use that. If the names are the
    # same, use them. If the names conflict, then they are ignored, and the
    # default "" is used.
    statenames = [s1.statenames; s2.statenames]
    if all(s1.inputnames .== "")
        inputnames = s2.inputnames
    elseif all(s2.inputnames .== "") || (s1.inputnames == s2.inputnames)
        inputnames = s1.inputnames
    else
        inputnames = fill(String(""),s1.ny)
    end
    if all(s1.outputnames .== "")
        outputnames = s2.outputnames
    elseif all(s2.outputnames .== "") || (s1.outputnames == s2.outputnames)
        outputnames = s1.outputnames
    else
        outputnames = fill(String(""),s1.nu)
    end
    return StateSpace(A, B, C, D, s1.Ts, statenames, inputnames, outputnames)
end

+(s::StateSpace, n::Real) = StateSpace(s.A, s.B, s.C, s.D .+ n, s.Ts,
        s.statenames, s.inputnames, s.outputnames)
+(n::Real, s::StateSpace) = +(s, n)

## SUBTRACTION ##
-(s1::StateSpace, s2::StateSpace) = +(s1, -s2)
-(s::StateSpace, n::Real) = +(s, -n)
-(n::Real, s::StateSpace) = +(-s, n)

## NEGATION ##
-(s::StateSpace) = StateSpace(s.A, s.B, -s.C, -s.D, s.Ts, s.statenames,
        s.inputnames, s.outputnames)

## MULTIPLICATION ##
function *(s1::StateSpace, s2::StateSpace)
    #Check dimension alignment
    #Note: s1*s2 = y <- s1 <- s2 <- u
    if s1.nu != s2.ny
        error("s1*s2: s1 must have same number of inputs as s2 has outputs")
    elseif s1.Ts != s2.Ts
        error("Sampling time mismatch")
    end

    A = [s1.A    s1.B*s2.C;
         zeros(s2.nx, s1.nx)  s2.A]
    B = [s1.B*s2.D ; s2.B]
    C = [s1.C   s1.D*s2.C;]
    D = [s1.D*s2.D;]

    statenames = [s1.statenames; s2.statenames]
    return StateSpace(A, B, C, D, s1.Ts, statenames, s2.inputnames,
            s1.outputnames)
end

*(s::StateSpace, n::Real) = StateSpace(s.A, s.B, s.C*n, s.D*n, s.Ts,
        s.statenames, s.inputnames, s.outputnames)
*(n::Real, s::StateSpace) = *(s, n)

## DIVISION ##
/(s1::StateSpace, s2::StateSpace) = s1*(1/s2)

function /(n::Real, s::StateSpace)
    # Ensure s.D is invertible
    Dinv = try
        inv(s.D)
    catch
        error("D isn't invertible")
    end
    return StateSpace(s.A - s.B*Dinv*s.C, s.B*Dinv, -n*Dinv*s.C, n*Dinv, s.Ts,
            s.statenames, s.outputnames, s.inputnames)
end

/(s::StateSpace, n::Real) = StateSpace(s.A, s.B, s.C/n, s.D/n, s.Ts,
        s.statenames, s.inputnames, s.outputnames)

#####################################################################
##                       Indexing Functions                        ##
#####################################################################
Base.ndims(::StateSpace) = 2
Base.size(s::StateSpace) = (s.ny, s.nu)
Base.size(s::StateSpace, d) = d <= 2 ? size(s)[d] : 1

function Base.getindex(s::StateSpace, inds...)
    if size(inds, 1) != 2
        error("Must specify 2 indices to index statespace model")
    end
    rows, cols = ControlSystems.index2range(inds...)
    return StateSpace([s.A;], [s.B[:, cols];], [s.C[rows, :];], [s.D[rows, cols];],
            s.Ts, [s.statenames;], [s.inputnames[cols];], [s.outputnames[rows];])
end

#####################################################################
##                        Display Functions                        ##
#####################################################################

# TODO : this is a very hacky way of handling StateSpace printing.
function _string_mat_with_headers(X::Matrix, cols::Vector{String},
                                rows::Vector{String})
    mat = [[""] reshape(cols,1,length(cols));
           rows X]
    p = (io, m) -> Base.showarray(io, m, false, header=false)
    return replace(sprint(p, mat), "\"", " ")
end

Base.print(io::IO, s::StateSpace) = show(io, s)

function Base.show(io::IO, s::StateSpace)
    # Compose the name vectors
    inputs = format_names(s.inputnames, "u", "?")
    outputs = format_names(s.outputnames, "y", "?")
    println(io, "StateSpace:")
    if s.nx > 0
        states = format_names(s.statenames, "x", "?")
        println(io, "A = \n", _string_mat_with_headers(s.A, states, states))
        println(io, "B = \n", _string_mat_with_headers(s.B, inputs, states))
        println(io, "C = \n", _string_mat_with_headers(s.C, states, outputs))
    end
    println(io, "D = \n", _string_mat_with_headers(s.D, inputs, outputs), "\n")
    # Print sample time
    if s.Ts > 0
        println(io, "Sample Time: ", s.Ts, " (seconds)")
    elseif s.Ts == -1
        println(io, "Sample Time: unspecified")
    end
    # Print model type
    if s.nx == 0
        print(io, "Static gain")
    elseif iscontinuous(s)
        print(io, "Continuous-time state-space model")
    else
        print(io, "Discrete-time state-space model")
    end
end


using InterPol
using Base.Test

@test interpolateAt([0], [0], 1) == 0

@test interpolateAt([0], [1], 1) == 1

@test interpolateAt([0,1,2,3], [0,1,2,3], 0) == 0

@test interpolateAt([0,1,2,3], [0,1,2,3], 4) == 4

@test interpolateAt([0,1,2], [0,1,4], 0) == 0

@test interpolateAt([0,1,2], [0,1,4], 1) == 1

@test interpolateAt([0,1,2], [0,1,4], 3) == 9

@test interpolateDerivativeAt([0], [0], 1) == 0

@test interpolateDerivativeAt([0], [1], 1) == 0

@test interpolateDerivativeAt([0,1,2,3], [0,1,2,3], 0) == 1

@test interpolateDerivativeAt([0,1,2,3], [0,1,2,3], 4) == 1

@test interpolateDerivativeAt([0,1,2], [0,1,4], 0) == 0

@test interpolateDerivativeAt([0,1,2], [0,1,4], 1) == 2

@test interpolateDerivativeAt([0,1,2], [0,1,4], 3) == 6


module Subsequences

export longest_common_subsequence, longest_contiguous_subsequence

# -------

function longest_common_subsequence(a, b; result_base = "", join_fn = string)
    lengths = zeros(length(a) + 1, length(b) + 1)

    for (i, x) in enumerate(a)
        for (j, y) in enumerate(b)
            if x == y
                lengths[i+1, j+1] = lengths[i, j] + 1
            else
                lengths[i+1, j+1] = max(lengths[i+1, j], lengths[i, j+1])
            end
        end
    end

    x, y = length(a) + 1, length(b) + 1
    result = result_base
    a_start, a_end = 0, 0
    b_start, b_end = 0, 0

    while x != 1 && y != 1
        if lengths[x, y] == lengths[x-1, y]
            x -= 1
        elseif lengths[x, y] == lengths[x, y-1]
            y -= 1
        else
            result = join_fn(a[x-1], result)

            a_start = x - 1
            b_start = y - 1

            if (a_end == 0)
                a_end = x - 1
                b_end = y - 1
            end

            x -= 1
            y -= 1
        end
    end

    result, a_start:a_end, b_start:b_end
end

longest_common_subsequence(a::Array, b::Array) = longest_common_subsequence(a, b, result_base = [], join_fn = vcat)

# -------

function longest_contiguous_subsequence(a, b)
    m = zeros(Int, length(a) + 1, length(b) + 1)
    longest, x_longest, y_longest = 0, 0, 0

    for x in 2:(length(a) + 1)
        for y in 2:(length(b) + 1)

            if a[x - 1] == b[y - 1]
                m[x, y] = m[x - 1, y - 1] + 1

                if m[x, y] > longest
                    longest = m[x, y]
                    x_longest = x
                    y_longest = y
                end
            else
                m[x, y] = 0
            end

        end
    end

    a_range = (x_longest - longest):(x_longest - 1)
    b_range = (y_longest - longest):(y_longest - 1)
    a[a_range], a_range, b_range
end

end


using IterativeSolvers
using Test

@testset "Hessenberg" begin

    # Some well-conditioned Hessenberg matrix
    H1 = [
        1.19789 1.42354 -0.0371401  0.0118481 -0.0362113  0.00269463;
        1.46142 4.01953  0.890729  -0.0157701 -0.0300656 -0.0191307;
        0.0     1.08456  3.35179    0.941966   0.0439339 -0.072888;
        0.0     0.0      1.29071    3.1746     0.853378   0.0202058;
        0.0     0.0      0.0        1.32227    3.06086    1.18129;
        0.0     0.0      0.0        0.0        1.58682    2.99037;
        0.0     0.0      0.0        0.0        0.0        1.45345
    ]

    H2 = [
        1.14661+1.00222im  1.38611+0.0100307im  -0.0280623+0.00494128im    0.0205642-0.00867472im;
        1.43066+0.0im      4.11145+1.0144im       0.883563+0.00583816im  -0.00332255-0.0274278im;
            0.0+0.0im       1.0865+0.0im            3.2653+0.990342im       0.942434+0.00586481im;
            0.0+0.0im          0.0+0.0im           1.31146+0.0im             3.11491+1.01587im;
            0.0+0.0im          0.0+0.0im               0.0+0.0im             1.42175+0.0im
    ]

    for H = (H1, H2)
        T = eltype(H)

        # Fist standard basis vector as rhs
        rhs = [i == 1 ? one(T) : zero(T) for i = 1 : size(H, 1)]

        # Compare \ against the optimized version.
        solution_with_residual = copy(rhs)
        ldiv!(IterativeSolvers.FastHessenberg(copy(H)), solution_with_residual)
        solution = H \ rhs

        # First part is the solution
        @test solution_with_residual[1 : size(H, 2)] ≈ H \ rhs

        # Last element is the residual
        @test abs(last(solution_with_residual)) ≈ norm(H * solution - rhs)
    end
end


# using Pkg
# Pkg.add("DataFramesMeta")

using Optim
using HTTP
using GLM
using LinearAlgebra
using Random
using Statistics
using DataFrames
using CSV
using DataFramesMeta

function q3g()
# Question 1
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdataBeta0.csv"
df = CSV.read(HTTP.get(url).body)

df = @transform(df, bus_id = 1:size(df,1))

dfy = @select(df, :bus_id,:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20,:RouteUsage,:Branded)
dfy_long = DataFrames.stack(dfy, Not([:bus_id,:RouteUsage,:Branded]))
rename!(dfy_long, :value => :Y)
dfy_long = @transform(dfy_long, time = kron(collect([1:20]...),ones(size(df,1))))
select!(dfy_long, Not(:variable))

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

df_long = leftjoin(dfy_long, dfx_long, on = [:bus_id,:time])
sort!(df_long,[:bus_id,:time])

# Question 2
sigma_hat = glm(@formula(Y ~ Odometer + Branded ), df_long, Binomial(), LogitLink())
println(sigma_hat)

# Question 3a
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS5-ddc/busdata.csv"
df = CSV.read(HTTP.get(url).body)

Y = Matrix(df[:,[:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20]])
X1 = Matrix(df[:,[:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20]])
X2 = Matrix(df[:,[:Xst1,:Xst2,:Xst3,:Xst4,:Xst5,:Xst6,:Xst7,:Xst8,:Xst9,:Xst10,:Xst11,:Xst12,:Xst13,:Xst14,:Xst15,:Xst16,:Xst17,:Xst18,:Xst19,:Xst20]])

# Question 3b
function create_grids()

    function xgrid(theta,xval)
        N      = length(xval)
        xub    = vcat(xval[2:N],Inf)
        xtran1  = zeros(N,N)
        xtran1c = zeros(N,N)
        lcdf   = zeros(N)
        for i=1:length(xval)
            xtran1[:,i]   = (xub[i] .>= xval).*(1 .- exp.(-theta*(xub[i] .- xval)) .- lcdf)
            lcdf        .+= xtran1[:,i]
            xtran1c[:,i] .+= lcdf
        end
        return xtran1, xtran1c
    end

    zval   = collect([.25:.01:1.25]...)
    zbin   = length(zval)
    xval   = collect([0:.125:25]...)
    xbin   = length(xval)
    tbin   = xbin*zbin
    xtran  = zeros(tbin,xbin)
    xtranc = zeros(xbin,xbin,xbin)
    for z=1:zbin
        xtran[(z-1)*xbin+1:z*xbin,:],xtranc[:,:,z] = xgrid(zval[z],xval)
    end

    return zval,zbin,xval,xbin,xtran
end
zval,zbin,xval,xbin,xtran = create_grids()

@views @inbounds function q3e()
# Question 3c
T=20
future_value_array=zeros(size(xtran,1),2,T+1)
β=0.9

for t=T+1:1
  for b=1:size(df,1)
   for z=1:zbin
    for x=1:xbin
       row=xval[x] + (zval[z]-1)*xbin
       v1[t]= 1.92596-0.148154*xval[x]+1.05919*Branded[b]+xtran[row,:]'*FV[(zval[z]-1)*xbin+1:zval[z]*xbin,b+1,t+1]
       v0[t]= xtran[1+(zval[z]-1)*xbin,:]'*FV[(zval[z]-1)*xbin+1:zval[z]*xbin,b+1,t+1]
       futurevaluearray[:,:,t]=β*log(exp(v0[t])+exp(v1[t]))
end
end
end
end

# Q3d
log_likelihood_value=zeros(T+1)
for t=T+1:1
row0=1+(:Zst-1)*xbin
row1=:Xst+(:Zst-1)*xbin
v1t-v0t
(xtran[row1,:].-xtran[row0,:])'*FV[row0:row0+xbin-1,B[i]+1,t+1]

P_i1t =exp(v1t −v0t)/(1+exp(v1t −v0t))
P_i0t = 1−P_i1t

log_likelihood_value[t]=d_ijt*log(P_ijt)
loglikehood=sum(log_likelihood_value)
end

return -loglikehood
end

end

# Question 3e
println(q3e())

# Question 3g
println(q3g())


export ProjectionVI

function ProjectionVI(model::AbstractVIModel, x0::AbstractVector; rho0::Float64 = 0.5, kwargs...)
  stp = GenericStopping(model, x0; kwargs...)
  return ProjectionVI(stp, rho0 = rho0)
end

function abresidual!(model, xk, rho, Fx) # xk + rho * F(xk)
  residual!(model, xk, Fx)
  Fx .*= rho
  Fx .+= xk
  return Fx
end

function ProjectionVI(stp::AbstractStopping; rho0::Float64 = 0.5)
  xk = stp.current_state.x
  xkp = similar(xk)
  rho = rho0
  Fx = similar(xk)

  OK = update_and_start!(stp)
  while !OK
    abresidual!(stp.pb, xk, rho, Fx)
    project!(stp.pb, Fx, xk) # possible failure here

    if norm(xk - xkp, Inf) < stp.meta.atol * rho
      stp.meta.optimal = true
    end
    OK = update_and_stop!(stp, x = xk)
  end

  return xk
end


using Test
using DataStructures

include("../optimizers/cma_es.jl")
include("optimizers/cma_es_deap.jl")
include("optimizers/compare_optimizer_states.jl")

number_generations = 100
population_size = 200
sigma = 1.5
free_parameters = 1000

tolerance = 0.00001

@testset "Optimizers" begin

    optimizer_configuration = OrderedDict("sigma" => sigma, "population_size" => population_size)

    # Initialize Optimizers
    # OptimizerCmaEsDeap: Original Deap CMA-ES optimizer implemented in Python using PyCall
    # OptimizerCmaEs: CMA-ES optimizer implemented in Julia that we use in our framework
    optimizer1, eigenvectors1, indx1 = OptimizerCmaEsDeap(free_parameters, optimizer_configuration)
    optimizer2, eigenvectors2, indx2, diagD = OptimizerCmaEs(free_parameters, optimizer_configuration, test = true, eigenvectors1 = eigenvectors1, indx1 = indx1)

    # Compare return values
    @test size(indx1) == size(indx2)
    @test diagD[indx1] ≈ diagD[indx2] atol = 0.00001
    @test size(eigenvectors1) == size(eigenvectors2)

    # Compare internal states of both optimizers
    compare_optimizer_states(optimizer1, optimizer2, tolerance)

    for generation = 1:number_generations

        # Ask optimizers for new population
        genomes1, randoms1 = ask(optimizer1)
        genomes2, randoms2 = ask(optimizer2, test = true, randoms1 = randoms1)

        # Compare return values
        @test genomes1 ≈ genomes2 atol = tolerance
        @test size(randoms1) == size(randoms2)
        @test mean(randoms1) ≈ mean(randoms2) atol = 0.1
        @test std(randoms1) ≈ std(randoms2) atol = 0.01

        # Generate random rewards
        rewards_training = rand(population_size)

        # Tell optimizers new rewards
        eigenvectors1, indx1 = tell(optimizer1, rewards_training)
        eigenvectors2, indx2, diagD = tell(optimizer2, rewards_training, test = true, eigenvectors1 = eigenvectors1, indx1 = indx1)

        # Compare return values
        @test size(indx1) == size(indx2)
        @test diagD[indx1] ≈ diagD[indx2] atol = 0.00001
        @test size(eigenvectors1) == size(eigenvectors2)

        # Compare internal states of both optimizers
        compare_optimizer_states(optimizer1, optimizer2, tolerance)

        # Test if C is a Hermitian matrix
        @test optimizer2.C ≈ Hermitian(optimizer2.C) atol = tolerance

    end
end


using DirectConvolution
using Test
using LinearAlgebra

@testset "DirectConvolution" begin

    include("utils.jl")

    include("linearFilter.jl")

    include("directConvolution.jl")

    include("SG_Filter.jl")

    include("udwt.jl")

end;
nothing


module PoincareInvariant2ndTest

    using GeometricIntegrators
    using PoincareInvariants
    using SymPy

    const Δt = 10.
    const B₀ = 1.
    const r₀ = 0.5
    const z₀ = 0.0
    const z₁ = 0.1
    const u₀ = 5E-1
    const u₁ = 5E-2
    const kz = 2
    const ku = 4

    function B(t, q)
        B₀ * (1 + q[1]^2 + q[2]^2)
    end

    function gcϑ(t, q, ϑ)
        ϑ[1] =-q[2] * (1 + B(t,q)) / 4
        ϑ[2] =+q[1] * (1 + B(t,q)) / 4
        ϑ[3] = q[4]
        ϑ[4] = 0
        nothing
    end

    function gcϑ(t, q)
        ϑ = zero(q)
        gcϑ(t, q, ϑ)
        return ϑ
    end

    function gcω(t, q, Ω)
        Ω[1,1] = 0
        Ω[1,2] =-B(t,q)
        Ω[1,3] = 0
        Ω[1,4] = 0

        Ω[2,1] =+B(t,q)
        Ω[2,2] = 0
        Ω[2,3] = 0
        Ω[2,4] = 0

        Ω[3,1] = 0
        Ω[3,2] = 0
        Ω[3,3] = 0
        Ω[3,4] =+q[4]

        Ω[4,1] = 0
        Ω[4,2] = 0
        Ω[4,3] =-q[4]
        Ω[4,4] = 0

        nothing
    end

    function gc_surface_q(s,t)
        x  = r₀*(s-0.5)
        y  = r₀*(t-0.5)
        z  = z₀ + z₁ * cos(2π*kz*s) * cos(2π*kz*t)
        u  = u₀ + u₁ * sin(2π*ku*s) * sin(2π*ku*t)

        [x, y, z, u]
    end

    function gc_surface_p(s,t)
        gcϑ(0, gc_surface_q(s,t))
    end

    function gc_dummy(t, a, b, c)
        nothing
    end

    function gc_dummy_pode(q₀, p₀)
        PODE(gc_dummy, gc_dummy, q₀, p₀)
    end

    function gc_dummy_iode(q₀)
        p₀ = [gcϑ(0,q) for q in q₀]
        IODE(gc_dummy, gc_dummy, gc_dummy, q₀, p₀; v̄=gc_dummy)
    end

    function compute_canonical_invariant_approxfun(nx, ny, nt=0)
        pinv = PoincareInvariant2ndCanonical(gc_dummy_pode, gc_surface_q, gc_surface_p, Δt, 4, nx, ny, nt)
        sol  = Solution(pinv.equ, Δt, nt)

        I = evaluate_poincare_invariant(pinv, sol)

        return I[0]
    end

    function compute_invariant_approxfun(nx, ny, nt=0)
        pinv = PoincareInvariant2ndApproxFun(gc_dummy_iode, gc_surface_q, gcω, Δt, 4, nx, ny, nt)
        sol  = Solution(pinv.equ, Δt, nt)

        I, J, ΔI, ΔJ = evaluate_poincare_invariant(pinv, sol)

        return I[0], J[0]
    end

    function compute_invariant_cop(nx, ny, nt=0)
        pinv = PoincareInvariant2ndCOP(gc_dummy_iode, gc_surface_q, gcω, Δt, 4, nx, ny, nt)
        sol  = Solution(pinv.equ, Δt, nt)

        I, J, ΔI, ΔJ = evaluate_poincare_invariant(pinv, sol)

        return I[0], J[0]
    end

    function compute_invariant_trapezoidal(nx, ny, nt=0)
        pinv = PoincareInvariant2ndTrapezoidal(gc_dummy_iode, gc_surface_q, gcω, Δt, 4, nx, ny, nt)
        sol  = Solution(pinv.equ, Δt, nt)

        I, J, ΔI, ΔJ = evaluate_poincare_invariant(pinv, sol)

        return I[0], J[0]
    end

    function compute_invariant_analytical()
        local verbosity = get_config(:verbosity)

        if get_config(:verbosity) > 1
            println()
            println("Second Euler-Poincaré Integral Invariant (Analytical)")
            println("=====================================================")
            println()
        end

        s, t = Sym("s, t")
        x, y, z, u = gc_surface_q(s,t)

        S = simplify(B₀ * (1 + x^2 + y^2) * ( diff(x, s) * diff(y, t) + diff(x, t) * diff(y, s) )
                   + u * ( diff(z, s) * diff(u, t) + diff(z, t) * diff(u, s) ) )
        I = N(SymPy.integrate(S, (s, 0, 1), (t, 0, 1)), 16)

        verbosity > 1 ? println("           I = ", I) : nothing

        return I
    end

    export compute_canonical_invariant_approxfun,
           compute_invariant_approxfun,
           compute_invariant_cop,
           compute_invariant_trapezoidal,
           compute_invariant_analytical

end


module TestTree0

using Test
import Pkg: TOML
import TreeLas
import GraphIdx

include("data.jl")

function readtoml(fname::String; nan="-1.0")
    txt = read(open(fname), String)
    txt = replace(txt, "nan" => nan)
    TOML.parse(txt)
end

function instance(tree)
    y = Float64.(tree["y"])
    n = length(y)
    mu = GraphIdx.Weights(Float64.(get(tree, "mu", 1.0)))
    pi = Int.(tree["parent"]) .+ 1
    root::Int = get(tree, "root", -1) + 1
    if root <= 0
        root = GraphIdx.Tree.find_root(pi)
    end
    lam = Float64.(tree["lam"])
    if lam isa Array
        @assert findall(lam .< 0) == [root] "$(typeof(lam)): $lam, $root"
        lam[root] = NaN
    end
    lam = GraphIdx.Weights(lam)
    t = GraphIdx.Tree.RootedTree(root, pi)
    return y, t, lam, mu
end

@testset "tree0.toml                     " begin
    for (i, ti) in enumerate(readtoml(data_dir("test", "tree0.toml"))["tree"])
        @testset "tree0.toml[$i]" begin
            y, t, lam, mu = instance(ti)
            @test TreeLas.TreeDP.tree_dp(y, t, lam, mu) ≈ ti["x"]
        end
    end
end

@testset "tree0.1.toml                   " begin
    ti = readtoml(data_dir("test", "tree0.1.toml"))["tree"][1]
    y, t, lam, mu = instance(ti)
    x = TreeLas.TreeDP.tree_dp(y, t, lam, mu)
    @test all(isfinite.(x))
    @test minimum(y) <= minimum(x)
    @test maximum(x) <= maximum(y)
end

end


using BlockwiseInversion
using Test

@testset "BlockwiseInversion.jl" begin

  t = 250

  b = rand(t)

  test_matrix = nothing

  test_matrix = zeros(t,t)

  for i in 1:t
    j = i
    test_matrix[i,j] = j + (i-1) * (t)

    if i > 1
      j = i - 1
      test_matrix[i,j] = j + (i-1) * (t)
    end

    if i < t
      j = i + 1
      test_matrix[i,j] = j + (i-1) * (t)
    end
  end

  test_vector = test_matrix \ b

  for n in 0:10:t

    m = t - n
    d = 1

    cur_matrix = FourBlockMatrix(n,m,d,d)

    @test size(cur_matrix.A) == (n,n)
    @test size(cur_matrix.D) == (m,m)

    for i in 1:t
      j = i
      cur_matrix[i,j] = j + (i-1) * (t)

      if i > 1
        j = i - 1
        cur_matrix[i,j] = j + (i-1) * (t)
      end

      if i < t
        j = i + 1
        cur_matrix[i,j] = j + (i-1) * (t)
      end
    end

    cur_vector = cur_matrix \ b

    @test sum( abs.( test_matrix .- cur_matrix ) ) / t < 1e-14
    @test sum( abs.( test_vector .- cur_vector ) ) / t < 1e-14

  end

end


struct SFloat{T} <: AbstractFloat
    value :: T
end
value(x::SFloat) = x.value
SFloat{T}(x::SFloat) where T = SFloat{T}(T(x.value))

"""
Stochastic Float64

This type represents a double-precision (64-bit) floating-point number, on which
every operation gets randomly rounded upwards or downwards.
"""
const SFloat64 = SFloat{Float64}

"""
Stochastic Float32

This type represents a single-precision (32-bit) floating-point number, on which
every operation gets randomly rounded upwards or downwards.
"""
const SFloat32 = SFloat{Float32}

det_type(::Type{T}) where {T<:Real} = T
det_type(::Type{SFloat{T}}) where {T} = T

sto_type(::Type{T}) where {T} = SFloat{T}

Base.promote_rule(::Type{T1}, ::Type{T2}) where {T1<:SFloat, T2} = sto_type(promote_type(det_type(T1), det_type(T2)))
result_type(x::T1, y::T2) where {T1, T2} = promote_type(T1, T2)

#binary operators
for op = (:+, :-, :*, :/)
    eval(quote
        Base.$op(a::SFloat,b::SFloat) = result_type(a,b)($op(RND, a.value,  b.value))
    end)
end

for op = (:<, :<=)
    eval(quote
        Base.$op(a::SFloat,b::SFloat) = $op(a.value,  b.value)
    end)
end

#unary operators
for op = (:-, :zero, :one, :abs)
    eval(quote
        Base.$op(a::SFloat) = SFloat($op(a.value))
    end)
end

value(x::Number) = x
deref(x) = x
deref(x::Array{T,0}) where T = x[]
macro reliable_digits(expr)
    return quote
        x = [value.($(esc(expr))) for i = 1:10]
        mu = mean(x)
        sigma = std(x)
        s = -log10.((/).(sigma, abs.(mu)))
        zip(mu, s) |> collect |> deref
    end
end


@testset "$TEST $G" begin

r1 = G(10,20)
r2 = DG(5,10)
@test typeof(r1) == G
@test typeof(r2) == DG
@test nv(r1) == 10
@test ne(r1) == 20
@test nv(r2) == 5
@test ne(r2) == 10

@test G(10, 20, seed=3) == G(10, 20, seed=3)
@test DG(10, 20, seed=3) == DG(10, 20, seed=3)
@test G(10, 20, seed=3) == erdos_renyi(10, 20, G, seed=3)
@test ne(G(10, 40, seed=3)) == 40
@test ne(DG(10, 80, seed=3)) == 80

er = erdos_renyi(100, 0.5, G)
@test er isa G
@test nv(er) == 100
@test 2300 < ne(er) < 2650        # 50*99/2 == 2475.0
@test is_directed(er) == false

er = erdos_renyi(100, 0.5, DG)
@test er isa DG
@test nv(er) == 100
@test 4600 < ne(er) < 5300                 # 50*99/2 == 4950.0
@test is_directed(er) == true

ws = watts_strogatz(10, 4, 0.2,  G)
@test nv(ws) == 10
@test ne(ws) == 20
@test is_directed(ws) == false

ws = watts_strogatz(10, 4, 0.2, DG)
@test nv(ws) == 10
@test ne(ws) == 20
@test is_directed(ws) == true

ba = barabasi_albert(10, 2,  G)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == false

ba = barabasi_albert(10, 2, 2,  G)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == false

ba = barabasi_albert(10, 4, 2, G)
@test nv(ba) == 10
@test ne(ba) == 12
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(2, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 17
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(2, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 17
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(4, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == false

ba = barabasi_albert(10, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == true

ba = barabasi_albert(10, 2, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == true

ba = barabasi_albert(10, 4, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 12
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(2, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(2, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(4, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 24
@test is_directed(ba) == true

fm = static_fitness_model(20, rand(10), G)
@test nv(fm) == 10
@test ne(fm) == 20
@test is_directed(fm) == false

fm = static_fitness_model(20, rand(10), rand(10), DG)
@test nv(fm) == 10
@test ne(fm) == 20
@test is_directed(fm) == true

sf = static_scale_free(10, 20, 2.0, G)
@test nv(sf) == 10
@test ne(sf) == 20
@test is_directed(sf) == false

sf = static_scale_free(10, 20, 2.0, 2.0, DG)
@test nv(sf) == 10
@test ne(sf) == 20
@test is_directed(sf) == true
@test sf isa DG

rr = random_regular_graph(5, 0, G)
@test nv(rr) == 5
@test ne(rr) == 0
@test is_directed(rr) == false
@test rr isa G

rd = random_regular_digraph(10,0)
@test nv(rd) == 10
@test ne(rd) == 0
@test is_directed(rd)

rr = random_regular_graph(6, 3, G, seed=1)
@test all(degree(rr) .== 3)
@test nv(rr) == 6
@test ne(rr) == 9
@test is_directed(rr) == false

g = random_bipartite_regular_graph(6, 8, 4, 3, G, seed=1)
ks = degree(g)
@test all(ks[1:6] .== 4)
@test all(ks[7:14] .== 3)
@test nv(g) == 14
@test ne(g) == 24
@test is_directed(g) == false
@test g isa G

rr = random_regular_graph(100, 5, G)
@test nv(rr) == 100
@test ne(rr) == 250
@test is_directed(rr) == false
for v in vertices(rr)
    @test degree(rr, v) == 5
end
@test rr isa G

rr = random_configuration_model(10, repeat([2,4] ,5), G, seed=3)
@test nv(rr) == 10
@test ne(rr) == 15
@test is_directed(rr) == false
num2 = 0; num4 = 0
for v in vertices(rr)
    d = degree(rr, v)
    @test  d == 2 || d == 4
    d == 2 ? num2 += 1 : num4 += 1
end
@test num4 == 5
@test num2 == 5
@test rr isa G

ks1 = [1,3,2,1]
ks2 = [1,2,2,1,1]

g = random_bipartite_configuration_model(4, 5, ks1, ks2, G, seed=1)
ks = degree(g)
@test all(ks[1:4] .== ks1)
@test all(ks[5:9] .== ks2)
@test nv(g) == 9
@test ne(g) == 7
@test is_directed(g) == false
@test g isa G

rr = random_configuration_model(100, zeros(Int,100), G)
@test nv(rr) == 100
@test ne(rr) == 0
@test is_directed(rr) == false
@test rr isa G

rr = random_configuration_model(3, [2,2,2], G, check_graphical=true)
@test nv(rr) == 3
@test ne(rr) == 3
@test is_directed(rr) == false
@test rr isa G

rd = random_regular_digraph(100, 4, DG)
@test nv(rd) == 100
@test ne(rd) == 400
@test is_directed(rd)
@test std(out_degree(rd)) == 0

rd = random_regular_digraph(100, 4, DG, dir=:in)
@test nv(rd) == 100
@test ne(rd) == 400
@test is_directed(rd)
@test std(in_degree(rd)) == 0

rr = random_regular_graph(10, 8, G, seed=4)
@test nv(rr) == 10
@test ne(rr) == 40
@test is_directed(rr) == false
for v in vertices(rr)
    @test degree(rr, v) == 8
end

rd = random_regular_digraph(10, 8, DG, dir=:out, seed=4)
@test nv(rd) == 10
@test ne(rd) == 80
@test is_directed(rd)

g = stochastic_block_model(2., 3., [100,100], G)
@test  3.5 < mean(degree(g)) < 6.5
g = stochastic_block_model(3., 4., [100,100,100])
@test  9.5 < mean(degree(g)) < 12.5

end # testset


# Monte Carlo Simulation of Berry 1994

function Berry1994(R, repNum;
                     βo = 5.00,
                     βx = 2.00,
                     α  = 1.00,
                     γo = 1.00,
                     γx = 0.50,
                     γw = 0.25,
                     σc = 0.25,
                     σw = 0.25,
                     σd = 3.00)

#Let's set up the remaining parameters
J  = 2

# Let's randomly draw observable and unobservable characteristics that affect
# demand: xₘⱼ and ξₘⱼ
xₘⱼ =  randn(R, J, repNum);
ξₘⱼ =  randn(R, J, repNum);

# Let's randomly draw observable and unobservable characteristics that affect
# marginal cost: wₘⱼ and ωₘⱼ
wₘⱼ =  randn(R, J, repNum);
ωₘⱼ =  randn(R, J, repNum);

# Let's compute the marginal cost (cₘⱼ) for each market m and product/firm j
# as Berry does:
cₘⱼ = exp.( γo.*ones(R,J,repNum) .+ γx.*xₘⱼ + σc.*ξₘⱼ .+ γw.*wₘⱼ .+ σw.*ωₘⱼ );

# Let's compute the equilibrium prices (pₘⱼ) and market share (sₘⱼ)
function eq_p_sh( xₘ, ξₘ, cₘ, sₘ)
    p = cₘ .+ ( α.*(1.0.-sₘ)).^(-1);
    v = βo .+ βx.*xₘ .+ σd.*ξₘ .- α.*p;
    F = (exp.(v))./(1.0 .+ sum(exp.(v)));
    return F
end

s₀ = [0.5, 0.5];
sₘⱼ = similar(xₘⱼ);
pₘⱼ = similar(xₘⱼ);

for k = 1:repNum
    for r = 1:R
        sol = nlsolve(y -> eq_p_sh( xₘⱼ[r,:,k], ξₘⱼ[r,:,k], cₘⱼ[r,:,k], y)-y, s₀)

        s_eq = sol.zero

        if ( sol.f_converged == false || sum(s_eq .< 0)>0 || sum(s_eq)>1 )
            s_eq = [NaN , NaN];
        end

        p_eq = cₘⱼ[r,:,k] .+ (α.*(1.0 .- s_eq) ).^(-1);

        sₘⱼ[r,:,k] = s_eq;
        pₘⱼ[r,:,k] = p_eq;

    end
end

@assert sum(isnan.(sₘⱼ)) == 0

# Let's calculate the mean utility of product j: (δₘⱼ)
δₘⱼ = similar(xₘⱼ);
for k = 1:repNum
for r=1:R
    δₘⱼ[r,1,k] = log(sₘⱼ[r,1,k])-log(1-sₘⱼ[r,1,k]-sₘⱼ[r,2,k]);
    δₘⱼ[r,2,k] = log(sₘⱼ[r,2,k])-log(1-sₘⱼ[r,1,k]-sₘⱼ[r,2,k]);
end
end

# Let's compute the OLS estimates
θhatᵒˡˢ = zeros(3,repNum);

for k = 1:repNum
    y = reshape(δₘⱼ[:,:,k],R*J,1);
    X = [ones(R*J,1)  reshape(xₘⱼ[:,:,k],R*J,1) reshape(pₘⱼ[:,:,k],R*J,1) ];
    θhatᵒˡˢ[:,k] = inv(X'*X)*(X'*y);
end

# Taking the average and the standard deviation of the simulations
m_θhatᵒˡˢ  = mean(θhatᵒˡˢ,dims=2);
sd_θhatᵒˡˢ = std(θhatᵒˡˢ,dims=2);

# Let's compute the IV estimates
θhatⁱᵛ = zeros(3,repNum);

aux_xₘⱼ = similar(xₘⱼ);
aux_xₘⱼ[:,1,:] = xₘⱼ[:,2,:];
aux_xₘⱼ[:,2,:] = xₘⱼ[:,1,:];

for k = 1:repNum
    y = reshape(δₘⱼ[:,:,k],R*J,1);
    X = [ones(R*J,1)  reshape(xₘⱼ[:,:,k],R*J,1)  reshape(pₘⱼ[:,:,k],R*J,1) ];
    Z = [ones(R*J,1)  reshape(xₘⱼ[:,:,k],R*J,1)  reshape(wₘⱼ[:,:,k],R*J,1)  reshape(aux_xₘⱼ[:,:,k],R*J,1) ];
    Xhat = Z*inv(Z'*Z)*(Z'*X);
    θhatⁱᵛ[:,k] = inv(Xhat'*Xhat)*(Xhat'*y);
end

# Taking the average and the standard deviation of the simulations
m_θhatⁱᵛ  = mean(θhatⁱᵛ, dims=2);
sd_θhatⁱᵛ = std(θhatⁱᵛ, dims=2);

# Organizing the outputs
θ = [βo ; βx ; -α];
comparison = (θ = θ, m_θhatᵒˡˢ= m_θhatᵒˡˢ,  sd_θhatᵒˡˢ=sd_θhatᵒˡˢ,  m_θhatⁱᵛ= m_θhatⁱᵛ,  sd_θhatⁱᵛ= sd_θhatⁱᵛ);
data = ( xₘⱼ=xₘⱼ, ξₘⱼ=ξₘⱼ, wₘⱼ=wₘⱼ, ωₘⱼ=ωₘⱼ, cₘⱼ=cₘⱼ, sₘⱼ=sₘⱼ, pₘⱼ=pₘⱼ, δₘⱼ=δₘⱼ );
params = (βo = βo, βx = βx, α  = α, γo = γo, γx = γx, γw = γw, σc = σc, σw = σw, σd = σd);
output = (comparison = comparison, data = data, params=params);

return output

end


include("euler/euler.jl")
using .BigIntegers: BigInteger
using .Numbers: digits_sum
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int, p::Int)::Int = digits_sum(BigInteger(n) ^ p)

compute(2, 15)

compute(2, 1_000)

@benchmark compute(2, 1_000)


"""
The two stage method in DiffEqParamEstim is great, but it employs (intelligent) estimation of the derivative of the data over time. If we actually have a nominal solution, we know exactly what the derivative is over time. So this exploits it
"""

"""
For a prob::ODEProblem, with nominal parameters p0, creates a cost function C(p)
Let pprob = remake(prob, p=p).
C(p) is the collocation cost associated with  pprob. Calculated by integrating the following over the trajectory of solve(prob; saveat=tsteps):

 int_u  sum(pprob.f(u) - prob.f(u)).^2

 with output map g(x), this turns into

 int_u  sum(dgdx(u)*pprob.f(u) - dgdx(u)*prob.f(u)).^2

"""
function build_injection_loss(prob::ODEProblem, solmethod::T, tpoints, output_map= x -> x) where T <: DiffEqBase.AbstractODEAlgorithm
    pdim = length(prob.u0)
    nom_sol = Array(solve(prob, solmethod, saveat=tpoints))
    n = length(tpoints)
    dgdx = x -> ForwardDiff.jacobian(output_map, x)
    dgdx_template = dgdx(prob.u0)
    dgdx_template2 = deepcopy(dgdx_template)

    function cost(p)
        pprob = remake(prob, p=p)
        du_nom = similar(prob.u0, promote_type(eltype(prob.u0), eltype(p)))
        du_p = similar(pprob.u0, promote_type(eltype(pprob.u0), eltype(p)))
        c = 0.
        @inbounds for i = 1:n
            prob.f(du_nom,nom_sol[:,i], prob.p, tpoints[i])
            pprob.f(du_p, nom_sol[:,i], p, tpoints[i])
            dgdx_template = dgdx(nom_sol[:,i])
            c += sum(abs2, dgdx_template*du_nom .- dgdx_template*du_p)
        end
        return c
    end

    function cost2(p,g)
        g[:] = ForwardDiff.gradient(cost, p)
        return cost(p)
    end
    return DiffCost(cost, cost2)
end


# ====================================================
# Homogeneization of linear systems
# ====================================================

# no-op
homogenize(ivp::IVP{LCS{N, MT}, ST}) where {N, MT<:AbstractMatrix{N}, ST} = ivp
homogenize(sys::LCS{N, MT}) where {N, MT<:AbstractMatrix{N}} = sys

"""
    homogenize(ivp::IVP{CLCCS{N,MT,IdentityMultiple{N},XT,ConstantInput{SI}},ST}) where {N, MT<:AbstractMatrix{N}, XT<:LazySet{N}, SI<:Singleton{N}, ST<:LazySet{N}}

Transform an inhomogeneous linear initial-value problem into an homogeneous one
by introducing auxiliary state variables.

### Input

- `ivp` -- initial-value problem

### Output

Homogeneous initial-value problem.

### Notes

This function transforms the canonical initial-value problem ``x' = Ax + u``,
``x ∈ X`` with ``u(0) ∈ U = {u}`` (a singleton) into an homogeneous problem
without inputs ``y' = Â * y``, ``y ∈ Y``.
"""
function homogenize(ivp::IVP{CLCCS{N,MT,IdentityMultiple{N},XT,ConstantInput{SI}},ST}) where {N, MT<:AbstractMatrix{N}, XT<:LazySet{N}, SI<:Singleton{N}, ST<:LazySet{N}}
    # homogenized state matrix
    U = inputset(ivp) |> ReachabilityAnalysis.next_set
    A = state_matrix(ivp)
    Â = _homogenize_state_matrix(A, U)

    # homogenized input set
    X0 = initial_state(ivp)
    Y0 = _homogenize_initial_state(X0)

    X = stateset(ivp)
    Y = _homogenize_stateset(X)
    ivph = IVP(CLCS(Â, Y), Y0)
    return ivph
end

"""
    homogenize(sys::SOACS)

Transform an inhomogeneous second order system into an homogeneous one
by introducing auxiliary state variables.

### Input

- `sys` -- second order system

### Output

First-order homogeneous system.

### Notes

This function transforms the second-order system ``Mx'' + Cx' + Kx = b`` into a
first-order, homogeneous one, ``y' = Â * y``. It is assumed that the matrix ``M``
is invertible.
"""
function homogenize(sys::SOACS)
    @unpack M, C, K, b = sys

    n = size(M, 1)
    invM = inv(M)
    Zn = spzeros(n, n)
    In = Matrix(1.0I, n, n)

    A = [Zn           In     ;
        -invM*K       -invM*C]
    b0 = vcat(zeros(n), invM * b)

    m = 2n
    Aext = spzeros(m+1, m+1)
    Aext[1:m, 1:m] .= A
    Aext[1:m, m+1] .= b0

    return @system(x' = Aext*x)
end

function _homogenize_state_matrix(A::AbstractMatrix, U::Singleton)
    u = element(U)
    n = size(A, 1)
    Â = zeros(n+1, n+1)
    Â[1:n, 1:n] .= A
    Â[1:n, n+1] .= u
    return Â
end

function _homogenize_initial_state(X0::Singleton{N}) where {N}
    x0 = element(X0)
    y0 = vcat(x0, one(N))
    Y0 = Singleton(y0)
    return Y0
end

function _homogenize_initial_state(X0::LazySet{N}, dims=1) where {N}
    Y0 = X0 × Singleton(ones(N, dims))
    return Y0
end

function _homogenize_stateset(X::Universe, ndims=1)
    return Universe(dim(X) + ndims)
end

function _homogenize_stateset(X::LazySet, ndims=1)
    return X × Universe(ndims)
end

function _homogenize_stateset(X::HalfSpace{N}, ndims=1) where {N}
    a′ = vcat(X.a, fill(zero(N), ndims))
    return HalfSpace(a′, X.b)
end


struct Affine{M,V}
    A::M
    u::V
    """
        Affine(A, b)

    Represents an (inverse) affine transform of the form `x -> A * (x - u)`.
    """
    function Affine(A::AbstractMatrix, u::AbstractVector)
        @assert size(A, 2) == length(u)
        return new{typeof(A), typeof(u)}(A, u)
    end
end

function Affine(A::AbstractMatrix)
    u = zeros(eltype(A), size(A, 2))
    return Affine(A, u)
end

function Affine(u::AbstractVector)
    A = Diagonal(ones(eltype(u), length(u)))
    return Affine(A, u)
end

Affine(::Type{T}, n::Int) where {T} = Affine(zeros(T, n))
Affine(n::Int) = Affine(Float64, n)

"""
    one(t::Affine)

Returns the identity affine transformation, which maps all points to themselves.
"""
Base.one(t::Affine) = Affine(one(t.A), zero(t.u))

"""
    zero(t::Affine)

Returns the zero affine transformation, which maps all points to zero.
"""
Base.zero(t::Affine) = Affine(zero(t.A), zero(t.u))

function Base.copy!(dst::Affine, src::Affine)
    dst.A .= src.A
    dst.u .= src.u
    return dst
end

Base.:*(t::Affine, x::AbstractVecOrMat) = t.A * (x .- t.u)
Base.:\(t::Affine, y::AbstractVecOrMat) = t.A \ y .+ t.u
Base.:*(s::Affine, t::Affine) = Affine(s.A * t.A, t \ s.u)
Base.:\(s::Affine, t::Affine) = Affine(s.A \ t.A, t.u - t.A \ s.A * s.u)

Base.:*(t::Affine, c::Number) = Affine(t.A * c, t.u / c)
Base.:*(c::Number, t::Affine) = Affine(c * t.A, t.u)
Base.:-(t::Affine) = Affine(-t.A, t.u)

Base.:+(s::Affine, t::Affine) = Affine(s.A + t.A, (s.A + t.A) \ (s.A * s.u + t.A * t.u))
Base.:-(s::Affine, t::Affine) = Affine(s.A - t.A, (s.A - t.A) \ (s.A * s.u - t.A * t.u))

Base.inv(t::Affine) = Affine(inv(t.A), -t.A * t.u)

"""
    whitening_transform(μ, C)

Returns the `Affine` transform that whitens data with mean `μ` and covariance `C`.
"""
whitening_transform(u::AbstractVector, C::AbstractMatrix) = Affine(inv(cholesky(C).L), u)
whitening_transform(C::AbstractMatrix) = Affine(inv(cholesky(C).L))
whitening_transform(u::AbstractVector) = Affine(u)

function whitening_transform!(affine::Affine, μ::AbstractVector, C::AbstractMatrix)
    affine.A .= inv(cholesky(C).L)
    affine.u .= μ
    return affine
end

function whitening_transform!(affine::Affine, μ::AbstractVector)
    copyto!(affine.A, I)
    affine.u .= μ
    return affine
end

function whitening_transform!(affine::Affine, C::AbstractMatrix)
    affine.A .= inv(cholesky(C).L)
    affine.u .= 0
    return affine
end


# The now-defunct JuMPeR formulations of the problems are included here!

""" Robust facility location model, with JuMPeR. """
function robust_facility_model(c::Matrix, f::Vector, rho::Real, Gamm::Real)
    n, m = size(c)
    @assert length(f) == n
    model = RobustModel(solver = GurobiSolver())
    # model = Model(solver = GurobiSolver())

    # VARIABLES
    @variable(model, x[1:n], Bin)                 # Facility locations
    @variable(model, y[1:n, 1:m] >= 0)            # Flow variables
    @uncertain(model, -rho <= z[1:m] <= rho)      #  Demand uncertainties
    @constraint(model, norm(z, 1) <= Gamm)

    # CONSTRAINTS
    @constraint(model, [j=1:m], sum(y[i, j] for i = 1:n) >= d[j] + (P*z)[j]) # demand constraint
    @constraint(model, [i=1:n], sum(y[i, j] for j = 1:m) <= s[i] * x[i]) # capacity constraint
    # OBJECTIVE
    @objective(model, Min,
            sum(f[j] * x[j] for j = 1:n) + sum(c[i, j] * y[i, j] for i=1:n, j=1:m))
    return model, x, y
end

""" Adaptive facility location model, with JuMPeR and relaxation of binary variables. """
function adaptive_facility_model(c::Matrix, f::Vector, rho::Real, Gamm::Real)
    n, m = size(c)
    @assert length(f) == n
    model = RobustModel(solver = GurobiSolver())

    # VARIABLES
    @variable(model, 0 <= x[1:n] <= 1)                 # Facility locations
    @uncertain(model, -rho <= z[1:m] <= rho)      #  Demand uncertainties
    @constraint(model, norm(z, 1) <= Gamm)
    @adaptive(model, y[i=1:n, j=1:m] >= 0, policy = Affine, depends_on=z) # Flow variables

    # CONSTRAINTS
    @constraint(model, [j=1:m], sum(y[i, j] for i = 1:n) >= d[j] + (P*z)[j]) # demand constraint
    @constraint(model, [i=1:n], sum(y[i, j] for j = 1:m) <= s[i] * x[i]) # capacity constraint
    # OBJECTIVE
    @variable(model, F)
    @constraint(model, F >= sum(f[j] * x[j] for j = 1:n) + sum(c[i, j] * y[i, j] for i=1:n, j=1:m))
    @objective(model, Min, F)
    return model, x, y
end

# Adaptive solution with greedy binary heuristic (3.4 and 3.6)
function iterate_adaptive_model(c, f, rho, Gamm)
    ints = []
    xvals = zeros(size(c, 1))
    while !all([val in [0,1] for val in xvals]) || isempty(ints)
        model, x, y = adaptive_facility_model(c, f, rho, Gamm)
        for int in ints
            @constraint(model, x[int] == 1)
        end
        optimize!(model)
        xvals = value.(x)
        nonint = [!(var in [0,1]) for var in xvals]
        if sum(nonint) == 0
            return model, x, y, ints
        end
        max_val, max_idx = findmax(nonint.*xvals)
        push!(ints, max_idx)
    end
    return
end

# model, x, y, ints = iterate_adaptive_model(c, f, rho, Gamm)

# # Plot generation for adaptive solution (3.4 and 3.6)
# cost = objective_value(model)
# nonints = [val for val in 1:n if !(val in ints)]
# plot_model, x, y = facility_model(c, f)
# @constraint(plot_model, x[ints] .== 1)
# @constraint(plot_model, x[nonints] .== 0)
# optimize!(plot_model)
# plt = plot_solution(model, x, y, cost)
# @show plt


using Test

using LinearAlgebra
using SymbolicWedderburn
using SymbolicWedderburn.FiniteFields
using LinearAlgebra

using GroupsCore
using PermutationGroups
using Cyclotomics
using SparseArrays

include("smallgroups.jl")

@testset "Characters" begin
    import SymbolicWedderburn.Characters
    include("gf.jl")
    include("eigenspacedecomposition.jl")
    include("ccmatrix.jl")
    include("dixon.jl")
    include("characters.jl")
end
include("projections.jl")
include("sa_basis.jl")
include("action_permutation.jl")
include("action_linear.jl")
include("action_dihedral.jl")

if VERSION >= v"1.7.0" && !haskey(ENV, "CI")
    @testset "Examples" begin
        include("../examples/run_examples.jl")
    end
end


include("1.1.3.1 (a+b x^n)^p.jl")
include("1.1.3.2 (c x)^m (a+b x^n)^p.jl")
include("1.1.3.3 (a+b x^n)^p (c+d x^n)^q.jl")
include("1.1.3.4 (e x)^m (a+b x^n)^p (c+d x^n)^q.jl")
include("1.1.3.5 (a+b x^n)^p (c+d x^n)^q (e+f x^n)^r.jl")
include("1.1.3.6 (g x)^m (a+b x^n)^p (c+d x^n)^q (e+f x^n)^r.jl")
include("1.1.3.7 P(x) (a+b x^n)^p.jl")
include("1.1.3.8 P(x) (c x)^m (a+b x^n)^p.jl")

src/rules/1 Algebraic functions/1.1 Binomial products/1.1.3 General/.jl