# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module OrthoPolynomials
using Nemo, Triangles

export ModuleOrthoPolynomials
export OrthoPoly, InvOrthoPoly
export T053121, T216916, T217537, T064189, T202327, T111062, T099174
export T066325, T049310, T137338, T104562, T037027, T049218, T159834, T137286
export T053120, T053117, T111593, T059419
export L217924, L005773, L108624, L005425, L000085, L001464, L003723, L006229

"""

* OrthoPoly, InvOrthoPoly, T053121, T216916, T217537, T064189, T202327, T111062, T099174, T066325, T049310, T137338, T104562, T037027, T049218, T159834, T137286, T053120, T053117, T111593, T059419, L217924, L005773, L108624, L005425, L000085, L001464, L003723, L006229
"""
const ModuleOrthoPolynomials = ""

# Cf. http://oeis.org/wiki/User:Peter_Luschny/AignerTriangles

"""

By the theorem of Favard an orthogonal polynomial systems  ``p_{n}(x)`` is a sequence of real polynomials with deg``(p_{n}(x)) = n`` for all ``n`` if and only if

`` p_{n+1}(x) = (x - s_n)p_n(x) - t_n p_{n-1}(x) ``

with ``p_{0}(x)=1`` for some pair of seq's ``s_k`` and ``t_k``. Return the coefficients of the polynomials as a triangular array with `dim` rows.
"""
function OrthoPoly(dim::Int, s::Function, t::Function)
    dim ≤ 0 && return ZZ[]
    T = fill(ZZ(0), dim, dim)
    for n ∈ 1:dim T[n, n] = 1 end
    for n ∈ 2:dim, k ∈ 1:n-1
        T[n, k] = ((k > 1 ? T[n-1, k-1] : 0)
                + s(k - 1) * T[n-1, k] + t(k) * T[n-1, k+1])
    end
    [T[n, k] for n ∈ 1:dim for k ∈ 1:n] # flatt format
    # [[T[n, k] for k ∈ 1:n] for n ∈ 1:dim] # triangle format
end

"""

Return the inverse of the coefficients of the orthogonal polynomials generated by ``s`` and ``t`` as a triangular array with `dim` rows.
"""
function InvOrthoPoly(dim::Int, s::Function, t::Function)
    dim ≤ 0 && return ZZ[]
    T = fill(ZZ(0), dim, dim)
    for n ∈ 1:dim T[n, n] = 1 end
    for n ∈ 1:dim-1, k ∈ 1:n+1
        T[n+1, k] = ((k > 1 ? T[n, k-1] : 0)
        - s(n - 1) * T[n, k] - (n > 1 ? t(n - 1) * T[n-1, k] : 0))
    end
    [T[n, k] for n ∈ 1:dim for k ∈ 1:n]
end

"""

Return the Catalan triangle (with 0's) read by rows.
"""
T053121(dim::Int) = OrthoPoly(dim, n -> 0, n -> 1)

# """
# binomial(n, floor(n/2)).
# """
# L001405(len::Int) = RowSums(T053121(len))

"""

Return the coefficients of some orthogonal polynomials related to set partitions without singletons (cf. A000296).
"""
T216916(dim::Int) = OrthoPoly(dim, n -> n + 1, n -> n + 1)

"""

Return the triangle ``T(n,k)`` of tangent numbers, coefficient of ``x^n/n!`` in the expansion of ``(tan x)^k/k!``.
"""
T059419(dim::Int) = OrthoPoly(dim, n -> 0, n -> n * (n - 1))

"""

Return the expansion of exp(tan(x)).
"""
L006229(len::Int) = RowSums(T059419(len))

"""

Return the first len integers defined as ``a(n) = n! [x^n] \\exp(2 \\exp (x) - x - 2)``.
"""
L217924(len::Int) = RowSums(T217537(len))

"""

Return the coefficients of some orthogonal polynomials related to indecomposable set partitions without singletons (cf. A098742).
"""
T217537(dim::Int) = OrthoPoly(dim, n -> n, n -> n)

"""

Return the (reflected) Motzkin triangle.
"""
T064189(dim::Int) = OrthoPoly(dim, n -> 1, n -> 1)

"""

Return the number of directed animals of size n as an array of length len.
"""
L005773(len::Int) = RowSums(T064189(len))

"""

Return the coefficients of ``x^n`` in the expansion of ``((-1-x+√(1+2x+5x^2))/2)^k`` as a triangle with dim rows.
"""
T202327(dim::Int) = OrthoPoly(dim, n -> -1, n -> -1)

"""

Return the sequence with generating function satisfying ``x = (A(x)+(A(x))^2)/(1-A(x)-(A(x))^2)``.
"""
L108624(len::Int) = RowSums(T202327(len))

"""

Return the triangle ``T(n, k) = \\binom{n}{k} \\times`` involutions``(n - k)``.
"""
T111062(dim::Int) = OrthoPoly(dim, n -> 1, n -> n)

"""

Return the number of self-inverse partial permutations.
"""
L005425(len::Int) = RowSums(T111062(len))

"""

Return the coefficients of the modified Hermite polynomials.
"""
T099174(dim::Int) = OrthoPoly(dim, n -> 0, n -> n)

# Also
# T099174(dim::Int) = InvOrthoPoly(dim, n -> 0, n -> -n)

"""

Return the number of involutions.
"""
L000085(len::Int) = RowSums(T099174(len))

"""

Return the coefficients of unitary Hermite polynomials He``_n(x)``.
"""
T066325(dim::Int) = InvOrthoPoly(dim, n -> 0, n -> n)

"""

Return the sequence defined by ``a(n) = n! [x^n] \\exp(-x-(x^2)/2)``.
"""
L001464(len::Int) = RowSums(T066325(len), true)

"""

Return the triangle of tanh numbers.
"""
T111593(dim::Int) = OrthoPoly(dim, n -> 0, n -> -n * (n - 1))

"""

Return the sequence defined by ``A(n) = n! [x^n] \\exp \\tan(x)`` as an array of length `len`.
"""
L003723(len::Int) = RowSums(T111593(len))

"""

Return the coefficients of Chebyshev's U``(n, x/2)`` polynomials.
"""
T049310(dim::Int) = InvOrthoPoly(dim, n -> 0, n -> 1)

"""

Return the coefficients of the Charlier polynomials with parameter ``a = 1``.
"""
T137338(dim::Int) = InvOrthoPoly(dim, n -> n + 1, n -> n + 1)

"""

Return the inverse of the Motzkin triangle (cf. A064189).
"""
T104562(dim::Int) = InvOrthoPoly(dim, n -> 1, n -> 1)

"""

Return the skew Fibonacci-Pascal triangle with `dim` rows.
"""
T037027(dim::Int) = InvOrthoPoly(dim, n -> -1, n -> -1)

"""

Return the arctangent numbers (expansion of arctan``(x)^n/n!``).
"""
T049218(dim::Int) = InvOrthoPoly(dim, n -> 0, n -> n * (n + 1))

"""

Return the coefficients of Hermite polynomials ``H(n, (x-1)/√(2))/(√(2))^n``.
"""
T159834(dim::Int) = InvOrthoPoly(dim, n -> 1, n -> n)

"""

Return the coefficients of a variant of the Hermite polynomials.
"""
T137286(dim::Int) = InvOrthoPoly(dim, n -> 0, n -> n + 1)

"""

Return the coefficients of the Chebyshev-T polynomials.
"""
function T053120(len)
    T = ZTriangle(len)
    R, x = PolynomialRing(ZZ, "x")
    m = 1
    for n ∈ 0:len-1
        f = chebyshev_t(n, x)
        for k ∈ 0:n
            T[m] = coeff(f, k)
            m += 1
        end
    end
    T
end

"""

Return the coefficients of the Chebyshev-U polynomials.
"""
function T053117(len)
    T = ZTriangle(len)
    R, x = PolynomialRing(ZZ, "x")
    m = 1
    for n ∈ 0:len-1
        f = chebyshev_u(n, x)
        for k ∈ 0:n
            T[m] = coeff(f, k)
            m += 1
        end
    end
    T
end

#START-TEST-########################################################

using Test, SeqTests

function test()

    @testset "OrthoPoly" begin

        @test isa(OrthoPoly(10, n -> 1, n -> n + 1)[end], fmpz)
        @test isa(InvOrthoPoly(10, n -> 1, n -> n + 1)[end], fmpz)
        @test RowSums(T217537(8)) == L217924(8)

        if data_installed()

            T = [
                T066325,
                T049310,
                T137338,
                T104562,
                T037027,
                T049218,
                T159834,
                T137286,
                T053120,
                T053117,
                T053121,
                T216916,
                T217537,
                T064189,
                T202327,
                T111062,
                T099174,
                T111593,
                T064189
            ]
            SeqTest(T, 'T')

            L = [L217924, L005425, L000085, L001464, L003723, L108624, L006229]
            SeqTest(L, 'L')
        end
    end
end

function demo()
    T = T111593(8)
    ShowAsΔ(T)
    println(RowSums(T))

    T = T217537(8)
    ShowAsΔ(T)
    println(RowSums(T))

    T = T053117(8)
    ShowAsΔ(T)
    println(RowSums(T))
end

"""

T111062(500) :: 0.339080 seconds (750.52 k allocations: 15.375 MiB)
T066325(500) :: 0.157202 seconds (751.50 k allocations: 13.374 MiB)
T053120(500) :: 0.061058 seconds (375.75 k allocations:  6.705 MiB)

"""
function perf()
    GC.gc()
    @time T111062(500)
    GC.gc()
    @time T066325(500)
    GC.gc()
    @time T053120(500)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


using DecisionMakingProblems
using Test
using Random
using LinearAlgebra
using GridInterpolations

const p = DecisionMakingProblems

# MDP

@testset "2048.jl" begin
    m = TwentyFortyEight()
    mdp = MDP(m)
    @test length(mdp.𝒜) == 4
    @test mdp.γ == 1.0
    init_state = p.initial_board()
    s′, r = mdp.TR(init_state, p.DOWN)
    @test s′ != init_state || isdone(s′)
    @test r >= -1.0
end

@testset "cart_pole.jl" begin
    # m = CartPole(1.0, 10.0, 1.0, 1.0, 0.1, 9.8, 0.02, 4.8, deg2rad(24))
    m = CartPole()
    @test p.n_actions(m) == 2
    @test p.discount(m) == 1.0
    @test p.ordered_actions(m) == 1:2
    min_state = [-4.8, -Inf, -24 * pi/180, -Inf]
    max_state = [4.8, Inf, 24 * pi/180, Inf]
    state = p.generate_start_state(m)
    @test min_state <= p.vec(state) <= max_state
    @test !p.is_terminal(m, state)
    @test min_state <= p.vec(p.cart_pole_transition(m, state, rand(1:2))) <= max_state
    @test p.reward(m, state, rand(1:2)) in [0.0, 1.0]
    mdp = MDP(m)
end

@testset "collision_avoidance.jl" begin
    m = CollisionAvoidance()
    distrib = p.CollisionAvoidanceStateDistribution()
    s = p.rand(distrib)
    simple_pol = p.SimpleCollisionAvoidancePolicy()
    optimal_pol = p.OptimalCollisionAvoidancePolicy()
    @test length(p.vec(rand(p.transition(m, s, optimal_pol(s))))) == 4
    @test p.is_terminal(m, s) == (p.vec(s)[4] < 0.0)
    @test p.reward(m, rand(p.transition(m, s, optimal_pol(s))), rand(m.𝒜)) <= 0
    policy = p.CollisionAvoidanceValueFunction(m, simple_pol)
    mdp = MDP(m)
end

@testset "hexworld.jl" begin
    m = HexWorld()
    hexes = m.hexes
    @test p.n_states(m) == length(hexes) + 1 && p.ordered_states(m) == 1:length(hexes) + 1
    @test p.n_actions(m) == 6 && p.ordered_actions(m) == 1:6
    @test p.discount(m) == 0.9
    @test p.state_index(m, 1) == 1
    state = rand(1:p.n_states(m))
    action = rand(1:p.n_actions(m))
    p.transition(m, state, action)
    state_ = p.generate_s(m, state, action)
    reward = p.reward(m, state, action)
    @test state_ in p.ordered_states(m)
    @test reward <= 10
    @test p.generate_sr(m, state, action)[1] in p.ordered_states(m) && p.generate_sr(m, state, action)[2] <= 10
    @test p.generate_start_state(m) in p.ordered_states(m)
    @test p.hex_distance(rand(hexes), rand(hexes)) >= 0
    mdp = MDP(m)
end
@testset "simple_lqr.jl" begin
    m = LQR()
    @test p.discount(m) == 1.0
    state = p.generate_start_state(m)
    @test -10 <= rand(p.transition(m, state, rand())) <= 10
    @test p.reward(m, state, rand()) <= 0
    mdp = MDP(m)
end

@testset "mountain_car.jl" begin
    m = MountainCar()
    @test p.n_actions(m) == 3 && p.ordered_actions(m) == [1, 2, 3]
    @test p.discount(m) == 1.0
    state_min = [-1.2, -0.07]
    state_max = [0.6, 0.07]
    start_state = p.generate_start_state(m)
    @test all(state_min <= start_state <= state_max)
    @test all(state_min <= p.mountain_car_transition(start_state, 1) <= state_max)
    @test p.reward(m, start_state, 1) <= 0
    mdp = MDP(m)
end

# POMDP

@testset "crying_baby.jl" begin
    m = CryingBaby()
    @test p.n_states(m) == 2 && p.ordered_states(m) == [1, 2]
    @test p.n_actions(m) == 3 && p.ordered_actions(m) == [1, 2, 3]
    @test p.n_observations(m) == 2 && p.ordered_observations(m) == [true, false]
    @test p.discount(m) == 0.9
    p.transition(m, rand(1:2), rand(1:3))
    @test 0 <= p.observation(m, rand(1:3), rand(1:2)).p <= 1
    @test p.reward(m, rand(1:2), rand(1:3)) <= 0
    @test p.reward(m, [0.1, 0.9], rand(1:3)) <= 0
    pomdp = p.POMDP(m)
end

@testset "machine_replacement.jl" begin
    mdp = MachineReplacement()
    m = DiscretePOMDP(mdp)
    @test p.n_states(m) == 3 && p.ordered_states(m) == 1:3
    @test p.n_actions(m) == 4 && p.ordered_actions(m) == 1:4
    @test p.n_observations(m) == 2 && p.ordered_observations(m) == 1:2
    @test p.discount(m) == 1.0
    @test rand(p.transition(m, rand(1:3), rand(1:4))) in 1:3
    @test rand(p.observation(m, rand(1:4), rand(1:3))) in 1:2
    @test p.reward(m, rand(1:3), rand(1:4)) <= 1.0
    pomdp = p.POMDP(m)
end

@testset "catch.jl" begin
    mdp = Catch()
    m = DiscretePOMDP(mdp)
    @test p.n_states(m) == 4 && p.ordered_states(m) == 1:4
    @test p.n_actions(m) == 10 && p.ordered_actions(m) == 1:10
    @test p.n_observations(m) == 2 && p.ordered_observations(m) == 1:2
    @test p.discount(m) == 0.9
    @test rand(p.transition(m, rand(1:4), rand(1:10))) in 1:4
    @test rand(p.observation(m, rand(1:10), rand(1:4))) in 1:2
    @test p.reward(m, rand(1:4), rand(1:10)) >= 0
    pomdp = p.POMDP(m)
end

# Simple Game

@testset "prisoners_dilemma.jl" begin
    m = PrisonersDilemma()
    @test p.n_agents(m) == 2
    @test length(p.ordered_actions(m, rand(1:2))) == 2 && length(p.ordered_joint_actions(m)) == 4
    @test p.n_actions(m, rand(1:2)) == 2 && p.n_joint_actions(m) == 4
    @test p.reward(m, rand(1:2), [rand(p.ordered_actions(m, 0)), rand(p.ordered_actions(m, 0))]) <= 0.0
    @test p.joint_reward(m, [rand(p.ordered_actions(m, 0)), rand(p.ordered_actions(m, 0))]) <= [0.0, 0.0]
    simplegame = p.SimpleGame(m)
end

@testset "rock_paper_scissors.jl" begin
    m = RockPaperScissors()
    @test p.n_agents(m) == 2
    @test length(p.ordered_actions(m, rand(1:2))) == 3 && length(p.ordered_joint_actions(m)) == 9
    @test p.n_actions(m, rand(1:2)) == 3 && p.n_joint_actions(m) == 9
    @test -1.0 <= p.reward(m, rand(1:2), [rand(p.ordered_actions(m, 0)), rand(p.ordered_actions(m, 0))]) <= 1.0
    @test [-1.0, -1.0] <= p.joint_reward(m, [rand(p.ordered_actions(m, 0)), rand(p.ordered_actions(m, 0))]) <= [1.0, 1.0]
    simplegame = p.SimpleGame(m)
end

@testset "travelers.jl" begin
    m = Travelers()
    @test p.n_agents(m) == 2
    @test length(p.ordered_actions(m, rand(1:2))) == 99 && length(p.ordered_joint_actions(m)) == 99^2
    @test p.n_actions(m, rand(1:2)) == 99 && p.n_joint_actions(m) == 99^2
    @test 0.0 <= p.reward(m, rand(1:2), [rand(p.ordered_actions(m, 0)), rand(p.ordered_actions(m, 0))]) <= 102
    @test [0.0, 0.0] <= p.joint_reward(m, [rand(p.ordered_actions(m, 0)), rand(p.ordered_actions(m, 0))]) <= [102, 102]
    simplegame = p.SimpleGame(m)
end

# Markov Game

@testset "predator_prey.jl" begin
    m = PredatorPreyHexWorld()
    hexes = m.hexes
    @test p.n_agents(m) == 2
    @test length(p.ordered_states(m, rand(1:2))) == length(hexes) && length(p.ordered_states(m)) == length(hexes)^2
    @test length(p.ordered_actions(m, rand(1:2))) == 6 && length(p.ordered_joint_actions(m)) == 36
    @test p.n_actions(m, rand(1:2)) == 6 && p.n_joint_actions(m) == 36

    @test 0.0 <= p.transition(m, rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m)), rand(p.ordered_states(m))) <= 1.0
    @test -1.0 <= p.reward(m, rand(1:2), rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m))) <= 10.0
    @test [-1.0, -1.0] <= p.joint_reward(m, rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m))) <= [10.0, 10.0]
    mg = p.MG(m)
end

# POMG

@testset "multicaregiver.jl" begin
    m = MultiCaregiverCryingBaby()
    @test p.n_agents(m) == 2
    @test length(p.ordered_states(m)) == 2
    @test length(p.ordered_actions(m, rand(1:2))) == 3 && length(p.ordered_joint_actions(m)) == 9
    @test p.n_actions(m, rand(1:2)) == 3 && p.n_joint_actions(m) == 9
    @test length(p.ordered_observations(m, rand(1:2))) == 2 && length(p.ordered_joint_observations(m)) == 4
    @test p.n_observations(m, rand(1:2)) == 2 && p.n_joint_observations(m) == 4

    @test 0.0 <= p.transition(m, rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m)), rand(p.ordered_states(m))) <= 1.0
    @test 0.0 <= p.joint_observation(m, rand(p.ordered_joint_actions(m)), rand(p.ordered_states(m)), rand(p.ordered_joint_observations(m))) <= 1.0
    @test p.joint_reward(m, rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m))) <= [0.0, 0.0]
    @test p.joint_reward(m, rand(Float64, 2), rand(p.ordered_joint_actions(m))) <= [0.0, 0.0]
    pomg = p.POMG(m)
end

# DecPOMDP

@testset "collab_predator_prey.jl" begin
    m = CollaborativePredatorPreyHexWorld()
    hexes = m.hexes
    @test p.n_agents(m) == 2
    @test length(p.ordered_states(m, rand(1:2))) == length(hexes) && length(p.ordered_states(m)) == length(hexes)^3
    @test length(p.ordered_actions(m, rand(1:2))) == 6 && length(p.ordered_joint_actions(m)) == 36
    @test p.n_actions(m, rand(1:2)) == 6 && p.n_joint_actions(m) == 36
    @test length(p.ordered_observations(m, rand(1:2))) == 2 && length(p.ordered_joint_observations(m)) == 4
    @test p.n_observations(m, rand(1:2)) == 2 && p.n_joint_observations(m) == 4

    @test 0.0 <= p.transition(m, rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m)), rand(p.ordered_states(m))) <= 1.0
    @test 0.0 <= p.joint_observation(m, rand(p.ordered_joint_actions(m)), rand(p.ordered_states(m)), rand(p.ordered_joint_observations(m))) <= 1.0
    @test -1.0 <= p.reward(m, rand(p.ordered_states(m)), rand(p.ordered_joint_actions(m))) <= 10.0
    decpomdp = p.DecPOMDP(m)
end


export joMultigrid, multigrid_multiply, multigrid_vcycle

function joMultigrid(Hs,S,R,P,C,coarse_solver;
                     recursive_vcycle::Bool=false)
    n = size(Hs[1],1)
    nlevels = length(Hs)
    if recursive_vcycle
        coarse_solves = Array{Function,1}()
        prepend!(coarse_solves,[solvesystem(Hs[nlevels],coarse_solver)])
        for i=nlevels-1:-1:1
            old_cs = coarse_solves[1]
            Hc = Hs[i+1]
            Hf = Hs[i]
            Rf = R[i]
            Pf = P[i]
            Sf = S[i]
            coarse_solver = deepcopy(coarse_solver)
            coarse_solver.precond = old_cs
            CS = solvesystem(Hc,coarse_solver)
            prepend!(coarse_solves,[(b,x,mode::Bool)->multigrid_vcycle(Hf,Sf,Rf,Pf,CS,b,x,forw_mode=mode)])
        end
        solver = coarse_solves[1]
        return joLinearFunctionFwdCT(n,n,
                                     x->solver(x,zeros(x),true),
                                     x->solver(x,zeros(x),false),
                                     eltype(Hs[1]),name="Recursive Multigrid Preconditioner")
    else
        return joLinearFunctionFwdCT(n,n,
                             x->multigrid_multiply(Hs,S,R,P,C,x,forw_mode=true),
                             x->multigrid_multiply(Hs,S,R,P,C,x,forw_mode=false),
                             eltype(Hs[1]),name="V cycle Multigrid preconditioner")
    end

end

function multigrid_vcycle(H,S,R,P,C,b,x;forw_mode::Bool=true)
    xf = S(b,x,forw_mode)
    if forw_mode
        r = b-H*xf
    else
        r = b-H'*xf
    end
    xc = R*r
    xc = C(xc,zeros(xc),forw_mode)
    xf .+= P*xc
    xf = S(b,xf,forw_mode)
end

function multigrid_multiply(Hs,S,R,P,C,b;forw_mode::Bool=true)
    x = zeros(b)
    x_lvl = Array{typeof(b),1}()
    b_lvl = Array{typeof(b),1}()
    push!(x_lvl,x)
    push!(b_lvl,b)
    nlevels = length(Hs)
    for i=1:nlevels-1
        xf = x_lvl[i]
        bf = b_lvl[i]
        xf = S[i](bf,xf,forw_mode)
        if forw_mode
            rf = bf-Hs[i]*xf
        else
            rf = bf-Hs[i]'*xf
        end
        push!(x_lvl,R[i]*rf)
        push!(b_lvl,R[i]*bf)
    end
    x = x_lvl[end]
    x = C(x,zeros(x),forw_mode)
    for i=nlevels-1:-1:1
        xf = x_lvl[i] + P[i]*x
        x = S[i](b_lvl[i],xf,forw_mode)
    end
    return x

end


Base.@kwdef mutable struct NewtonSolver{T, LS<:AbstractLinearSolver{T}} <: AbstractNonlinearSolver{T}
    prob::NonlinearProblem{T}
    linalg::LS
    maxiters::Int
    P_tol::T
    r_tol::T
    iter::Int
    norm_r::T
    dP_rel::T
    P::Vector{T}
    P_old::Vector{T}
    r::Vector{T}
    success::Bool
end

function NewtonSolver{T}(prob::NonlinearProblem{T}, linalg::AbstractLinearSolver{T}, opts) where {T}
    Nt = size(prob.C, 2)
    kwargs = (
        maxiters = opts["maxiters"],
        P_tol = opts["P_tol"],
        r_tol = opts["r_tol"],
        iter = 0,
        norm_r = zero(T),
        dP_rel = zero(T),
        P = Array{T}(undef, Nt),
        P_old = Array{T}(undef, Nt),
        r = Array{T}(undef, Nt),
        success = false,
    )
    NewtonSolver{T, typeof(linalg)}(; prob, linalg, kwargs...)
end

function perform_step!(solver::NewtonSolver, Δt, n)

    @unpack P, P_old, r, linalg, prob, P_tol, r_tol = solver

    iter, success = 0, false
    for outer iter = 1:solver.maxiters
        J = getjac(linalg, n)
        val_and_jac!(r, J, P, Δt, prob, n)

        copyto!(P_old, P)
        solve!(P, r, n, linalg)
        @turbo for i = 1:length(P)
            P[i] = P_old[i] - P[i]
        end

        x_converged, f_converged = assess_convergence(P, P_old, r, P_tol, r_tol)
        (success = x_converged | f_converged) && break
    end
    setjac!(linalg, n)
    norm_r = maximum(abs, r)
    dP_rel = maxdiff(P, P_old) / maximum(abs, P)

    @pack! solver = iter, norm_r, dP_rel, success

    return solver
end


module Day02

using AdventOfCode2019

function day02(input::String = readInput(joinpath(@__DIR__, "..", "data", "day02.txt")))
    data = parse.(Int, split(input, ","))
    dataCopy = copy(data)

    data[1+1] = 12
    data[2+1] = 2
    part1 = _run_program(data)

    copy!(data, dataCopy)
    part2 = 0
    for noun = 0:99
        for verb = 0:99
            copy!(data, dataCopy)
            data[2] = noun
            data[3] = verb
            result = _run_program(data)
            if result == 19690720
                part2 = 100 * noun + verb
            end
        end
    end
    return [part1, part2]
end

function _run_program(data::Array{Int, 1})
    for i = 1:4:length(data)
        optcode = data[i]
        if optcode == 99
            break
        elseif optcode == 1 || optcode == 2
            pos1 = data[i+1] + 1
            pos2 = data[i+2] + 1
            pos3 = data[i+3] + 1
            if optcode == 1
                data[pos3] = data[pos1] + data[pos2]
            elseif optcode == 2
                data[pos3] = data[pos1] * data[pos2]
            end
        else
            println("Invalid optcode")
            break
        end
    end
    return data[1]
end

end # module


#
# Implementation of PRS similar to lassosum
#
# Mak et al.(2018), Polygenic scores via penalized regression on
# summary statistics
#

export z2cor,
    elnetg!,
    elnetg_path

"""
    z2cor(z, n)

Convert GWAS z-value to phenotype-genotype correlation coefficient.
"""
z2cor(z, n) = z ./ sqrt.(n .- 1 .+ z.^2)

"""
    elnetg!(X, r, λ, s, β = zeros(Float64, length(r));
            maxiter = 10_000, thresh = 1e-4)

Solve elastic net via coordinate descent.

# Arguments

- `X::Matrix` nsubj x nmarkers normalized genotype matrix
              need to pass in ac*sqrt((1-s)/nsubj),
              where ac is column normalized (μ=0, σ=1) allele count matrix
- `r::Vector` nmarkers x 1 vector of correlation between phenotype and genotypes
- `λ::Real`   shrinkage parameter for 1-norm
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnetg!(X, r, λ, s, β;
                 maxiter = 10_000, thresh = 1e-4,
                 verbose = false)

    p = length(r) # number of markers
    @assert p == size(X, 2) "Size mismatch: length(r) = $p, size(X, 2) = $(size(X, 2))"

    sd = std(X; dims = 1)[:]
    diag = ones(Float64, p) * (1.0 - s)
    diag[sd .== 0.0] .= 0.0

    yhat = X * β

    converged = false

    for k = 1:maxiter
        Δ = 0.0
        for j = 1:p
            β_j = β[j]
            β[j] = 0.0
            u_j = diag[j]*β_j + r[j]  - yhat'*X[:,j]

            abs_u_j = abs(u_j)
            #println("j=$j, abs_u_j = $abs_u_j")
            if abs_u_j > λ
                β[j] = sign(u_j) * (abs_u_j - λ) / (diag[j] + s)
            end
            Δ_j = β[j] - β_j
            Δ = max(Δ, abs(Δ_j))
            yhat += Δ_j * X[:,j]
        end

        if Δ < thresh
            converged = true
            break
        end

        if verbose
            print("\rλ = $λ, s = $s: Δ = $Δ  ($k/$maxiter)")
        end

    end
    obj = yhat'*yhat - 2*β'*r + 2*λ*sum(β) + s*β'*β

    if verbose
        println("\rλ = $λ, s = $s: obj = $obj, converged = $converged")
    end

    (converged = converged, obj = obj)
end

"""
    elnetg_path(X, r, λs, s;
                maxiter = 10_000, thresh = 1e-4)

Solve elastic net for path along `λs` using warm starts.

# Arguments

- `X::Matrix` nsubj x nmarkers normalized genotype matrix
              column normalized (μ = 0, σ = 1)
- `r::Vector` nmarkers x 1 vector of correlation coefficients
              between phenotype and genotypes
- `λs::Vector` shrinkage parameter path for 1-norm in decreasing order
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnetg_path(X, r, λs, s;
                     maxiter = 1000, thresh = 1e-4, verbose = false)

    nl = length(λs)
    p = length(r)

    β = zeros(Float64, p, nl)
    converged = zeros(Bool, nl)

    nsubj = size(X, 1)
    Xs = X .* sqrt((1.0 - s)/nsubj)
    ret = elnetg!(Xs, r, λs[1], s, view(β, :, 1);
                  maxiter = maxiter, thresh = thresh, verbose = verbose)
    converged[1] = ret.converged

    for l = 2:nl
        @assert λs[l] < λs[l-1] "λs not in decreasing order"
        β[:, l] .= β[:, l-1] # use warm start
        ret = elnetg!(Xs, r, λs[l], s, view(β, :, l);
                     maxiter = maxiter, thresh = thresh, verbose = verbose)
        converged[l] = ret.converged

    end

    (β = β,
     λs = λs,
     converged = converged,
     s = s)

end

#
# Implementations using sparse LD matrix instead of genotype matrix.
# These have convergence problems, though.
#

"""
    elnet(R, r, λ, s, β = zeros(Float64, length(r));
          maxiter = 10_000, thresh = 1e-4)

Solve elastic net via coordinate descent.

# Arguments

- `R::Matrix` p x p LD matrix
- `r::Vector` p x 1 vector of correlation between phenotype and genotypes
- `λ::Real`   shrinkage parameter for 1-norm
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnet!(R, r, λ, s, β;
                maxiter = 10_000, thresh = 1e-4,
                verbose = false)

    p = length(r)
    @assert p == size(β, 1) "Size mismatch: length(r) = $p, size(β, 1) = $(size(β, 1))"

    converged = false

    for k = 1:maxiter
        Δ = 0.0   # maximum change in β
        for j = 1:p
            β_j = β[j]
            β[j] = 0.0
            u_j = s ≈ 1 ? r[j] : (r[j] - (1.0-s) * β' * R[:,j])
            abs_u_j = abs(u_j)
            if abs_u_j > λ
                β[j] = sign(u_j) * (abs_u_j - λ) / (s + (1-s)*R[j,j])
            end
            Δ_j = β[j] - β_j
            Δ = max(Δ, abs(Δ_j))
        end

        if Δ < thresh
            converged = true
            break
        end

        if verbose
            print("\rλ = $λ, s = $s: Δ = $Δ  ($k/$maxiter)")
        end

    end

    converged
end

"""
    elnet_path(R, r, λs, s;
               maxiter = 10_000, thresh = 1e-4)

Solve elastic net for path along λs using warm starts.

HAS CONVERGENCE PROBLEMS. USE elnetg_path INSTEAD

# Arguments

- `R::Matrix` p x p LD matrix
- `r::Vector` p x 1 vector of correlation coefficients
              between phenotype and genotypes
- `λs::Vector` shrinkage parameter path for 1-norm in decreasing order
- `s::Real`   shrinkage parameter for 2-norm (LD)
- `β::Vector` warm start (in) and result (out) for solution vector
- `maxiter::Int` maximum number of iterations
- `thresh::Real` maximum change in β to be called converged

"""
function elnet_path(R, r, λs, s;
                    maxiter = 10_000, thresh = 1e-4)

    nl = length(λs)
    p = length(r)

    β = zeros(Float64, p, nl)
    converged = zeros(Bool, nl)

    converged[1] = elnet!(R, r, λs[1], s, view(β, :, 1);
                         maxiter = maxiter, thresh = thresh)

    for l = 2:nl
        @assert λs[l] < λs[l-1] "λs not in decreasing order"
        β[:, l] .= β[:, l-1] # use warm start
        converged[l] = elnet!(R, r, λs[l], s, view(β, :, l);
                              maxiter = maxiter, thresh = thresh)
    end

    (β = β,
     λs = λs,
     converged = converged,
     s = s)

end


# block_diag.jl
# test

using LinearMapsAA: LinearMapAA, LinearMapAM, LinearMapAO
using SparseArrays: blockdiag, sparse
using Test: @test, @testset

M1 = rand(3,2)
M2 = rand(5,4)
M = [M1 zeros(size(M1,1),size(M2,2));
	zeros(size(M2,1),size(M1,2)) M2]

A1 = LinearMapAA(M1) # WrappedMap
A2 = LinearMapAA(x -> M2*x, y -> M2'y, size(M2), T=eltype(M2)) # FunctionMap
B = block_diag(A1, A2)

@testset "block_diag for AM" begin
	@test Matrix(B) == M
	@test Matrix(B)' == Matrix(B')
end

# test LMAO
@testset "block_diag for AO" begin
	Ao = LinearMapAA(M2 ; odim=(1,5))
	@test block_diag(A1, A1) isa LinearMapAM
	@test block_diag(A1, A1 ; tryop=true) isa LinearMapAO
	@test block_diag(A1, Ao) isa LinearMapAM
	Bo = block_diag(Ao, Ao)
	@test Bo isa LinearMapAO

	Md = blockdiag(sparse(M2), sparse(M2))
	X = rand(4,2)
	Yo = Bo * X
	Yd = Md * vec(X)
	@test vec(Yo) ≈ Yd

	@test Matrix(Bo)' == Matrix(Bo')
end


using QuasiMonteCarlo, Distributions, Test
using Test

#1D
lb = 0.0
ub = 5.0
n = 5
d = 1
QuasiMonteCarlo.sample(n,lb,ub,GridSample(0.1))
QuasiMonteCarlo.sample(n,lb,ub,UniformSample())
QuasiMonteCarlo.sample(n,lb,ub,SobolSample())
QuasiMonteCarlo.sample(n,lb,ub,LatinHypercubeSample())
QuasiMonteCarlo.sample(20,lb,ub,LowDiscrepancySample(10))
QuasiMonteCarlo.sample(5,d,Cauchy())
QuasiMonteCarlo.sample(5,d,Normal(0,4))

#ND
lb = [0.1,-0.5]
ub = [1.0,20.0]
n = 5
d = 2

#GridSample{T}
s = QuasiMonteCarlo.sample(n,lb,ub,GridSample([0.1,0.5]))
@test isa(s,Matrix{typeof(s[1][1])}) == true

#UniformSample()
s = QuasiMonteCarlo.sample(n,lb,ub,UniformSample())
@test isa(s,Matrix{typeof(s[1][1])}) == true

#SobolSample()
s = QuasiMonteCarlo.sample(n,lb,ub,SobolSample())
@test isa(s,Matrix{typeof(s[1][1])}) == true

#LHS
s = QuasiMonteCarlo.sample(n,lb,ub,LatinHypercubeSample())
@test isa(s,Matrix{typeof(s[1][1])}) == true

#LDS
s = QuasiMonteCarlo.sample(n,lb,ub,LowDiscrepancySample([10,3]))
@test isa(s,Matrix{typeof(s[1][1])}) == true

#Distribution 1
s = QuasiMonteCarlo.sample(n,d,Cauchy())
@test isa(s,Matrix{typeof(s[1][1])}) == true

#Distribution 2
s = QuasiMonteCarlo.sample(n,d,Normal(3,5))
@test isa(s,Matrix{typeof(s[1][1])}) == true


"""
Subroutine dtrqsol

This subroutine computes the largest (non-negative) solution
of the quadratic trust region equation

  ||x + sigma*p|| = delta.

The code is only guaranteed to produce a non-negative solution
if ||x|| <= delta, and p != 0. If the trust region equation has
no solution, sigma = 0.

MINPACK-2 Project. March 1999.
Argonne National Laboratory.
Chih-Jen Lin and Jorge J. More'.
"""
function dtrqsol(n,x,p,delta)
    zero = 0.0
    sigma = zero

    ptx = ddot(n,p,1,x,1)
    ptp = ddot(n,p,1,p,1)
    xtx = ddot(n,x,1,x,1)
    dsq = delta^2

    # Guard against abnormal cases.
    rad = ptx^2 + ptp*(dsq - xtx)
    rad = sqrt(max(rad,zero))

    if ptx > zero
        sigma = (dsq - xtx)/(ptx + rad)
    elseif rad > zero
        sigma = (rad - ptx)/ptp
    else
        sigma = zero
    end

    return sigma
end


using BarycentricLagrange
using Compat
using Compat.Test
using Compat.LinearAlgebra

@testset "BarycentricLagrange" begin

    function dotest(F)
        # Chebyshev source nodes
        n = 100
        x = cos.((0:n)*pi/n)
        w = bclag_interp_weights(x)
        # Equispaced evluation nodes
        xx = collect(-1:0.01:1)

        f = F.(x)
        ref = F.(xx)

        @test norm(bclag_interp_eval(x, f, xx)-ref, Inf) <  2e-15
        @test norm(bclag_interp_eval(x, f, xx, w)-ref, Inf) <  2e-15

        M1 = bclag_interp_matrix(x, xx)
        @test norm(M1*f-ref, Inf) <  2e-15

        M2 = bclag_interp_matrix(x, xx, w)
        @test norm(M2*f-ref, Inf) <  2e-15
    end

    @testset "Real-valued" begin
        F = x -> cos(x)
        dotest(F)
    end

    @testset "Complex-valued" begin
        F = x -> exp(1im*x)
        dotest(F)
    end

end


# Copyright 2020-2021 John T. Foster
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
using Test
using TaD
import TaD: evaluate

@testset "evaluate bspline basis functions" begin
    kv = [0, 0, 0, 1, 2, 3, 4, 4, 5, 5, 5.0]
    @test evaluate(BSplineBasis(kv, 2), 5 / 2)[1] ≈ 1 / 8
    @test evaluate(BSplineBasis(kv, 2), 5 / 2)[2] ≈ 3 / 4
    @test evaluate(BSplineBasis(kv, 2), 5 / 2)[3] ≈ 1 / 8
end
@testset "evaluate bspline basis function derivatives" begin
    kv = [0, 0, 0, 0, 0.5, 1, 1, 1, 1]
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[1, 1] ≈ 0.25
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[1, 2] ≈ 0.5
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[1, 3] ≈ 0.25
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[1, 4] ≈ 0.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[2, 1] ≈ -1.5
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[2, 2] ≈ 0.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[2, 3] ≈ 1.5
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[2, 4] ≈ 0.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[3, 1] ≈ 6.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[3, 2] ≈ -12.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[3, 3] ≈ 6.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[3, 4] ≈ 0.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[4, 1] ≈ -12.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[4, 2] ≈ 48.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[4, 3] ≈ -84.0
    @test evaluate(BSplineBasis(kv, 3, k=3), 0.5)[4, 4] ≈ 48.0
end
@testset "evaluate bspline curve and derivatives" begin
    kv = [0, 0, 0, 1, 1, 1]
    b = BSplineBasis(kv, 2, k=1)
    cp = [0 0; 0.5 1; 1 0]
    c = BSplineCurve(b, cp)
    @test evaluate(c, 0.05)[1, 1] ≈ 0.05
    @test evaluate(c, 0.05)[1, 2] ≈ 0.095
    @test evaluate(c, 0.05)[2, 1] ≈ 1.0
    @test evaluate(c, 0.05)[2, 2] ≈ 1.8
    @test evaluate(c, 0.5)[1, 1] ≈ 0.5
    @test evaluate(c, 0.5)[1, 2] ≈ 0.5
    @test evaluate(c, 0.5)[2, 1] ≈ 1.0
    @test evaluate(c, 0.5)[2, 2] ≈ 0.0
end


quadratic_field(10)
k, a = cyclotomic_field(11)
automorphism_group(k)
kt, t = k["t"]
factor(t^2-a)
factor(t^5-a)
k, a = wildanger_field(3, 13)
h = hilbert_class_field(k)
K = number_field(h)
L = simple_extension(K)[1]
absolute_simple_field(L)
discriminant(maximal_order(K))
norm_equation(k, 27)


using DifferentialEquations
using PyCall
using PyPlot

function tusharabmodels()
    # Quick plots of Tushar's a and b
    v = collect(LinRange(1e-15, 1e1, 100000))
    amin = 0.045
    ahighdiff =  0.0319 - amin
    δ = 0.7 # m/s
    aplot = @. amin + ahighdiff * 0.5 * (1 + tanh((log10(v) - log10(5e-3)) / δ))
    bmin = 0.085
    bhighdiff =  0.2811 - bmin
    ##  0.5 * (1 + tanh((log10(v)-np.log10(5e-3)) / 7e-1))
    bplot = @. bmin + bhighdiff * 0.5 * (1 + tanh((log10(v) - log10(5e-3)) / δ))
    plot(v, aplot, "-r", label="a")
    plot(v, bplot, "-b", label="b")
    xscale("log")
    xlabel(L"\log_{10} v")
    ylabel(L"a, b")
    legend()
end

function aginglaw(v, θ, dc)
    return 1 - θ * v / dc
end

function sliplaw(v, θ, dc)
    return -v * θ / dc * log(v * θ / dc)
end

function calcdadv(v)
    amin = 0.045
    ahighdiff =  0.0319 - amin
    δ = 0.7 # m/s
    # a = amin + ahighdiff * 0.5 * (1 + tanh((log10(v) - log10(5e-3)) / δ))
    dadv = ahighdiff * 0.5 * (sech(( -log10(v) + log10(5e-3)) / (δ * log(10))  ))^2 / (δ * v * log(10))
    return dadv
end

function calcdbdv(v)
    bmin = 0.085
    bhighdiff =  0.2811 - bmin
    δ = 0.7 # m/s
    #  bmin + bhighdiff * 0.5 * (1 + tanh((log10(v) - log10(5e-3)) / δ))
    dbdv = bhighdiff * 0.5 * (sech(( -log10(v) + log10(5e-3)) / (δ * log(10))  ))^2 / (δ * v * log(10))
    return dbdv
end

function calcdvθclassic!(du, u, p, t)
    dc, η, σn, a, b, μ, Vp, L, ρ, statelaw = p
    θ = u[1]
    v = u[2]
    du[1] = statelaw(v, θ, dc)
    du[2] = 1 / (η / σn + a / v) * (μ * (Vp - v) / (L * σn) - b * du[1] / θ)
    return nothing
end

function calcdvθab!(du, u, p, t)
    dc, η, σn, μ, vp, L, ρ, statelaw, fstar, vstar, θstar = p
    θ = u[1]
    v = u[2]
    a = u[3]
    b = u[4]
    dθ = du[1]
    dv = du[2]
    du[1] = statelaw(v, θ, dc)
    da = calcdadv(v) * dv
    db = calcdbdv(v) * dv
    numclassic = (μ * (vp - v) / (L * σn) - b * du[1] / θ)
    numab = -da * log(abs(v) / vstar) - db * log(abs(θ) / θstar)
    numσn = 0
    denom = (η / σn + a / v)
    du[2] = (numclassic + numab + numσn) / denom
    du[3] = da
    du[4] = db
    return nothing
end

function plottimeseries(sol, siay, titlelabel)
    t1 = [x / siay for x in sol.t]
    θ1 = [x[1] for x in sol.u]
    v1 = [x[2] for x in sol.u]

    figure(figsize = (12, 6))
    subplot(2, 2, 1)
    plot(t1, θ1, "-b", linewidth = 0.5)
    yscale("log")
    ylabel(L"\theta")

    subplot(2, 2, 2)
    plot(1:1:length(t1), θ1, "-b", linewidth = 0.5)
    yscale("log")
    ylabel(L"\theta")

    subplot(2, 2, 3)
    plot(t1, v1, "-b", linewidth = 0.5)
    yscale("log")
    xlabel("time (years)")
    ylabel("v (m/s)")

    subplot(2, 2, 4)
    plot(1:1:length(t1), v1, "-b", linewidth = 0.5)
    yscale("log")
    xlabel("time (step)")
    ylabel("v (m/s)")
    suptitle(titlelabel)

    return nothing
end

function sliding()
    # Model parameters
    siay = 365.25 * 24 * 60 * 60
    tspan = (0.0, siay * 330.0) # This will run
    # tspan = (0.0, siay * 302.0) # This will run but fail for the a, b case
    μ = 3e10
    ν = 0.25
    ρ = 2700.0
    η = μ / (2.0 * sqrt(μ / ρ))
    L = 60 * 1e3
    a = 0.015
    b = 0.02
    vp = 1e-9
    σn = 30e6
    dc = 0.1
    abstol = 1e-4
    reltol = 1e-4

    # Parameters need for d(a, b σn)/dt
    fstar = 0.6 # ???
    vstar = vp # ???
    θstar = 1e9 # ???

    # Show Tushar scalings for a and b
    tusharabmodels()

    # Time integrate - classic
    icsclassic = [1e8; vp / 1000]
    pclassic = (dc, η, σn, a, b, μ, vp, L, ρ, aginglaw)
    probclassic = ODEProblem(calcdvθclassic!, icsclassic, tspan, pclassic)
    solclassic = solve(probclassic, RK4(), abstol = abstol, reltol = reltol)
    plottimeseries(solclassic, siay, "RSF (classic)")

    # Time integrate - a, b can evolve
    icsab = [1e8 ; vp / 1000 ; a ; b]
    pab = (dc, η, σn, μ, vp, L, ρ, aginglaw, fstar, vstar, θstar)
    probab = ODEProblem(calcdvθab!, icsab, tspan, pab)
    solab = solve(probab, RK4(), abstol = abstol, reltol = reltol)
    plottimeseries(solab, siay, "RSF (ab)")

    return nothing
end
sliding()


#=
Implements the Kalman filter for a linear Gaussian state space model.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-07-29

References
----------

https://julia.quantecon.org/tools_and_techniques/kalman.html

TODO: Do docstrings here after implementing LinerStateSpace
=#

mutable struct Kalman
    A
    G
    Q
    R
    k
    n
    cur_x_hat
    cur_sigma
end

# Initializes current mean and cov to zeros
function Kalman(A, G, Q, R)
    k = size(G, 1)
    n = size(G, 2)
    xhat = n == 1 ? zero(eltype(A)) : zeros(n)
    Sigma = n == 1 ? zero(eltype(A)) : zeros(n, n)
    return Kalman(A, G, Q, R, k, n, xhat, Sigma)
end

function set_state!(k::Kalman, x_hat, Sigma)
    k.cur_x_hat = x_hat
    k.cur_sigma = Sigma
    Nothing
end

@doc doc"""
Updates the moments (`cur_x_hat`, `cur_sigma`) of the time ``t`` prior to the
time ``t`` filtering distribution, using current measurement ``y_t``.
The updates are according to

```math
    \hat{x}^F = \hat{x} + \Sigma G' (G \Sigma G' + R)^{-1}
                    (y - G \hat{x}) \\

    \Sigma^F = \Sigma - \Sigma G' (G \Sigma G' + R)^{-1} G
               \Sigma
```

#### Arguments

- `k::Kalman` An instance of the Kalman filter
- `y` The current measurement

"""
function prior_to_filtered!(k::Kalman, y)
    # simplify notation
    G, R = k.G, k.R
    x_hat, Sigma = k.cur_x_hat, k.cur_sigma

    # and then update
    if k.k > 1
        reshape(y, k.k, 1)
    end
    A = Sigma * G'
    B = G * Sigma * G' + R
    M = A / B
    k.cur_x_hat = x_hat + M * (y .- G * x_hat)
    k.cur_sigma = Sigma - M * G * Sigma
    Nothing
end

"""
Updates the moments of the time ``t`` filtering distribution to the
moments of the predictive distribution, which becomes the time
``t+1`` prior

#### Arguments
- `k::Kalman` An instance of the Kalman filter

"""
function filtered_to_forecast!(k::Kalman)
    # simplify notation
    A, Q = k.A, k.Q
    x_hat, Sigma = k.cur_x_hat, k.cur_sigma

    # and then update
    k.cur_x_hat = A * x_hat
    k.cur_sigma = A * Sigma * A' + Q
    Nothing
end

"""
Updates `cur_x_hat` and `cur_sigma` given array `y` of length `k`.  The full
update, from one period to the next

#### Arguments

- `k::Kalman` An instance of the Kalman filter
- `y` An array representing the current measurement

"""
function update!(k::Kalman, y)
    prior_to_filtered!(k, y)
    filtered_to_forecast!(k)
    Nothing
end

function stationary_values(k::Kalman)
    # simplify notation

    A, Q, G, R = k.A, k.Q, k.G, k.R

    # solve Riccati equation, obtain Kalman gain
    Sigma_inf = solve_discrete_riccati(A', G', Q, R)
    K_inf = A * Sigma_inf * G' * inv(G * Sigma_inf * G' .+ R)
    return Sigma_inf, K_inf
end

"""
computes log-likelihood of period ``t``

##### Arguments
- `kn::Kalman`: `Kalman` specifying the model. Current values must be the
                forecast for period ``t`` observation conditional on ``t-1``
                observation.
- `y::AbstractVector`: Respondentbservations at period ``t``

##### Returns
- `logL::Real`: log-likelihood of observations at period ``t``
"""
function log_likelihood(k::Kalman, y::AbstractVector)
    eta = y - k.G*k.cur_x_hat # forecast error
    P = k.G*k.cur_sigma*k.G' + k.R # covariance matrix of forecast error
    logL = - (length(y)*log(2pi) + logdet(P) .+ eta'/P*eta)[1]/2
    return logL
end

"""
computes log-likelihood of entire observations

##### Arguments
- `kn::Kalman`: `Kalman` specifying the model. Initial value must be the prior
                for t=1 period observation, i.e. ``x_{1|0}``.
- `y::AbstractMatrix`: `n x T` matrix of observed data.
                       `n` is the number of observed variables in one period.
                       Each column is a vector of observations at each period.

##### Returns
- `logL::Real`: log-likelihood of all observations
"""
function compute_loglikelihood(kn::Kalman, y::AbstractMatrix)
    T = size(y, 2)
    logL = 0
    # forecast and update
    for t in 1:T
        logL = logL + log_likelihood(kn, y[:, t])
        update!(kn, y[:, t])
    end
    return logL
end

"""
##### Arguments
- `kn::Kalman`: `Kalman` specifying the model. Initial value must be the prior
                for t=1 period observation, i.e. ``x_{1|0}``.
- `y::AbstractMatrix`: `n x T` matrix of observed data.
                       `n` is the number of observed variables in one period.
                       Each column is a vector of observations at each period.

##### Returns
- `x_smoothed::AbstractMatrix`: `k x T` matrix of smoothed mean of states.
                                `k` is the number of states.
- `logL::Real`: log-likelihood of all observations
- `sigma_smoothed::AbstractArray` `k x k x T` array of smoothed covariance matrix of states.
"""
function smooth(kn::Kalman, y::AbstractMatrix)
    G, R = kn.G, kn.R

    T = size(y, 2)
    n = kn.n
    x_filtered = Matrix{Float64}(undef, n, T)
    sigma_filtered = Array{Float64}(undef, n, n, T)
    sigma_forecast = Array{Float64}(undef, n, n, T)
    logL = 0
    # forecast and update
    for t in 1:T
        logL = logL + log_likelihood(kn, y[:, t])
        prior_to_filtered!(kn, y[:, t])
        x_filtered[:, t], sigma_filtered[:, :, t] = kn.cur_x_hat, kn.cur_sigma
        filtered_to_forecast!(kn)
        sigma_forecast[:, :, t] = kn.cur_sigma
    end
    # smoothing
    x_smoothed = copy(x_filtered)
    sigma_smoothed = copy(sigma_filtered)
    for t in (T-1):-1:1
        x_smoothed[:, t], sigma_smoothed[:, :, t] =
            go_backward(kn, x_filtered[:, t], sigma_filtered[:, :, t],
                        sigma_forecast[:, :, t], x_smoothed[:, t+1],
                        sigma_smoothed[:, :, t+1])
    end

    return x_smoothed, logL, sigma_smoothed
end

"""
##### Arguments
- `kn::Kalman`: `Kalman` specifying the model.
- `x_fi::Vector`: filtered mean of state for period ``t``
- `sigma_fi::Matrix`: filtered covariance matrix of state for period ``t``
- `sigma_fo::Matrix`: forecast of covariance matrix of state for period ``t+1``
                      conditional on period ``t`` observations
- `x_s1::Vector`: smoothed mean of state for period ``t+1``
- `sigma_s1::Matrix`: smoothed covariance of state for period ``t+1``

##### Returns
- `x_s1::Vector`: smoothed mean of state for period ``t``
- `sigma_s1::Matrix`: smoothed covariance of state for period ``t``
"""
function go_backward(k::Kalman, x_fi::Vector,
                     sigma_fi::Matrix, sigma_fo::Matrix,
                     x_s1::Vector, sigma_s1::Matrix)
    A = k.A
    temp = sigma_fi*A'/sigma_fo
    x_s = x_fi + temp*(x_s1-A*x_fi)
    sigma_s = sigma_fi + temp*(sigma_s1-sigma_fo)*temp'
    return x_s, sigma_s
end


struct Camera
    origin::Vector
    # orthonormal basis
    u::Vector
    v::Vector
    w::Vector
    lensradius::Real
    # focus plane window
    lowerleftcorner::Vector
    horizontal::Vector
    vertical::Vector

    """
        Camera(lookfrom, lookat, vup, vfov, aspect, aperture, focus_dist)

    Initialize camera position, orientation, field of view and aperture.

    # Arguments:
    - lookfrom: camera location within scene
    - lookat: coordinates towards which camera is oriented
    - vup: "up" direction
    - vfov: "field of view" angle, in rad
    - aspect: aspect ratio (width / height) of focus window
    - aperture: aperture (diameter) of camera lense
    - focus_dist: distance of focus plane from camera
    """
    function Camera(lookfrom::Vector, lookat::Vector, vup::Vector, vfov::Real, aspect::Real, aperture::Real, focus_dist::Real)
        # vfov is top to bottom in rad
        half_height = tan(vfov/2)
        half_width = aspect * half_height
        # orthonormal basis
        w = unitvector(lookfrom - lookat)
        u = unitvector(cross(vup, w))
        v = cross(w, u)
        # define the focus plane window
        lowerleftcorner = (lookfrom
            - half_width *focus_dist*u
            - half_height*focus_dist*v
            -             focus_dist*w)

        new(lookfrom, u, v, w, aperture / 2, lowerleftcorner,
            2*half_width *focus_dist*u,
            2*half_height*focus_dist*v)
    end
end

"""
    getray(camera, s, t)

Get a ray originating from a random position on the lens (to imitate
depth of field), targeting the focus window at relative coordinates.

# Arguments:
- `camera::Camera`: camera
- `s::Real`: relative x-coordinate within focus window (real number between 0 and 1)
- `t::Real`: relative y-coordinate within focus window (real number between 0 and 1)
"""
function getray(camera::Camera, s::Real, t::Real)
    rd = camera.lensradius * random_in_unit_disk()
    offset = rd[1]*camera.u + rd[2]*camera.v
    ray_origin = camera.origin + offset
    direction = camera.lowerleftcorner + s*camera.horizontal + t*camera.vertical - ray_origin
    return Ray(ray_origin, direction)
end


push!(LOAD_PATH, string(pwd(),"/src"))

import ModifiedStokesSolver
using Base.Test

mss = ModifiedStokesSolver

@testset "Split" begin

    srand(1)
    r = 1 .- 2.*(rand(), rand())
    f = 1 .- 2.*(rand(), rand())
    n = 1 .- 2.*(rand(), rand())
    rnorm = sqrt(r[1]^2+r[2]^2)

    cutoff = mss.POWERSERIES_ZMAX

    @testset "T1-T3" begin
        # Test split of stresslet kernel functions
        z = 0.987 # This value should work well for both
        (T1,T2,T3) = mss.Ti_direct(z)
        (T1S, T1L, T2S, T2L, T3S, T3L) =
            mss.Ti_split(z)
        E1 = (T1S + T1L*log(z) - T1)/T1
        E2 = (T2S + T2L*log(z) + 1/(8*pi*z^2) - 1/(pi*z^4)
              - T2)/T2
        E3 = (T3S + T3L*log(z) - T3)/T3

        @test abs(E1) < 1e-14
        @test abs(E2) < 1e-14
        @test abs(E3) < 1e-14

        if cutoff>0
            # Test with dispatcher for large and small arguments
            # 3.5 is the largest argument that works with explicit split
            for z=[cutoff/2 cutoff*2 1e-5 3.5]
                @show z
                (T1,T2,T3) = mss.Ti(z)
                (T1S, T1L, T2S, T2L, T3S, T3L) =
                    mss.Ti_split(z)
                E1 = (T1S + T1L*log(z) - T1)/T1
                E2 = (T2S + T2L*log(z) + 1/(8*pi*z^2) - 1/(pi*z^4)
                      - T2)/T2
                E3 = (T3S + T3L*log(z) - T3)/T3

                @test abs(E1) < 2e-14
                @test abs(E2) < 1e-14
                @test abs(E3) < 1e-14

            end
        end
    end

    @testset "Derivatives" begin
        @testset "Split" begin
            z = 0.9999*ModifiedStokesSolver.POWERSERIES_ZMAX_PRIME
            err = abs.( (mss.Ti_split_pow(z, true) .- mss.Ti_split_direct(z, true)) ./ mss.Ti_split_direct(z, true) )
            @test maximum(err) < 1e-13
        end
        @testset "Recombine" begin
            z = 0.99*ModifiedStokesSolver.POWERSERIES_ZMAX_PRIME
            err = abs.( (mss.Ti_prime(z) .- mss.Ti_direct_prime(z)) ./ mss.Ti_direct_prime(z) )
            @test maximum(err) < 3e-14
        end
    end

    @testset "Stresslet" begin
        # Test split stresslet
        r = rand(2)
        f = rand(2)
        n = rand(2)
        alpha = rand()
        u1 = mss.stresslet(r, f, n, alpha)
        (uS, uL, uC, uQ) = mss.stresslet_split(r, f, n, alpha)
        rnorm = norm(r)
        rdotf = dot(r,f)
        rdotn = dot(r,n)
        u2 = uS + uL*log(rnorm) + uC*rdotn/rnorm^2 +
            uQ*rdotf*r*rdotn/rnorm^4
        E = norm(u1-u2,Inf)/norm(u2,Inf)
        @test E < 1e-14
    end
end


#!/usr/bin/julia

using GLM, CSV, StatsModels, DataFrames, DataFramesMeta, Random, Statistics

mse(y1, y2) = mean((y1 - y2) .^ 2)

df = CSV.read("../csv/corei7_long.csv")

iterations = 100

cross_validation = DataFrame(MSE = [],
                             R2 = [],
                             PredictedParallel = [],
                             RealCompileTime = [],
                             PredictedCompileTime = [])

for i = 1:iterations
    println(i)
    train_ratio = 0.1
    train_size = round(Int, train_ratio * nrow(df))
    train_rows = shuffle(1:nrow(df))[1:train_size]
    df_train = df[train_rows, :]

    test_rows = [x for x in 1:nrow(df) if !(x in train_rows)]
    df_test = df[test_rows, :]

    regression = lm(@formula(CompileTime ~ ((Expected_insns +
                                             Expected_insns ^ 2) +
                                            (Functions +
                                             Functions ^ 2) +
                                            (Inlined_percentage +
                                             Inlined_percentage ^ 2) +
                                            (Num_partitions +
                                             Num_partitions ^ 2)) *
                             (Parallel + Parallel ^ 2)), df_train)

    prediction = predict(regression, select(df_test, Not(:CompileTime)))

    best_prediction = df_test[findmin(prediction)[2], :]

    push!(cross_validation, (MSE = mse(df_test[:, :CompileTime], prediction),
                             R2 = adjr2(regression),
                             PredictedCompileTime = findmin(prediction)[1],
                             RealCompileTime = best_prediction[:CompileTime],
                             PredictedParallel = best_prediction[:Parallel]))

end

CSV.write("../csv/cross_validation.csv", cross_validation)

cv_df = CSV.read("../csv/cross_validation.csv")
cv_df

println(mean(cv_df.MSE))
println((1.96 * std(cv_df.MSE)) / sqrt(length(cv_df.MSE)))

println(mean(cv_df.R2))
println((1.96 * std(cv_df.R2)) / sqrt(length(cv_df.R2)))


function centroid_(points)
    com = sum(points)
    return com/size(points, 1)
end

function covmat_(points, com)
    @assert size(points,1) > 1
    covmat = zeros(3,3)
    for i in eachindex(points)
        pt = points[i]-com

        covmat[2,2] += pt[2]*pt[2]
        covmat[2,3] += pt[2]*pt[3]
        covmat[3,3] += pt[3]*pt[3]

        covmat[1,1] += pt[1]*pt[1]
        covmat[1,2] += pt[1]*pt[2]
        covmat[1,3] += pt[1]*pt[3]
    end
    covmat[2,1] = covmat[1,2]
    covmat[3,1] = covmat[1,3]
    covmat[3,2] = covmat[2,3]
    return covmat
end

function covmat_(points)
    c = centroid_(points)
    return covmat_(points, c)
end

"""
    normedcovmat_(points, centr)

Compute the normed covariance matrix with the given centroid.
"""
function normedcovmat_(points, c)
    m = covmat_(points, c)
    return m/size(points,1)
end

"""
    normedcovmat_(points)

Compute the normed covariance matrix of the points.
"""
function normedcovmat_(points)
    c = centroid_(points)
    m = covmat_(points, c)
    return m/size(points,1)
end

"""
    vec2homvec_(ps)

Convert 3D vectors to homogeneous vectors: append 1 to each vector.
Return a vector of SVectors.
"""
function vec2homvec_(ps)
    return [SVector{4,Float64}(p..., 1) for p in ps]
end

"""
    homvec2vec_(ps)

Convert 3D homogeneous vectors to "normal" vectors: drop the 1 from the end.
Return a vector of SVectors.
"""
function homvec2vec_(ps)
    return [SVector{3,Float64}(p[1:3]) for p in ps]
end

"""
    findOBB(points)

Compute the oriented bounding box of points (in 3D).
Return the cornerpoints.
"""
function findOBB_(points)
    c = centroid_(points)
    m = normedcovmat_(points, c)
    evs = eigvecs(m)
    zv = cross(evs[:,1], evs[:,2])
    if abs(dot(zv, evs[:,3])) < cosd(5)
        @warn "evs[3]: $(evs[3]) and zv: $zv are not parallel!"
    end
    evs[:,3] = zv
    # transformation matrix
    rmat = zeros(4,4)
    rmat[4,4] = 1
    rmat[1:3,1:3] = evs'
    rmat[1:3,4] = -1 .* evs' * c
    trm = SMatrix{4,4,Float64}(rmat)

    p_transformed = [trm*p for p in vec2homvec_(points)]
    aabb_tr = homvec2vec_(findAABB(p_transformed))
    meanDiagonal = sum(aabb_tr)/2

    # size of oriented bounding box
    bbs = aabb_tr[2]-aabb_tr[1]
    nv = -bbs/2
    # box around origo
    bi = 0:1
    corners_ = [nv+bbs .*[i,j,k] for i in bi for j in bi for k in bi]
    corners = vec2homvec_(corners_)

    # final transformation
    finalTR = zeros(4,4)
    finalTR[4,4] = 1
    finalTR[1:3,1:3] = evs
    finalTR[1:3,4] = evs*meanDiagonal+c
    ftr = SMatrix{4,4,Float64}(finalTR)

    obbcorners_ = [ftr*p for p in corners]
    obbcorners = homvec2vec_(obbcorners_)

    # points are selected so that the normals point outside from the cube
    #planetris = [[1,5,2], [1,2,3], [1,3,5],
    #            [8,4,6], [8,6,7], [8,7,4]]

    return obbcorners
end


#Gingrich stocking chart calculated using the original upland oak parameters
#for more information see:http://oak.snr.missouri.edu/silviculture/tools/gingrich.html

using RecipesBase

amd_convert(qmd) = -0.259 + 0.973qmd
get_tpa(j, i, a) = (1*j*10)/(a[1]+a[2]*amd_convert(i)+a[3]*i^2)

function stklines(stk, dia, a)
  tpa = get_tpa.(stk', dia, Ref(a))
  ba = pi * (dia ./ 24).^2 .* tpa
  tpa, ba
end

@userplot gingrich_chart
@recipe function f(dat::gingrich_chart)

  #define parameters
  a = [-0.0507, 0.1698, 0.0317]
  b = [0.175, 0.205, 0.06]
  ba_ticks= 0:20:160
  tpa_ticks= 0:50:450
  #QMD lines
  dia = [7, 8, 10, 12, 14, 16, 18, 20, 22]
  #stocking percent horizontal lines
  stk_percent= 20:10:110

  # plot attributes
  xlabel --> "Trees Per Acre"
  ylabel --> "Basal Area (sq ft./acre)"
  xticks --> 0:50:450
  yticks --> 0:20:200
  legend := false
  primary := false

  tpa, ba = stklines(stk_percent, dia, a)

  # horizontal lines
  @series begin
    linecolor := :black
    tpa, ba
  end

  # vertical lines
  @series begin
    linecolor := :black
    tpa', ba'
  end

  # upper red line
  @series begin
    linecolor := :red
    stklines([100], dia, a)
  end

  # lower red line
  @series begin
    linecolor := :red
    stklines([100], dia, b)
  end

  # the point and the annotations
  @series begin
    primary := true
    seriestype := :scatter
    annotations := (100:38:442, 20:10:110, ["$i%" for i in 20:10:110])

    a,b = dat.args
    [a], [b]
  end
end


using CharFuncPricing, Test
using TaylorSeries

@testset "SZLogContinuity" begin
 v0 = 0.08
 ρ=-0.8
 σ=0.25
 θ=0.1
 κ=3.0
 #params from Cui et al for Heston.
 params = SchobelZhuParams{Float64}(sqrt(v0), κ, sqrt(θ), ρ, σ)
 cf = DefaultCharFunc(params)
 #with the original formula from SZ, there is a discontinuity in the log.
 sz = CharFuncPricing.evaluateLogCharFunc( cf, 2.0+0.0im, 15.0)
 lk = CharFuncPricing.evaluateLogCharFuncLK( cf, 2.0+0.0im, 15.0)
 @test isapprox(lk,sz, atol=1e-13)
 for u = 0.1:0.1:4.0
      lk = CharFuncPricing.evaluateLogCharFuncLK(cf, u+0.0im, 15.0)
      sz=CharFuncPricing.evaluateLogCharFunc( cf, u+0.0im, 15.0)
      @test isapprox(lk,sz, atol=1e-13)
  end

  v0 = 0.010201
  ρ  = -0.7
  σ = 0.61
  κ = 6.21
  θ = 0.019
  α = 3.35861
  params = SchobelZhuParams{Float64}(sqrt(v0), κ, sqrt(θ), ρ, σ)
  cf = DefaultCharFunc(params)
   #with the original formula from SZ, there is a discontinuity in the log.
  sz = CharFuncPricing.evaluateLogCharFunc(cf, 10.0+α*1im, 10.0)
  lk = CharFuncPricing.evaluateLogCharFuncLK(cf, 10.0+α*1im, 10.0)
  @test isapprox(lk,sz, atol=1e-13)

end

@testset "SZCumulantsJoshi" begin
    κ = 4.0
    θ = 0.25
    σ = 1.0
    ρ = -0.5
    v0 = 0.1
    τ = 0.01
    params = SchobelZhuParams{Float64}(v0, κ, θ, ρ, σ)
    cf = DefaultCharFunc{SchobelZhuParams{Float64}, Taylor1{Complex}}(params)
    t = Taylor1(Float64, 4)
    cft = evaluateLogCharFunc(cf, t, τ)
    c1, c2, c4 = computeCumulants(params, τ)
    println(c1, " ", c2 / 2, " ", c4 / (2 * 3 * 4), " ", cft)
    @test isapprox(imag(cft.coeffs[2]), c1, atol = 1e-12)
    @test isapprox(-real(cft.coeffs[3]) * 2, c2, atol = 1e-12)
    @test isapprox(cft.coeffs[5] * 2 * 3 * 4, c4, atol = 1e-16)
end

@testset "SchobelZhuPaper" begin
    #reference values from "Stochastic Volatility with an Ornstein-Uhlenbeck Process: An Extension" (1998)
    κ = 4.0
    θ = 0.2
    σ = 0.1
    ρ = -0.5
    v0 = 0.2
    τ = 0.5
    spot = 100.0
    r = 0.0953
    forward = spot * exp(r * τ)
    df = exp(-r * τ)
    params = SchobelZhuParams{Float64}(v0, κ, θ, ρ, σ)
    cf = DefaultCharFunc(params)
    m = 128
    l = 6
    pricer =
        makeCosCharFuncPricer(cf, τ, m, l)
    strikes = [90.0, 95.0, 100.0, 105.0, 110.0, 115.0, 120.0]
    refPrices = [15.29, 11.50, 8.24, 5.60, 3.58, 2.16, 1.22]
    for (strike, refPrice) in zip(strikes, refPrices)
        price = priceEuropean(pricer, true, strike, forward, τ, df)
        println(strike, " ", price)
        @test isapprox(refPrice, price, atol = 0.01)
    end
    θ = 0.0001
    v0 = 0.15
    params = SchobelZhuParams{Float64}(v0, κ, θ, ρ, σ)
    cf = DefaultCharFunc(params)

    pricer =
        makeCosCharFuncPricer(cf, τ, m, l)
    refPrices = [14.22, 9.60, 5.37, 2.15, 0.50, 0.06, 0.004]
    for (strike, refPrice) in zip(strikes, refPrices)
        price = priceEuropean(pricer, true, strike, forward,τ, df)
        println(strike, " ", price)
        @test isapprox(refPrice, price, atol = 0.01)
    end

end


# ____________________________________________________________________
# SpecLineAsymmGauss
#
mutable struct SpecLineAsymmGauss <: AbstractComponent
    norm::Parameter
    center::Parameter
    fwhm::Parameter
    voff::Parameter
    asymm::Parameter
    index::Vector{Int}  # optimization

    function SpecLineAsymmGauss(center::Number)
        out = new(Parameter(1),
                  Parameter(center),
                  Parameter(3000),
                  Parameter(0),
                  Parameter(0),
                  Vector{Float64}(), true)

        @assert center > 0
        out.norm.low = 0
        out.center.low = 0
        out.fwhm.low = 0
        out.voff.low = 0
        out.center.fixed = true
        return out
    end
end

function prepare!(comp::SpecLineAsymmGauss, domain::Domain{1})
    comp.index = collect(1:length(domain))
    return fill(NaN, length(domain))
end

function evaluate!(buffer, comp::SpecLineAsymmGauss, x::Domain{1},
                   norm, center, fwhm, voff, asymm)
    buffer[comp.index] .= 0.
    empty!(comp.index)

    x0 = center - (voff / 3.e5) * center
    hwhm = fwhm / 3.e5 * center / 2  # Note: this is in `center` units

    sigma0 = hwhm / (2.355 / 2)
    sigma = 2. * sigma0 ./ (1 .+ exp.(asymm .* (x .- x0) ./ 2 ./ sigma0))
    X = (x .- x0) ./ sigma
    i = findall(abs.(X) .< 4) # optimization
    append!(comp.index, i)
    buffer[i] .= norm * exp.(-X[i].^2 ./ 2) ./ sqrt(2pi) * sigma0
end


free_energies(chem_potentials::AbstractVector, stoich_matrix::AbstractMatrix) =
    vec(chem_potentials'*stoich_matrix)

function free_energies(chem_potentials::AbstractVector, reactions::Vector{Reaction})
    ∇r, ∇p = stoichiometric_matrix(ReactionNetwork(reactions))

    return free_energies(chem_potentials, ∇r - ∇p)
end

# compute reaction rates from:
# free energy,
# overall rate (sqrt(kf*kr)),
# inverse temperature
function reaction_rates(free_e, ovr_rate, beta)
    kf = ovr_rate.*exp.(-0.5*beta*free_e)
    kr = ovr_rate.*exp.( 0.5*beta*free_e)
    return kf,kr
end


using JuMP, Ipopt

@testset "dat export import consistency check" begin
    # min x - y
    # st  x + x^2 + x*y + y^2 <= 1
    #     -2 <= x, y <= 2
    # solution: x+y = -1/3
    # optimal objective -1-4/sqrt(3)

    qcqp = Problem()

    x = MPC.Variable("x", Complex)
    y = MPC.Variable("y", Real)
    b = MPC.Variable("b", Bool)

    set_objective!(qcqp, x - y)

    add_constraint!(qcqp, "ctr1", (x + Base.power_by_squaring(x,2) + x*y + Base.power_by_squaring(y,2)) << 1)
    add_constraint!(qcqp, "ctr2", (Base.power_by_squaring(x,3) + (1.0+3im)*conj(Base.power_by_squaring(x,2))*x*y + Base.power_by_squaring(y,2)) << (1 + 2im))
    add_constraint!(qcqp, "ctr3", -3im << (5*b + 5im*y) << 1)
    add_constraint!(qcqp, "x_bounds", -2 << x << 2)
    add_constraint!(qcqp, "y_bounds", -2 << y << 2)

    point = Point([x], [3])
    dat_exportpath = joinpath(pwd(),"tmp_datexport")
    # dat_exportpath = joinpath(Pkg.dir("MathProgComplex"), "tmp_datexport")
    mkpath(splitdir(dat_exportpath)[1])

    export_to_dat(qcqp, dat_exportpath, filename="POP.dat", point=point)

    qcqp2, point2 = import_from_dat(joinpath(dat_exportpath, "POP.dat"))

    rm(joinpath(dat_exportpath, "POP.dat"))

    @test qcqp.variables == qcqp2.variables
    @test qcqp.objective == qcqp2.objective

    @test Set(keys(qcqp.constraints)) == Set(keys(qcqp2.constraints))

    for (ctrname, ctr) in qcqp.constraints
        @test qcqp.constraints[ctrname].ub == qcqp2.constraints[ctrname].ub
        @test qcqp.constraints[ctrname].lb == qcqp2.constraints[ctrname].lb
        @test qcqp.constraints[ctrname].p == qcqp2.constraints[ctrname].p
    end

    @test point == point2

    rm(dat_exportpath, recursive=true)
end

@testset "WB5.dat import and Ipopt solve" begin
    instancepath = joinpath(pwd(), "test", "instances")
    # instancepath = joinpath(Pkg.dir("MathProgComplex"), "test", "instances")
    WB5_cplx, initpt = import_from_dat(joinpath(instancepath, "WB5.dat"))

    WB5 = pb_cplx2real(WB5_cplx)

    mysolver = Ipopt.Optimizer
    m, variables_jump, ctr_jump, ctr_exp = get_JuMP_cartesian_model(WB5, mysolver)
    optimize!(m)

    sol = get_JuMP_solution(m, variables_jump, WB5)

    @test get_objective(WB5, sol) ≈ 1146.478 atol=1e-3
end


export sdiag, spoutput

import Base.kron
import SparseArrays
export kron

function sdiag(a::Vector)
# S = sdiag(s) builds sparse diagonal matrix
	n = length(a)
	i = collect(1:n+1) # colptr
	j = collect(1:n)   # rowval
	return SparseMatrixCSC(n,n,i,j,a)
end

# kron(v::SparseVector,A::SparseMatrixCSC) = kron(SparseMatrixCSC(v),A)
# kron(A::SparseMatrixCSC,v::SparseVector) = kron(A,SparseMatrixCSC(v))

function kron(v1::SparseVector,v2::SparseVector)
  v = kron(SparseMatrixCSC(v1),SparseMatrixCSC(v2))
   return SparseVector(v.n,v.nzind,v.nzval)
end

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

function spoutput( filename::String,
                   A::SparseMatrixCSC )
# Output a 3 (or 4 for complex) column sparse matrix file.

f = open(filename, "w")
n = size(A,2)

complexvalue = typeof(A.nzval[1]) == ComplexF64

for ir = 1:n
   j1 = A.colptr[ir]
   j2 = A.colptr[ir+1] - 1

   if complexvalue
      for ic = j1:j2
         println(f, A.rowval[ic], " ", ir, " ", real(A.nzval[ic]), " ", imag(A.nzval[ic]) )
      end # ic
   else
      for ic = j1:j2
         println(f, A.rowval[ic], " ", ir, " ", A.nzval[ic] )
      end
   end # ic

end # ir

close(f)
return
end # function spoutput


# NegativeBinomial is the distribution of the number of failures
# before the r-th success in a sequence of Bernoulli trials.
# We do not enforce integer size, as the distribution is well defined
# for non-integers, and this can be useful for e.g. overdispersed
# discrete survival times.

immutable NegativeBinomial <: DiscreteUnivariateDistribution
    r::Float64
    prob::Float64

    function NegativeBinomial(r::Real, p::Real)
        zero(p) < p <= one(p) || error("prob must be in (0, 1].")
        zero(r) < r || error("r must be positive.")
        new(float64(r), float64(p))
    end

    NegativeBinomial() = new(1.0, 0.5)
end

@_jl_dist_2p NegativeBinomial nbinom

isupperbounded(::Union(NegativeBinomial, Type{NegativeBinomial})) = false
islowerbounded(::Union(NegativeBinomial, Type{NegativeBinomial})) = true
isbounded(::Union(NegativeBinomial, Type{NegativeBinomial})) = false

minimum(::Union(NegativeBinomial, Type{NegativeBinomial})) = 0
maximum(::Union(NegativeBinomial, Type{NegativeBinomial})) = Inf

insupport(::NegativeBinomial, x::Real) = isinteger(x) && zero(x) <= x
insupport(::Type{NegativeBinomial}, x::Real) = isinteger(x) && zero(x) <= x

function probs(d::NegativeBinomial, rgn::UnitRange)
    r = d.r
    p0 = 1.0 - d.prob
    f, l = rgn[1], rgn[end]
    0 <= f <= l || throw(BoundsError())
    res = Array(Float64, l - f + 1)
    res[1] = v = pdf(d, f)
    b = f - 1
    for x = f+1:l
        c = (x + r - 1) * p0 / x
        res[x-b] = (v *= c)
    end
    return res
end

function mgf(d::NegativeBinomial, t::Real)
    r, p = d.r, d.prob
    return ((1.0 - p) * exp(t))^r / (1.0 - p * exp(t))^r
end

function cf(d::NegativeBinomial, t::Real)
    r, p = d.r, d.prob
    return ((1.0 - p) * exp(im * t))^r / (1.0 - p * exp(im * t))^r
end

function mean(d::NegativeBinomial)
    p = d.prob
    (1.0 - p) * d.r / p
end

function var(d::NegativeBinomial)
    p = d.prob
    (1.0 - p) * d.r / (p * p)
end

function std(d::NegativeBinomial)
    p = d.prob
    sqrt((1.0 - p) * d.r) / p
end

function skewness(d::NegativeBinomial)
    p = d.prob
    (2.0 - p) / sqrt((1.0 - p) * d.r)
end

function kurtosis(d::NegativeBinomial)
    p = d.prob
    6.0 / d.r + (p * p) / ((1.0 - p) * d.r)
end

function mode(d::NegativeBinomial)
    p = d.prob
    ifloor((1.0 - p) * (d.r - 1.) / p)
end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ 5eebac68-891b-11eb-37ba-c1d9481c6134
begin
  using Pkg; Pkg.activate("MLJ_env", shared=true)
	using Test
	using Random
	using PlutoUI

	push!(LOAD_PATH, "./src")
	using YaCounter
end

# ╔═╡ 87e5e2ec-8915-11eb-362b-6bf13a36b8e4
md"""
## Decision Trees (DT)

ref. from book **"Data Science from Scratch"**, Chap 17

$(html"
© Pascal, April 2021
")
"""

# ╔═╡ 8e2aca8e-8a8f-11eb-1e68-1316740f4697
begin
	const F = Float64
	const VF = AbstractVector{F}
	const AVT{T} = AbstractVector{T} where {T <: Any};
end

# ╔═╡ cb25501a-8915-11eb-3626-a1ae6d233f92
PlutoUI.TableOfContents(indent=true, depth=4, aside=true)

# ╔═╡ cb0838ae-8915-11eb-082f-3991192a101f
html"""

"""

# ╔═╡ caf03ada-8915-11eb-2935-01e110accba6
html"""

"""

# ╔═╡ cad66356-8915-11eb-0173-ada200678f63
md"""
### Entropy

In mathematical terms, if pᵢ is the proportion of data labeled as class cᵢ, then the entropy is defined as:

$$H(S) = \sum_i -p_i \times log_2(p_i)$$

With the standard convention: $0 \times log_2(0) = 0$

Each term is non-negative and is close to 0 when pᵢ is either close 0 or close to 1.
This means the entropy will be small when every pᵢ is close to 0 or 1 (*i.e.* when most of the data is in 1 class) and it will be larger when many of the pᵢ's are close to 0 (*i.e.* when the data is spread across multiple classes).
"""

# ╔═╡ dc3b7280-8915-11eb-258a-1789476897b9
function entropy(class_prob::VF)::F
	λ = p -> p > zero(F) ? -p * log(2, p) : zero(F)
	sum(λ.(class_prob))
end

# ╔═╡ dc209b9a-8915-11eb-0f63-ab3b7c33793c
begin
	@test entropy([1.]) ≈ 0.        # minimal entropy (max. certainty)
	@test entropy([.5, .5]) ≈ 1.    # maximal entropy for 2 classes
	@test 0.81 < entropy([.25, .75]) < 0.82
 end

# ╔═╡ d7aade44-891b-11eb-00d3-abe7406c4b09
md"""
Our data will consist of pairs(input, label) for which we will need to compute the class probabilities.
"""

# ╔═╡ 9ae99efe-891c-11eb-3d1f-89182e3f9ff8
begin
	function class_prob(labels::AVT)::VF
		@assert length(labels) > 0
		values(Counter(labels)) / length(labels)
	end

	function data_entropy(labels::AVT)::F
		class_prob(labels) |> entropy
	end
end

# ╔═╡ ea1f4a2c-891d-11eb-3027-bd447dd9d1e1
begin
	@test data_entropy(["a"]) ≈ 0.
	@test data_entropy([:a]) ≈ 0.
	@test data_entropy([:b, :a, :b]) ≈ entropy([1. / 3., 2. / 3.])
	@test data_entropy([true, false]) == 1.
	@test data_entropy([2, 1, 2, 2]) == entropy([0.25, 0.75])
	#
	bv = BitVector[[1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 0, 1, 0]]
	@test data_entropy(bv) ≈ 1.5849625007
end

# ╔═╡ 13e3ccb6-8aae-11eb-3a9e-85b5afc74097
class_prob([1, 1, 1, 1, 1])

# ╔═╡ 13c729da-8aae-11eb-29aa-b171020ad8dd
class_prob([0, 0, 0, 1, 1]), class_prob([1, 1, 0, 1, 0])

# ╔═╡ 13b2940c-8aae-11eb-2dd2-8dec30dd4881
data_entropy([0, 0, 0, 1, 1]), data_entropy([1, 1, 0, 1, 0])

# ╔═╡ dc0c5a5e-8915-11eb-2c62-5328b42add69
html"""

"""

# ╔═╡ dbefe77a-8915-11eb-39f1-5d0c84729739
md"""
### Entropy Partition

Mathematically, if we split our data S into partitions S₁, ...Sₘ containing proportions q₁, ...qₘ of the data, then we compute the entropy of the partition as a weighted sum:

$$H = \sum_{i=1}^m q_i \times H(S_i)$$
"""

# ╔═╡ dbd86924-8915-11eb-351f-8362f09ba984
function partition_entropy(subsets)::F
	"""
	Given the partition into subsets, calc. its entropy
	"""
	tot_cnt = sum(length.(subsets))
	λ = s -> (data_entropy(s) * length(s)) / tot_cnt
	sum(λ.(subsets))
end

# ╔═╡ d32b86e0-8928-11eb-193c-e3d7c85d23dd
begin
	a_ = BitVector[[1, 1, 1, 1], [0, 0, 0, 1, 1], [1, 1, 0, 1, 0]]
	@test abs(partition_entropy(a_) - 0.6935361388961) ≤ 1e-6
end

# ╔═╡ dbbf22a2-8915-11eb-00eb-4b0278c0283d
html"""

"""

## ================================

# ╔═╡ 5f5e6d02-892a-11eb-32d8-dfc1a7f259fa
md"""
### Creating our DT
"""

# ╔═╡ 073f6442-892e-11eb-0fc2-bbeb1d12619e
begin
const NB = Union{Nothing, Bool}

struct Candidate
	level::Symbol
	lang::Symbol
	tweets::Bool
	phd::Bool
	did_well::NB
	function Candidate(level::Symbol, lang::Symbol, tweets::Bool,
				phd::Bool, did_well::NB=nothing)
		new(level, lang, tweets, phd, did_well)
	end
end

end

# ╔═╡ 2ec41c48-8a95-11eb-13b6-61f4a07a31db
inputs = [
	#         level    lang     tweets phd    did_well
	Candidate(:Senior, :Java,   false, false, false),
	Candidate(:Senior, :Java,   false, true,  false),
	Candidate(:Mid,    :Python, false, false, true),
	Candidate(:Junior, :Python, false, false, true),
	Candidate(:Junior, :R,      true,  false, true),
	Candidate(:Junior, :R,      true,  true,  false),
	Candidate(:Mid,    :R,      true,  true,  true),
	Candidate(:Senior, :Python, false, false, false),
	Candidate(:Senior, :R,      true,  false, true),
	Candidate(:Junior, :Python, true,  false, true),
	Candidate(:Senior, :Python, true,  true,  true),
	Candidate(:Mid,    :Python, false, true,  true),
	Candidate(:Mid,    :Java,   true,  false, true),
	Candidate(:Junior, :Python, false, true,  false)
];

# ╔═╡ 25fb8f9a-8952-11eb-35bb-b173ac466835
md"""
We will build a decision tree (DT) following ID3 algorithm, which works as follows:
  - if the data have all the same label, create a leaf node that predicts that label and stops.
  - if the list of attributes is empty (*i.e* no more questions to split the data on), create a leaf that predicts the most common lable and stops
  - otherwise try partitionning the data by each of the attributes
  - choose the partition with the lowest entropy
  - add a decision node based on the chosen attribute
  - using the remaining attributes, recursively apply previous steps on each subset

First let's go manually through those steps using our toy dataset.
"""

# ╔═╡ 045023be-8aac-11eb-29c1-3f4da066ff9b
## Tuple
fieldnames(Candidate)[1:end-1]

# ╔═╡ b1c0a416-8aba-11eb-140a-9f3c4eb93346
## Vector
Symbol[fieldnames(Candidate)[1:end-1]...]

# ╔═╡ 65ec8554-8953-11eb-3973-b56039754312
function partition_by(inputs::AVT, attr::Symbol)::Dict{Symbol, AVT}
	parts = Dict{Symbol, AVT}()
	for input ∈ inputs
		kval = getfield(input, attr)
		parts[Symbol(kval)] = push!(get(parts, Symbol(kval), []), input)
	end
	parts
end

# ╔═╡ 0c8055b8-89b2-11eb-3046-331119a0dc9b
function partition_entropy_by(inputs::AVT, attr::Symbol, label_attr::Symbol)::F
	"""
	Given the partition, calc. its entropy
	"""
	parts = partition_by(inputs, attr)
	# @show(parts, attr)
	# println("-----------------------")

	λ = input -> getfield(input, label_attr)
	labels = [λ.(vp) for vp ∈ values(parts)]
	# @show(labels)
	# println("-----------------------")

	partition_entropy(AVT[labels...])
end

# ╔═╡ d728d526-8956-11eb-3c22-e788d025e8b4
with_terminal() do
	attr = :did_well
	for key ∈ fieldnames(Candidate)[1:end-1]
		r = partition_entropy_by(inputs, key, attr)
		println("$(key)/$(attr) => $(r)")
	end
end

# ╔═╡ 17768334-8a98-11eb-1172-3d5079435dcf
begin
	@test 0.69 ≤ partition_entropy_by(inputs, :level, :did_well) < 0.7
	@test 0.86 ≤ partition_entropy_by(inputs, :lang, :did_well) < 0.87

	@test 0.78 ≤ partition_entropy_by(inputs, :tweets, :did_well) < 0.79
	@test 0.89 ≤ partition_entropy_by(inputs, :phd, :did_well) < 0.90
end

# ╔═╡ 8365f5cc-8aaf-11eb-1928-b3a201b840a0
begin
	senior_inputs = filter(c -> getfield(c, :level) == :Senior, inputs)

	@test partition_entropy_by(senior_inputs, :lang, :did_well) ≈ 0.4
	@test partition_entropy_by(senior_inputs, :tweets, :did_well) ≈ 0.0
	@test 0.95 ≤ partition_entropy_by(senior_inputs, :phd, :did_well) ≤ 0.96
end

# ╔═╡ 432cd594-89b9-11eb-2d4a-6f46f08a511d
html"""

"""

# ╔═╡ 48db121e-8a8e-11eb-189f-d9f41e4b1d9d
md"""
### Putting it all Together

"""

# ╔═╡ 48ba469c-8a8e-11eb-0016-d784d1f60eb0
md"""
We are going to define out tree as either:
  - a :leaf that predicts a single value xor
  - a :split containing an attribute to split on, subtrees for specific values of that attribute and possibly a default value (if we see an unknown value)
"""

# ╔═╡ dd9487fa-9035-11eb-174b-fdc642e5f518
const T = Any

# ╔═╡ 704c988e-8ab3-11eb-2882-7bdab634fdb4
struct Leaf
	value::T
end

# ╔═╡ 702fd9e2-8ab3-11eb-0ac9-590fd1e3ca02
struct Split
	attr::Symbol
	subtrees::Dict
	defval::T
end

# ╔═╡ 7014500a-8ab3-11eb-13bc-69adb8a72f13
const DT = Union{Leaf, Split}

# ╔═╡ 6ff9bd6c-8ab3-11eb-1346-f1b274ec7283
function classify(dt::DT, input::T)::T
	"""
	Classify given input using given decision tree (dt)
	"""
	typeof(dt) == Leaf && (return dt.value)

	## Otherwise this tree consists of an attr to split on and a
	## dictionary whose keys are values of that attribute and whose
	## values are subtrees to consider next
	sdt_key = getfield(input, dt.attr)

	## no subtree for key => default value
	!haskey(dt.subtrees, Symbol(sdt_key)) && (return dt.defval)

	sdt = dt.subtrees[Symbol(sdt_key)] ## choose appropriate subtree and
	classify(sdt, input)               ## use it to classify the input
end

# ╔═╡ 6fdb83ce-8ab3-11eb-03a2-bf78a195426b
function build_tree_id3(inputs::AVT, split_attrs::Vector{Symbol},
		target::Symbol)::DT
	λ₁ = inp -> getfield(inp, target)
	label_cnt = λ₁.(inputs) |> Counter

	most_common_label = most_common(label_cnt, 1)[1][1]

	## If unique label, predict it
	length(label_cnt) == 1 && (return Leaf(most_common_label))

	## no split attributes left => return the majority label
	length(split_attrs) == 0 && (return Leaf(most_common_label))

	## otherwise split by best attribute
	best_attr = reduce(
		(t_attr, c_attr) -> (p = partition_entropy_by(inputs, c_attr, target);
		t_attr = p < t_attr[2] ? (c_attr, p) : t_attr), split_attrs;
		init=(nothing, Inf)
	)[1]

	parts = partition_by(inputs, best_attr)
	new_attrs = filter(a -> a ≠ best_attr, split_attrs)

	## Recursively build the subtrees
	subtrees = Dict(attr_val => build_tree_id3(subset, new_attrs, target)
		for (attr_val, subset) ∈ parts)

	return Split(best_attr, subtrees, most_common_label)
end

# ╔═╡ f21e8240-8ab9-11eb-1223-cbc0f2b1aa8c
dtree = build_tree_id3(inputs, Symbol[fieldnames(Candidate)[1:end-1]...],
	:did_well)

# ╔═╡ 4891273a-8a8e-11eb-32a5-5b7fb191b833
@test classify(dtree, Candidate(:Junior, :Java, true, false))
# Should predict true

# ╔═╡ 2a823d90-8abd-11eb-0e52-d53dfd7d28af
@test !classify(dtree, Candidate(:Junior, :Java, true, true))
# Should predict false

# ╔═╡ 2a68888e-8abd-11eb-396d-2dbac68b7419
@test classify(dtree, Candidate(:Intern, :Java, true, true))
## Should predict true - :Intern is unknwon to decison tree

# ╔═╡ Cell order:
# ╟─87e5e2ec-8915-11eb-362b-6bf13a36b8e4
# ╠═5eebac68-891b-11eb-37ba-c1d9481c6134
# ╠═8e2aca8e-8a8f-11eb-1e68-1316740f4697
# ╟─cb25501a-8915-11eb-3626-a1ae6d233f92
# ╟─cb0838ae-8915-11eb-082f-3991192a101f
# ╟─caf03ada-8915-11eb-2935-01e110accba6
# ╟─cad66356-8915-11eb-0173-ada200678f63
# ╠═dc3b7280-8915-11eb-258a-1789476897b9
# ╠═dc209b9a-8915-11eb-0f63-ab3b7c33793c
# ╟─d7aade44-891b-11eb-00d3-abe7406c4b09
# ╠═9ae99efe-891c-11eb-3d1f-89182e3f9ff8
# ╠═ea1f4a2c-891d-11eb-3027-bd447dd9d1e1
# ╠═13e3ccb6-8aae-11eb-3a9e-85b5afc74097
# ╠═13c729da-8aae-11eb-29aa-b171020ad8dd
# ╠═13b2940c-8aae-11eb-2dd2-8dec30dd4881
# ╟─dc0c5a5e-8915-11eb-2c62-5328b42add69
# ╟─dbefe77a-8915-11eb-39f1-5d0c84729739
# ╠═dbd86924-8915-11eb-351f-8362f09ba984
# ╠═d32b86e0-8928-11eb-193c-e3d7c85d23dd
# ╟─dbbf22a2-8915-11eb-00eb-4b0278c0283d
# ╟─5f5e6d02-892a-11eb-32d8-dfc1a7f259fa
# ╠═073f6442-892e-11eb-0fc2-bbeb1d12619e
# ╠═2ec41c48-8a95-11eb-13b6-61f4a07a31db
# ╟─25fb8f9a-8952-11eb-35bb-b173ac466835
# ╠═045023be-8aac-11eb-29c1-3f4da066ff9b
# ╠═b1c0a416-8aba-11eb-140a-9f3c4eb93346
# ╠═65ec8554-8953-11eb-3973-b56039754312
# ╠═0c8055b8-89b2-11eb-3046-331119a0dc9b
# ╠═d728d526-8956-11eb-3c22-e788d025e8b4
# ╠═17768334-8a98-11eb-1172-3d5079435dcf
# ╠═8365f5cc-8aaf-11eb-1928-b3a201b840a0
# ╟─432cd594-89b9-11eb-2d4a-6f46f08a511d
# ╟─48db121e-8a8e-11eb-189f-d9f41e4b1d9d
# ╟─48ba469c-8a8e-11eb-0016-d784d1f60eb0
# ╠═dd9487fa-9035-11eb-174b-fdc642e5f518
# ╠═704c988e-8ab3-11eb-2882-7bdab634fdb4
# ╠═702fd9e2-8ab3-11eb-0ac9-590fd1e3ca02
# ╠═7014500a-8ab3-11eb-13bc-69adb8a72f13
# ╠═6ff9bd6c-8ab3-11eb-1346-f1b274ec7283
# ╠═6fdb83ce-8ab3-11eb-03a2-bf78a195426b
# ╠═f21e8240-8ab9-11eb-1223-cbc0f2b1aa8c
# ╠═4891273a-8a8e-11eb-32a5-5b7fb191b833
# ╠═2a823d90-8abd-11eb-0e52-d53dfd7d28af
# ╠═2a68888e-8abd-11eb-396d-2dbac68b7419


function table_lookup(datin,datout,in)
    out=datout[1]
    if in<=datin[1]
        done=1
    else
        done=0
    end
    k=1
    while done==0
        k=k+1
        if in<=datin[k]
            out=datout[k-1]+(datout[k]-datout[k-1])/(datin[k]-datin[k-1])*(in-datin[k-1])
            done=1
        elseif k>=length(datin)
            out=datout[k]
            done=1
        end
    end
    return out
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@stochastic_model vec_simple begin
    @stage 1 begin
        @parameters begin
            lb = [40.0, 20.0]
            c  = [100., 150.]
            A  = [1. 1.]
            b  = [120.]
        end
        @decision(vec_simple, x[i in 1:2] >= lb[i])
        @objective(vec_simple, Min, dot(c, x))
        @constraint(vec_simple, A * x - b in MOI.Nonpositives(1))
    end
    @stage 2 begin
        @parameters begin
            lb = [0., 0.]
            T  = [-60. 0.;
                  0.  -80.]
            W  = [6. 10.;
                  8. 5.]
        end
        @uncertain ξ[1:4]
        q  = [ξ[1], ξ[2]]
        ub = [ξ[3], ξ[4]]
        @recourse(vec_simple, lb[i] <= y[i in 1:2] <= ub[i])
        @objective(vec_simple, Max, dot(q, y))
        @constraint(vec_simple, T * x + W * y in MOI.Nonpositives(2))
    end
end

ξ₁ = @scenario ξ[1:4] = [24.0, 28.0, 500.0, 100.0] probability = 0.4
ξ₂ = @scenario ξ[1:4] = [28.0, 32.0, 300.0, 300.0] probability = 0.6

vec_simple_res = SPResult([46.67, 36.25], Dict(1 => [300., 100.], 2 => [300., 100.]), -855.83, -1518.75, 662.92, 286.92, -1445.92, -568.92)
push!(problems, (vec_simple, [ξ₁,ξ₂], vec_simple_res, "Vectorized Simple"))


immutable PolyUnion
    vars::Vector{Symbol}
    mat::Array{Int, 2}
    eltype::Type
end

function PolyUnion(ps)
    newvars = isempty(ps) ? [] : union([vars(p) for p in ps]...)
    mat = Array(Int, length(newvars), length(ps))
    for pindex = 1:length(ps)
        for i = 1:length(newvars)
            mat[i, pindex] = findfirst(vars(ps[pindex]), newvars[i])
        end
    end
    PolyUnion(newvars, mat, promote_type([eltype(p) for p in ps]...))
end

eltype(pu::PolyUnion) = pu.eltype

vars(pu::PolyUnion) = pu.vars

nvars(pu::PolyUnion) = length(vars(pu))

newexps(pu::PolyUnion, exps::Vector{Int}, polindex::Int) =
    [pu.mat[i, polindex] == 0 ? 0 : exps[pu.mat[i, polindex]] for i = 1:nvars(pu)]

function oldexps(pu::PolyUnion, exps::Vector{Int}, polindex::Integer)
    v = zeros(Int, length(exps))
    l = 0
    for i = 1:length(exps)
        t = pu.mat[i, polindex]
        if t != 0
            v[t] = exps[i]
            l += 1
        end
    end
    v[1:l]
end

zero(pu::PolyUnion) = zero(MPoly{eltype(pu)}, vars=vars(pu))

one(pu::PolyUnion) = one(MPoly{eltype(pu)}, vars=vars(pu))


# test_convergence.jl

# example problem from On the Robustness of a Multiple Semicoarsened Grid Method, Oosterlee and Wesseling.

function get_problem(ϵ,β)
    n = 256
    m = n
    A = anisotropic2d(ϵ, β*π/180, n, m)
    A = V_cycle(A, (n,m))
    b = ones((n-1)*(m-1))
    return (A,b)
end

@testset "Convergence for example from Oosterlee et. al." begin
    for ϵ in [1 1e-2 1e-8]
        for β in [0 15 30 45]
            (A,b) = get_problem(ϵ,β)
            A.grids[1].b .= b # copy rhs
            for item in Base.Iterators.take(A,15) end # iterate
            if !( ( ϵ == 1e-8 ) && ( β == 45 ) ) # known failure for last case in 15 iters
                @test A.resnorm[end] < 1/256^2
            end
            log(A,ϵ,β)
        end
    end
end


# Normalized S₂ in spatial domain
s2(x :: AbstractArray{Bool}) =
    (x |> fft .|> abs2 |> ifft |> real) / length(x)

function read_manifest(root)
    manifest = find_file(root, datafile("manifest.json"))
    return JSON.parse(manifest)
end

function create_hints(slices :: AbstractArray{Bool, 3},
                      ratio  :: Rational)
    n, m = numerator(ratio), denominator(ratio)
    x, y, z = size(slices)

    hints = rand(Bool, (x, y, Int(ratio * z)))

    # = Original has each n-th sample
    if m == 1
        for k in 1:z
            hints[:,:,(k-1)*n+1] = slices[:,:,k]
        end
    # = Original has each n-th sample dropped
    elseif n == m + 1
        for k in 1:z
            hints[:,:,(k-1)÷m + k] = slices[:,:,k]
        end
    else
        throw("Hints are not supported for this ratio")
    end

    # KLUDGE: without this magic line the last slice of a
    # reconstructed image will be a noisy mix of the first and the
    # last slice of the original.
    hints[:,:,end] = slices[:,:,end]

    return hints
end

"""
    reconstruct(archivename; radius = 0.6, use_hints = false)

Reconstruct two-phase porous medium from an archive of
slices. `archivename` is the name of the archive. `radius` is a filter
width used in Phase Reconstruction algorithm and the default value is
usually good for all cases. When `use_hints` is `true` the slices are
used not only for creation of autocorrelation function, but also in
initial approximation of the porous medium. This results in a
reconstruction with more visual resemblance with the original.
"""
function reconstruct(archivename :: AbstractString;
                     radius      :: AbstractFloat = 0.6,
                     use_hints   :: Bool          = false)
    local slices, ratio

    ZipFile.Reader(archivename) do root
        manifest = read_manifest(root)

        ratio  = manifest["numerator"] // manifest["denominator"]
        prefix = manifest["prefix"]
        n      = manifest["amount"]

        slices = mapreduce(stich, 1:n) do k
            pic  = find_file(root, datafile("$(prefix)-$(k).png"))
            # Here read all data and wrap in in another IO stream
            # because ZipFile does not work well with FileIO. Julia
            # cannot into abstractions.
            rawdata = Vector{UInt8}(undef, pic.uncompressedsize)
            read!(pic, rawdata)
            rawdata |> IOBuffer |> Stream{format"PNG"} |> load .|> Bool
        end
    end

    # Calculate normalized S₂ function in spatial
    # domain. Normalization means that the result of F¯¹[|F(f)|²] is
    # divided on number of elements in the array.
    s2cut = s2(slices)

    # Resample correlation function along the third dimension. The
    # number of slices in the third dimension is multiplied by ratio.
    s2restored = AutoCorrelationResampling.ac_resample(s2cut, ratio)

    # Transform to unnormalized representation in the frequency
    # domain. Sometimes small negative values may occur after
    # resampling. Just clamp all values to the range [0, ∞).
    s2ft = max.(s2restored |> rfft |> real, 0) * length(s2restored)

    # Create hints if needed
    hints = use_hints ? create_hints(slices, ratio) : nothing

    return PhaseRec.phaserec(s2ft, size(s2restored);
                             radius = radius,
                             noise  = hints)
end


export insolation, solar_flux_and_cos_sza

"""
    insolation(θ::FT, d::FT, param_set::APS) where {FT <: Real}

Returns the insolation given the zenith angle and earth-sun distance
param_set is an AbstractParameterSet from CLIMAParameters.jl.
"""
function insolation(θ::FT, d::FT, param_set::APS) where {FT <: Real}
    S0::FT = tot_solar_irrad(param_set)
    AU::FT = astro_unit()
    # set max. zenith angle to π/2, insolation should not be negative
    if θ > FT(π)/FT(2)
        θ = FT(π)/FT(2)
    end
    # weighted irradiance (3.12)
    S = S0 * (AU / d)^2
    # TOA insolation (3.15)
    F = S * cos(θ)
    return F
end

"""
    solar_flux_and_cos_sza(date::DateTime,
                      longitude::FT,
                      latitude::FT,
                      param_set::APS) where {FT <: Real}

Returns the top-of-atmosphere (TOA) solar flux, i.e.
the total solar irradiance (TSI) weighted by the earth-sun distance
and cos(solar zenith angle) for input to RRTMGP.jl
param_set is an AbstractParameterSet from CLIMAParameters.jl.
"""
function solar_flux_and_cos_sza(date::DateTime,
                           longitude::FT,
                           latitude::FT,
                           param_set::APS) where {FT <: Real}
    S0::FT = tot_solar_irrad(param_set)
    AU::FT = astro_unit()
    # θ = solar zenith angle, ζ = solar azimuth angle, d = earth-sun distance
    θ, ζ, d = instantaneous_zenith_angle(date, longitude, latitude, param_set)
    # set max. zenith angle to π/2, insolation should not be negative
    if θ > FT(π)/FT(2)
        θ = FT(π)/FT(2)
    end
    μ = cos(θ)
    # TOA solar flux (3.12)
    S = S0 * (AU / d)^2
    # return inputs needed for RRTMGP.jl
    return S, μ
end


#= -------------------------------------------------------------------------
# @file problema2_simulacion_calentamiento_invierno.jl
#
# @date 12/07/17 22:21:52
# @author Martin Noblia
# @email mnoblia@disroot.org
#
# @brief
#
# @detail
#
#  Licence:
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
---------------------------------------------------------------------------=#
using DifferentialEquations
using PyPlot
#=------------------------------------------------------------------------------
                              constants
------------------------------------------------------------------------------=#
const  A      =  2.0
const  R      =  2.390057361376673e-5
const  cₑ     =  4.184e6
const  V      =  0.5
const  θ_ref  =  333.0
const  θ₀     =  281.0
const  θ_a    =  281.0
q = 1841012.0
K_p = 0.0001
G_inv = 0.4 * 836800.0
tspan = (0.0, 1000.0)
θ_a_f(t) = (100 < t > 300 ) ? θ_a + 100: θ_a
# calculo del q_comp
q_nec = ((θ_ref - θ_a) / (R)) - A * G_inv
# sin compensar en invierno
θ_dot(t, θ) = ((G_inv * A) / (cₑ * V)) + ((θ_a / (R * cₑ * V))) - ((θ / (R * cₑ * V)))
# compensado en invierno
θ_dot_comp(t, θ) = ((G_inv * A) / (cₑ * V)) + ((θ_a_f(t) / (R * cₑ * V))) - ((θ / (R * cₑ * V))) + (q_nec / (cₑ * V)) + (sol(t)-θ_ref) * K_p
θ_dot(t, θ) = ((q / (cₑ * V)) + (G_inv * A) / (cₑ * V)) + ((θ_a / (R * cₑ * V))) - ((θ / (R * cₑ * V)))
prob = ODEProblem(θ_dot, θ₀, tspan)
prob_comp = ODEProblem(θ_dot_comp, θ₀, tspan)
sol = solve(prob, Tsit5(), reltol=1e-12,abstol=1e-12)
sol_comp = solve(prob_comp, Tsit5(), reltol=1e-12,abstol=1e-12)
#=------------------------------------------------------------------------------
                           plotting
------------------------------------------------------------------------------=#
plot(sol.t, sol[1, :] - 273.0, label=L"\theta_{inv}(t)", linewidth=3)
plot(sol_comp.t, sol_comp[1, :] - 273.0, label=L"\theta_{inv-comp}(t)", linewidth=3)
xlabel(L"t\,[seg]")
ylabel(L"T\,[^{\circ}C]")
grid("on")
matplotlib["rcParams"][:update](["font.size" => 14, "font.family" => "serif"])
legend()


function medianSD(x)

    x = vec(x)

    if mod.(length(x), 2) == 0

        #do find mid location
        midLocation = (length(x) + 1) / 2
        # evaluate mean of the upper and lower index
        meanSD([sort(x)[Int64(floor(midLocation))], sort(x)[Int64(ceil(midLocation))]])

    elseif mod.(length(x), 2) != 0

        midLocation = (length(x) + 1) / 2
        sort(x)[Int64(midLocation)]

    else

        println("No median")

    end

end