export
    LineSegment,
    parallel

struct LineSegment
    A::VecE2
    B::VecE2
end
Base.:+(seg::LineSegment, V::VecE2) = LineSegment(seg.A + V, seg.B + V)
Base.:-(seg::LineSegment, V::VecE2) = LineSegment(seg.A - V, seg.B - V)
Base.convert(::Type{Line}, seg::LineSegment) = Line(seg.A, seg.B)
Base.convert(::Type{LineSegment}, line::Line) = LineSegment(line.A, line.B)

get_polar_angle(seg::LineSegment) = mod2pi(atan(seg.B.y - seg.A.y, seg.B.x - seg.A.x))

"""
The distance between the line segment and the point P
"""
function get_distance(seg::LineSegment, P::VecE2)

    ab = seg.B - seg.A
    pb = P - seg.A

    denom = normsquared(ab)
    if denom == 0.0
        return 0.0
    end

    r = (ab⋅pb)/denom

    if r ≤ 0.0
        LinearAlgebra.norm(P - seg.A)
    elseif r ≥ 1.0
        LinearAlgebra.norm(P - seg.B)
    else
        LinearAlgebra.norm(P - (seg.A + r*ab))
    end
end

"""
What side of the line you are on, based on A → B
"""
get_side(seg::LineSegment, p::VecE2) = sign((seg.B.x-seg.A.x) * (p.y-seg.A.y) - (seg.B.y-seg.A.y) * (p.x-seg.A.x))

"""
The angular distance between the two line segments
"""
function angledist(segA::LineSegment, segB::LineSegment)
    u = segA.B - segA.A
    v = segB.B - segB.A
    sqdenom = (u⋅u)*(v⋅v)
    if isapprox(sqdenom, 0.0, atol=1e-10)
        return NaN
    end
    return acos((u⋅v) / sqrt(sqdenom))
end

"""
True if the two segments are parallel
"""
function parallel(segA::LineSegment, segB::LineSegment, ε::Float64=1e-10)
    θ = angledist(segA, segB)
    return isapprox(θ, 0.0, atol=ε)
end

"""
Given P colinear with seg, the function checks if
point P lies on the line segment.
"""
function on_segment(P::VecE2, seg::LineSegment)
    return P.x ≤ max(seg.A.x, seg.B.x) && P.x ≥ min(seg.A.x, seg.B.x) &&
           P.y ≤ max(seg.A.y, seg.B.y) && P.y ≥ min(seg.A.y, seg.B.y)
end
on_segment(A::VecE2, P::VecE2, B::VecE2) = on_segment(P, LineSegment(A,B))

"""
returns true if line segments segP and segQ intersect
"""
function intersects(segP::LineSegment, segQ::LineSegment)

    p1 = segP.A
    q1 = segP.B
    p2 = segQ.A
    q2 = segQ.B

    # Find the four orientations needed for general and
    # special cases
    o1 = orientation(p1, q1, p2)
    o2 = orientation(p1, q1, q2)
    o3 = orientation(p2, q2, p1)
    o4 = orientation(p2, q2, q1)

    # Special Cases
    if (o1 != o2 && o3 != o4) || # General case
       (o1 == 0 && on_segment(p1, p2, q1)) || # p1, q1 and p2 are colinear and p2 lies on segment p1q1
       (o2 == 0 && on_segment(p1, q2, q1)) || # p1, q1 and p2 are colinear and q2 lies on segment p1q1
       (o3 == 0 && on_segment(p2, p1, q2)) || # p2, q2 and p1 are colinear and p1 lies on segment p2q2
       (o4 == 0 && on_segment(p2, q1, q2))    # p2, q2 and q1 are colinear and q1 lies on segment p2q2

        return true
    end
    return false # Doesn't fall in any of the above cases
end


__precompile__()

module LorentzVectors

import LinearAlgebra: dot, ⋅, cross, ×, norm, normalize
import Random: rand, AbstractRNG

import Base: +, -, *, /, ==, isapprox, ≈, zero

export LorentzVector, SpatialVector, Vec4, Vec3
export dot, ⋅, cross, ×, norm, normalize
export boost, rotate

"""
    LorentzVector(t, x, y, z)

Lorentz 4-vector, as used in Special Relativity.

The metric convention is g = diag(+1,-1,-1,-1). No distinction is made between
co- and contravariant vectors.
"""
struct LorentzVector{T <: AbstractFloat}
    t :: T
    x :: T
    y :: T
    z :: T
end

"""
    SpatialVector(x, y, z)

Spatial part of a Lorentz 4-vector.
"""
struct SpatialVector{T <: AbstractFloat}
    x :: T
    y :: T
    z :: T
end

"""
    LorentzVector(t, x, y, z)

Promoting constructors for LorentzVector{T}.
"""
LorentzVector(t, x, y, z) = LorentzVector(promote(t, x, y, z)...)
LorentzVector(t::T, x::T, y::T, z::T) where {T <: Union{Integer, Rational, Irrational}} =
    LorentzVector(float(t), x, y, z)

"""
    SpatialVector(x, y, z)

Promoting constructors for SpatialVector{T}.
"""
SpatialVector(x, y, z) = SpatialVector(promote(x, y, z)...)
SpatialVector(x::T, y::T, z::T) where {T <: Union{Integer, Rational, Irrational}} =
    SpatialVector(float(x), y, z)

"""
    SpatialVector(u)

Construct a 3-vector from the spatial part of a 4-vector.
"""
SpatialVector(u::LorentzVector) = SpatialVector(u.x, u.y, u.z)

"""
    LorentzVector(t, u)

Construct a 4-vector from a time component and a 3-vector.
"""
LorentzVector(t::T, u::SpatialVector{U}) where {T,U} =
    LorentzVector(t, u.x, u.y, u.z)

"""
    zero(LorentzVector{T})
    zero(LorentzVector)

Constructs a zero four-vector.
"""
zero(::Type{LorentzVector{T}}) where {T} = LorentzVector{T}(zero(T), zero(T), zero(T), zero(T))
zero(::Type{LorentzVector}) = zero(LorentzVector{Float64})

"""
    zero(SpatialVector{T})
    zero(SpatialVector)

Constructs a zero three-vector.
"""
zero(::Type{SpatialVector{T}}) where {T} = SpatialVector{T}(zero(T), zero(T), zero(T))
zero(::Type{SpatialVector}) = zero(SpatialVector{Float64})

"Alias of LorentzVector"
const Vec4 = LorentzVector

"Alias of SpatialVector"
const Vec3 = SpatialVector

function +(u::LorentzVector)
    u
end

function +(u::LorentzVector, v::LorentzVector)
    @fastmath LorentzVector(u.t + v.t, u.x + v.x, u.y + v.y, u.z + v.z)
end

function -(u::LorentzVector)
    @fastmath LorentzVector(-u.t, -u.x, -u.y, -u.z)
end

function -(u::LorentzVector, v::LorentzVector)
    @fastmath u + (-v)
end

function *(λ::Number, u::LorentzVector)
    @fastmath LorentzVector(λ*u.t, λ*u.x, λ*u.y, λ*u.z)
end

function *(u::LorentzVector, λ::Number)
    @fastmath λ * u
end

function /(u::LorentzVector, λ::Number)
    @fastmath u * (one(λ) / λ)
end

function +(u::SpatialVector)
    u
end

function +(u::SpatialVector, v::SpatialVector)
    @fastmath SpatialVector(u.x + v.x, u.y + v.y, u.z + v.z)
end

function -(u::SpatialVector)
    @fastmath SpatialVector(-u.x, -u.y, -u.z)
end

function -(u::SpatialVector, v::SpatialVector)
    @fastmath u + (-v)
end

function *(λ::Number, u::SpatialVector)
    @fastmath SpatialVector(λ*u.x, λ*u.y, λ*u.z)
end

function *(u::SpatialVector, λ::Number)
    @fastmath λ * u
end

function /(u::SpatialVector, λ::Number)
    @fastmath u * (one(λ) / λ)
end

function ==(u::LorentzVector, v::LorentzVector)
    u.t == v.t && u.x == v.x && u.y == v.y && u.z == v.z
end

function ==(u::SpatialVector, v::SpatialVector)
    u.x == v.x && u.y == v.y && u.z == v.z
end

function isapprox(u::LorentzVector, v::LorentzVector;
                  atol::Real=0,
                  rtol::Real=atol>0 ? 0 : √min(eps(typeof(u.x)), eps(typeof(v.x))))
    err = max(abs(u.t - v.t), norm(Vec3(u)-Vec3(v)))
    err <= max(atol, rtol*max(abs(u.t), abs(v.t), norm(Vec3(u)), norm(Vec3(v))))
end

function isapprox(u::SpatialVector, v::SpatialVector;
                  atol::Real=0,
                  rtol::Real=atol>0 ? 0 : √min(eps(typeof(u.x)), eps(typeof(v.x))))
    err = norm(u-v)
    err <= max(atol, rtol*max(norm(u), norm(v)))
end

"""
    dot(u, v)
    u⋅v

Inner product of 4-vectors, in the Minkowsky metric (+,-,-,-).
"""
function dot(u::LorentzVector, v::LorentzVector)
    @fastmath u.t*v.t - u.x*v.x - u.y*v.y - u.z*v.z
end

function dot(u::SpatialVector, v::SpatialVector)
    @fastmath u.x*v.x + u.y*v.y + u.z*v.z
end

"""
    cross(u, v)
    u×v

Cross product of two spatial vectors.
"""
function cross(u::SpatialVector, v::SpatialVector)
    @fastmath SpatialVector(
        u.y*v.z - u.z*v.y,
        u.z*v.x - u.x*v.z,
        u.x*v.y - u.y*v.x)
end

"""
    norm(x)
"""
function norm(v::SpatialVector)
    @fastmath sqrt(v⋅v)
end

"""
    normalize(x)
"""
function normalize(v::SpatialVector)
    v / norm(v)
end

"""
    rand(rng, SpatialVector)
"""
function rand(r::AbstractRNG, ::Type{V}) where {U <: Real, V <: SpatialVector{U}}
    cθ = 2*rand(r, U) - 1
    sθ = sqrt(1-cθ^2)
    φ = 2π * rand(r, U)
    SpatialVector{U}(sθ * cos(φ), sθ * sin(φ), cθ)
end

function boost(u::LorentzVector, β::LorentzVector)
    boost(u, SpatialVector(β))
end

function boost(u::LorentzVector, β::SpatialVector)
    γ = one(β.x) / sqrt(one(β.x) - β⋅β)
    if γ == one(γ)
        return u
    end
    x_old = Vec3(u)
    t_new = γ * (u.t - β⋅x_old)
    x_new = x_old + ((γ-one(γ)) * (x_old⋅β) / (β⋅β) - γ*u.t) * β
    LorentzVector(t_new, x_new.x, x_new.y, x_new.z)
end

"""
Rotate a vector v by an angle θ around the unit vector k̂, using Rodrigues' rotation formula.

Note: k̂ *must* be normalized, otherwise this method will return the wrong result.

Source: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
"""
function rotate(v::SpatialVector, k̂::SpatialVector, θ::Real)
    v*cos(θ) + (k̂×v)*sin(θ) + k̂*(k̂⋅v)*(1-cos(θ))
end

end # module LorentzVectors


include("../utils.jl")
include("group_utils.jl")

@testset "Product group" begin
    SOn = SpecialOrthogonal(3)
    Tn = TranslationGroup(2)
    Rn = Rotations(3)
    M = ProductManifold(SOn, Tn)
    G = ProductGroup(M)
    @test_throws ErrorException ProductGroup(ProductManifold(Rotations(3), Stiefel(3,2)))
    @test G isa ProductGroup
    @test submanifold(G, 1) === SOn
    @test submanifold(G, 2) === Tn
    @test base_manifold(G) === M
    @test sprint(show, G) == "ProductGroup($(SOn), $(Tn))"
    @test sprint(show, "text/plain", G) == "ProductGroup with 2 subgroups:\n $(SOn)\n $(Tn)"
    x = Matrix{Float64}(I, 3, 3)
    for f in [group_exp!, group_log!]
        @test Manifolds.decorator_transparent_dispatch(f,G,x,x) === Val{:transparent}()
    end
    t = Vector{Float64}.([1:2, 2:3, 3:4])
    ω = [[1.0, 2.0, 3.0], [3.0, 2.0, 1.0], [1.0, 3.0, 2.0]]
    tuple_pts = [(exp(Rn, x, hat(Rn, x, ωi)), ti) for (ωi, ti) in zip(ω, t)]
    tuple_v = (hat(Rn, x, [1.0, 0.5, -0.5]), [-1.0, 2.0])
    eA = [x, zeros(2)]
    shape_se = Manifolds.ShapeSpecification(Manifolds.ArrayReshaper(), M.manifolds...)
    e = Manifolds.prod_point(shape_se, eA...)

    @testset "identity specializations" begin
        @test inv(G, Identity(G,e)) === Identity(G,e)
        @test identity(G, Identity(G,e)) === Identity(G,e)
        @test submanifold_component(G, Identity(G,e), Val(1)) == Identity(SOn,x)
        @test submanifold_component(G, Identity(G,e), Val(2)) == Identity(Tn, zeros(2))
        @test submanifold_components(G, Identity(G,e)) == (Identity(SOn,x), Identity(Tn,zeros(2)))
        @test compose(G, Identity(G,e), Identity(G,e)) === Identity(G,e)
    end

    @testset "product point" begin
        reshapers = (Manifolds.ArrayReshaper(), Manifolds.StaticReshaper())
        for reshaper in reshapers
            shape_se = Manifolds.ShapeSpecification(reshaper, M.manifolds...)
            pts = [Manifolds.prod_point(shape_se, tp...) for tp in tuple_pts]
            v_pts = [Manifolds.prod_point(shape_se, tuple_v...)]
            @test compose(G, pts[1], Identity(G,e)) == pts[1]
            @test compose(G, Identity(G,e), pts[1]) == pts[1]
            test_group(G, pts, v_pts, v_pts; test_diff = true)
            @test isapprox(
                M,
                identity(M, pts[1]),
                group_exp(M, v_pts[1]),
                Manifolds.prod_point(
                    shape_se,
                    group_exp(SOn, v_pts[1].parts[1]),
                    group_exp(Tn, v_pts[1].parts[2]),
                ),
            )
            @test isapprox(
                M,
                identity(M, pts[1]),
                group_log(M, pts[1]),
                Manifolds.prod_point(
                    shape_se,
                    group_log(SOn, pts[1].parts[1]),
                    group_log(Tn, pts[1].parts[2]),
                ),
            )
        end
    end

    @testset "product repr" begin
        pts = [ProductRepr(tp...) for tp in tuple_pts]
        v_pts = [ProductRepr(tuple_v...)]
        @test compose(G, pts[1], Identity(G,e)) == pts[1]
        @test compose(G, Identity(G,e), pts[1]) == pts[1]
        test_group(G, pts, v_pts, v_pts; test_diff = true, test_mutating = false)
        @test isapprox(
            M,
            group_exp(M, v_pts[1]),
            ProductRepr(
                group_exp(SOn, v_pts[1].parts[1]),
                group_exp(Tn, v_pts[1].parts[2]),
            ),
        )
        @test isapprox(
            M,
            group_log(M, pts[1]),
            ProductRepr(group_log(SOn, pts[1].parts[1]), group_log(Tn, pts[1].parts[2])),
        )
    end
    @test sprint(show, "text/plain", G) === """
    ProductGroup with 2 subgroups:
     SpecialOrthogonal(3)
     TranslationGroup(2; field = ℝ)"""
end


# The old-gdemo model.
@model gdemo(x, y) = begin
  s ~ InverseGamma(2, 3)
  m ~ Normal(0, sqrt(s))
  x ~ Normal(m, sqrt(s))
  y ~ Normal(m, sqrt(s))
  return s, m
end

@model gdemo_d() = begin
  s ~ InverseGamma(2, 3)
  m ~ Normal(0, sqrt(s))
  1.5 ~ Normal(m, sqrt(s))
  2.0 ~ Normal(m, sqrt(s))
  return s, m
end

gdemo_default = gdemo_d()

@model MoGtest(D) = begin
    mu1 ~ Normal(1, 1)
    mu2 ~ Normal(4, 1)
    z1 ~ Categorical(2)
    if z1 == 1
        D[1] ~ Normal(mu1, 1)
    else
        D[1] ~ Normal(mu2, 1)
    end
    z2 ~ Categorical(2)
    if z2 == 1
        D[2] ~ Normal(mu1, 1)
    else
        D[2] ~ Normal(mu2, 1)
    end
    z3 ~ Categorical(2)
    if z3 == 1
        D[3] ~ Normal(mu1, 1)
    else
        D[3] ~ Normal(mu2, 1)
    end
    z4 ~ Categorical(2)
    if z4 == 1
        D[4] ~ Normal(mu1, 1)
    else
        D[4] ~ Normal(mu2, 1)
    end
    z1, z2, z3, z4, mu1, mu2
end

MoGtest_default = MoGtest([1.0 1.0 4.0 4.0])

# Declare empty model to make the Sampler constructor work.
@model empty_model() = begin x = 1; end

# A collection of models for which the mean-of-means for the posterior should
# be same.
@model function gdemo1(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
    # `dot_assume` and `observe`
    m = TV(undef, length(x))
    m .~ Normal()
    x ~ MvNormal(m, 0.5 * ones(length(x)))
end

@model function gdemo2(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
    # `assume` with indexing and `observe`
    m = TV(undef, length(x))
    for i in eachindex(m)
        m[i] ~ Normal()
    end
    x ~ MvNormal(m, 0.5 * ones(length(x)))
end

@model function gdemo3(x = 10 * ones(2))
    # Multivariate `assume` and `observe`
    m ~ MvNormal(length(x), 1.0)
    x ~ MvNormal(m, 0.5 * ones(length(x)))
end

@model function gdemo4(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
    # `dot_assume` and `observe` with indexing
    m = TV(undef, length(x))
    m .~ Normal()
    for i in eachindex(x)
        x[i] ~ Normal(m[i], 0.5)
    end
end

# Using vector of `length` 1 here so the posterior of `m` is the same
# as the others.
@model function gdemo5(x = 10 * ones(1))
    # `assume` and `dot_observe`
    m ~ Normal()
    x .~ Normal(m, 0.5)
end

@model function gdemo6()
    # `assume` and literal `observe`
    m ~ MvNormal(2, 1.0)
    [10.0, 10.0] ~ MvNormal(m, 0.5 * ones(2))
end

@model function gdemo7(::Type{TV} = Vector{Float64}) where {TV}
    # `dot_assume` and literal `observe` with indexing
    m = TV(undef, 2)
    m .~ Normal()
    for i in eachindex(m)
        10.0 ~ Normal(m[i], 0.5)
    end
end

@model function gdemo8()
    # `assume` and literal `dot_observe`
    m ~ Normal()
    [10.0, ] .~ Normal(m, 0.5)
end

@model function _prior_dot_assume(::Type{TV} = Vector{Float64}) where {TV}
    m = TV(undef, 2)
    m .~ Normal()

    return m
end

@model function gdemo9()
    # Submodel prior
    m = @submodel _prior_dot_assume()
    for i in eachindex(m)
        10.0 ~ Normal(m[i], 0.5)
    end
end

@model function _likelihood_dot_observe(m, x)
    x ~ MvNormal(m, 0.5 * ones(length(m)))
end

@model function gdemo10(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
    m = TV(undef, length(x))
    m .~ Normal()

    # Submodel likelihood
    @submodel _likelihood_dot_observe(m, x)
end

const gdemo_models = (gdemo1(), gdemo2(), gdemo3(), gdemo4(), gdemo5(), gdemo6(), gdemo7(), gdemo8(), gdemo9(), gdemo10())


"""
    Household_StochasticDP

Stochastic DP for a standard household life-cycle model with one shock on consumption.
developed based on `Household.jl` and `src.jl`
"""
module Household_StochasticDP
    import src # math dependency
    import Household # for basic methods

# --------------------- CRRA UTILITY
function u_CRRA(c::Real ; γ::Real = 1.5, ψ::Real = 1E-6)
    # assertion
    @assert( 0 <= c < Inf, "c must be non-negative and finite" )
    @assert( 0 <= γ < Inf, "γ must be non-negative and finite" )
    @assert( ψ >= 0, "ψ must be non-negative" )
    # branch
    γ == 1 ? (return log(c + ψ)::Real) : nothing
    return ( ((c+ψ)^(1-γ) - 1) / (1-γ) )::Real
end # u_CRRA

end # Household_StochasticDP


export GaussianNB

# X: Feature vector as a p x n matrix
# c: Vector of classes
# k: Number of distinct classes
import Distributions
import ScikitLearnBase

type GaussianNB
    # The model hyperparameters (not learned from data)
    bias::Float64

    # The parameters learned from data
    mu::Matrix{Float64} # Matrix of cluster centers: p x k
    sigma::Matrix{Float64} # Diagonals of covariance matrices for each class
	p::Vector{Float64} # Vector of probabilities for each class
	counts::Matrix{Int}

    # A constructor that accepts the hyperparameters as keyword arguments
    # with sensible defaults
    GaussianNB(; bias=0.0f0) = new(bias)
end

# This will define `clone`, `set_params!` and `get_params` for the model
ScikitLearnBase.@declare_hyperparameters(GaussianNB, [:bias])

# GaussianNB is a classifier
ScikitLearnBase.is_classifier(::GaussianNB) = true   # not required for transformers

function Distributions.rand(d::GaussianNB)
	p = size(d.mu, 1)
	c = Distributions.draw(d.drawtable)
	x = Array(Float64, p)
	for dim in 1:p
		x[dim] = rand(Normal(d.mu[dim, c], d.sigma[dim, c]))
	end
	x, c
end

function Distributions.rand(d::GaussianNB, n::Integer)
	p = size(d.mu, 1)
	X = Array(Float64, p, n)
	c = Array(Int, n)
	for obs in 1:n
		c[obs] = Distributions.draw(d.drawtable)
		for dim in 1:p
			X[dim, obs] = rand(Normal(d.mu[dim, c[obs]],
				                      d.sigma[dim, c[obs]]))
		end
	end
	X, c
end

function Distributions.mean(d::GaussianNB)
	p, c = size(d.mu)
	mx = zeros(Float64, p)
	for cl in 1:c
		mx += d.p[cl] * d.mu[:, c]
	end
	mx /= c
	mc = 0.0
	for cl in 1:c
		mc += d.p[cl] * cl
	end
	mx, mc
end

function ScikitLearnBase.fit!(model::GaussianNB, X::Matrix, c::Vector)
	p, n = size(X)
	nclasses = maximum(c)

	for i in 1:n
		model.mu[:, c[i]] += X[:, i]
		model.counts[c[i]] += 1
	end
	for cl in 1:nclasses
		model.mu[:, cl] /= model.counts[cl]
	end
	for i in 1:n
		model.sigma[:, c[i]] += (X[:, i] - model.mu[:, c[i]]).^2
	end
	for cl in 1:nclasses
		model.sigma[:, cl] = sqrt(model.sigma[:, cl] / (model.counts[cl] - 1)) + 1e-8
	end
	return model
end

function Distributions.logpdf(d::GaussianNB, x::Vector, c::Real)
	p = length(x)
	res = log(d.p[c])
	for dim in 1:p
		res += logpdf(Normal(d.mu[dim, c], d.sigma[dim, c]), x[dim])
	end
	return res
end

function Distributions.logpdf(d::GaussianNB, X::Matrix, c::Vector)
	p, n = size(X)
	res = zeros(Float64, n)
	for obs in 1:n
		res[obs] = log(d.p[c[obs]])
		for dim in 1:p
			res[obs] += logpdf(Normal(d.mu[dim, c[obs]],
				                      d.sigma[dim, c[obs]]),
			                   X[dim, obs])
		end
	end
	return res
end

function Distributions.loglikelihood(d::GaussianNB, X::Matrix, c::Vector)
	p, n = size(X)
	res = 0.0
	for obs in 1:n
		res += logpdf(d, X[:, obs], c[obs])
	end
	return res
end

function ScikitLearnBase.predict(d::GaussianNB, X::Matrix)
	nclasses = length(d.p)
	p, n = size(X)
	res = Array(Int, n)
	ll = Array(Float64, nclasses)
	for obs in 1:n
		for cl in 1:nclasses
			ll[cl] = logpdf(d, X[:, obs], cl)
		end
		res[obs] = indmax(ll)
	end
	return res
end


Base.@kwdef mutable struct Benchmark
    b::Vector{Float64} # benchmark weights
    c::Vector{Float64} # Climate costs per $ of revenue
    σ::Vector{Float64} # Volatility
    ρ::Matrix # Correlation matrix
    Σ::Union{Matrix,Nothing} = nothing # Covariance Matrix
end

ρ = [
    1 0.8 0.7 0.6 0.7 0.5 0.7 0.6
    0.8 1 0.75 0.65 0.5 0.6 0.5 0.65
    0.7 0.75 1 0.8 0.7 0.7 0.7 0.7
    0.6 0.65 0.8 1 0.85 0.8 0.75 0.75
    0.7 0.5 0.7 0.85 1 0.6 0.8 0.65
    0.5 0.6 0.7 0.8 0.6 1 0.5 0.7
    0.7 0.5 0.7 0.75 0.8 0.5 1 0.8
    0.6 0.65 0.7 0.75 0.65 0.7 0.8 1
]

using LinearAlgebra
function get_cov_from_corr!(bench::Benchmark)::Benchmark
    D = diagm(bench.σ) # diagonal matrix of volatiltiy
    bench.Σ = D * bench.ρ * D # multiplying columns and rows of correlation matrix by volatiltiy
    return bench
end

our_benchmark = Benchmark(b = [0.23, 0.19, 0.17, 0.13, 0.09, 0.08, 0.06, 0.05],
                        c= [-1.2, 0.8, 2.75, 1.60, -2.75, -1.3, 0.90, -1.70],
                        σ = [0.22, 0.20, 0.25, 0.18, 0.35, 0.23, 0.13, 0.29],
                        ρ = ρ)

get_cov_from_corr!(our_benchmark)

function get_benchmark_score(bench::Benchmark)::Float64
    return round(bench.b' * bench.c; digits= 2)
end

get_benchmark_score(our_benchmark)

Base.@kwdef mutable struct TiltedPortfolio
    x::Vector{Float64} # Optimal weights
end

# portfolio = TiltedPortfolio(x = [0.1, 0.1, 0.3, 0.3, 0.2])

function get_excess_score(portfolio::TiltedPortfolio, bench::Benchmark)::Float64
    return round((portfolio.x - bench.b)' * bench.c;digits = 5)
end

function tracking_error_volatility(portfolio::TiltedPortfolio, bench::Benchmark)::Float64
    return sqrt((portfolio.x - bench.b)' * bench.Σ * (portfolio.x - bench.b))
end

# get_excess_score(portfolio, our_benchmark)

# tracking_error_volatility(portfolio, our_benchmark)

using JuMP, COSMO # for the optimization

# tilt the benchmark according to gamma
function portfolio_tilting(bench::Benchmark, γ::Float64)::TiltedPortfolio
    n = length(bench.c) # number of assets
    model = JuMP.Model(COSMO.Optimizer)
    # the optimal weights we want to find
    @variable(model, x[1:n])
    @objective(model, Min, 1/2 * x' * bench.Σ * x - x' * (γ * bench.c + bench.Σ * bench.b))
    @constraint(model, zeros(n) .<= x .<= ones(n))
    @constraint(model, ones(n)' * x == 1)
    JuMP.optimize!(model)
    portfolio = TiltedPortfolio(x = JuMP.value.(x))
    return portfolio
end

# simulate for gammas between 0 and 10
using Plots
gammas = [i for i in 0.0:0.001:0.012]
excess_scores = zeros(length(gammas))
excess_te = zeros(length(gammas))

for i in 1:length(gammas)
    new_portfolio = portfolio_tilting(our_benchmark, gammas[i])
    excess_scores[i] = get_excess_score(new_portfolio, our_benchmark)
    excess_te[i] = tracking_error_volatility(new_portfolio, our_benchmark)
end

plot(excess_te, excess_scores, xlabel = "Tracking error volatility",
                                ylabel = "Excess ESG score",
                                label = "")


@testset "Simplex" begin
    s = Simplex(1, 2, 3)
    @test s == Simplex(1, 2, 3)
    @test sort!(values(s)) == [1,2,3]

    s = Simplex([1,2,3])
    @test dim(s) == 2
    @test sort!(values(s)) == [1,2,3]
    allfaces = [Simplex(2, 3), Simplex(1, 3), Simplex(1, 2)]
    @testset "Simplex faces" for o in faces(s)
        idx = findfirst(isequal(o), allfaces)
        @test values(o) == values(allfaces[idx])
    end

    ss = Set{Simplex}()
    push!(ss, s)
    push!(ss, s)
    @test length(ss) == 1

    s = Simplex([1,2], [1,3])
    @test dim(s) == 1

    X = hcat(zeros(3), Matrix(I,3,3))
    S = Simplex(1,2,3,4)
    @test volume(S,X) == 1/6

    @test faces(s)[1] ∪ faces(s)[2] == s
end

@testset "Cube" begin
    c = Cube{Vector{Int}}([0,0,0],Int[]) # 0-cube
    @test dim(c) == 0
    c = Cube{Vector{Int}}([0,0,0],Int[0,0,0]) # 0-cube
    @test dim(c) == 0
    c = Cube{Vector{Int}}([0,0,0],[1]) # 1-cube
    @test dim(c) == 1
    c = Cube{Vector{Int}}([0,0,0],[1,0,1]) # 2-cube
    @test dim(c) == 2
    c = Cube{Vector{Int}}([0,0,0],[1,1,1]) # 3-cube
    @test dim(c) == 3
    @test hash(c) == hash([0,0,0,1,1,1])
    @test c == Cube{Vector{Int}}([0,0,0],[1,1,1])

    c = Cube{Vector{Float64}}([0.,0.,0.],[.5,.0,.5]) # 2-cube
    @test volume(c) == 0.5^2
    c = Cube{Vector{Float64}}([0.,0.,0.],[.5,.5,.5]) # 2-cube
    @test volume(c) == 0.5^3

end

@testset "CW" begin
    a = Cell()
    @test hash(a) == hash(Cell[])
    @test a == Cell()
    @test dim(a) == 0

    b = Cell()
    c = Cell(1, a, b)
    @test dim(c) == 1
    d = Cell(1, a, b, a, b)
    e = c ∪ d
    @test dim(e) == 2

    @testset for (f1, f2) in zip(faces(c), [a, b])
        @test f1 == f2
    end

    @testset for (f1, f2) in zip(faces(e), [c, d])
        @test f1 == f2
    end
end


# TODO: add more target distributions and make them iteratable

# Dimension of testing distribution
const D = 5
# Tolerance ratio
const TRATIO = Int == Int64 ? 1 : 2
# Deterministic tolerance
const DETATOL = 1e-3 * D * TRATIO
# Random tolerance
const RNDATOL = 5e-2 * D * TRATIO

using Distributions: logpdf, MvNormal, InverseGamma, Normal
using DiffResults: GradientResult, JacobianResult, value, gradient, jacobian
using ForwardDiff: gradient!, jacobian!
using Bijectors: link, invlink, logpdf_with_trans

const dm = zeros(D)
const dσ = ones(D)
const MVN = MvNormal(dm, dσ)
ℓπ(θ) = logpdf(MVN, θ)

# Manual implementation of gradient so that it support matrix mode
function ∂ℓπ∂θ(θ::AbstractVecOrMat)
    d = MVN
    diff = θ .- dm
    dims = θ isa AbstractVector ? (:) : 1
    v = -(D * log(2π) + 2 * sum(log.(dσ)) .+ sum(abs2, diff ./ dσ; dims=dims)) / 2
    v = θ isa AbstractMatrix ? vec(v) : v
    g = -diff
    return (v, g)
end

function ∂ℓπ∂θ_ad(θ::AbstractVector)
    res = GradientResult(θ)
    gradient!(res, ℓπ, θ)
    return (value(res), gradient(res))
end

function ∂ℓπ∂θ_ad(θ::AbstractMatrix)
    v = similar(θ, size(θ, 2))
    g = similar(θ)
    for i in 1:size(θ, 2)
        res = GradientResult(θ[:,i])
        gradient!(res, ℓπ, θ[:,i])
        v[i] = value(res)
        g[:,i] = gradient(res)
    end
    return (v, g)
end

function ∂ℓπ∂θ_viajacob(θ::AbstractMatrix)
    jacob = similar(θ)
    res = JacobianResult(similar(θ, size(θ, 2)), jacob)
    jacobian!(res, ℓπ, θ)
    jacob_full = jacobian(res)
    d, n = size(jacob)
    for i in 1:n
        jacob[:,i] = jacob_full[i,1+(i-1)*d:i*d]
    end
    return (value(res), jacob)
end

# For the Turing model
# @model gdemo() = begin
#     s ~ InverseGamma(2, 3)
#     m ~ Normal(0, sqrt(s))
#     1.5 ~ Normal(m, sqrt(s))
#     2.0 ~ Normal(m, sqrt(s))
#     return s, m
# end

function invlink_gdemo(θ)
    s = invlink(InverseGamma(2, 3), θ[1])
    m = θ[2]
    return [s, m]
end

function ℓπ_gdemo(θ)
    s, m = invlink_gdemo(θ)
    logprior = logpdf_with_trans(InverseGamma(2, 3), s, true) + logpdf(Normal(0, sqrt(s)), m)
    loglikelihood = logpdf(Normal(m, sqrt(s)), 1.5) + logpdf(Normal(m, sqrt(s)), 2.0)
    return logprior + loglikelihood
end

function ∂ℓπ∂θ_gdemo(θ)
    res = GradientResult(θ)
    gradient!(res, ℓπ_gdemo, θ)
    return (value(res), gradient(res))
end


function digit_conversion(digit::Integer, one::Char, five::Char, ten::Char)
    if digit == 9
        return one*ten
    end
    if 9 > digit >= 5
        digit = digit % 5
        return five*one^digit
    end
    if digit == 4
        return one*five
    end
    if digit < 4
        return one^digit
    end
end

function digit_conversion(digit::Integer, one::Char)
    if digit >= 4
        error("Input number exceeds maximum.")
    end

    return one^digit
end

function to_roman(number)
    if number <= 0
        error("Input number must be larger than 0.")
    end

    output = ""

    thousand = div(number, 1000)
    output *= digit_conversion(thousand, 'M')

    hundred = div(number % 1000, 100)
    output *= digit_conversion(hundred, 'C', 'D', 'M')

    ten = div(number % 100, 10)
    output *= digit_conversion(ten, 'X', 'L', 'C')

    one = number % 10
    output *= digit_conversion(one, 'I', 'V', 'X')

end

to_roman(1024)


# workspace()
using Indicators
using PyPlot
using Base.Dates
using Temporal

aapl = yahoo("AAPL")
aapl = aapl["2015"]

t = aapl.index
aapl = aapl.values
op = aapl[:,1]
hi = aapl[:,2]
lo = aapl[:,3]
cl = aapl[:,4]
vo = aapl[:,5]
hlc = [hi lo cl]

subplot(411)
plot(t, cl, lw=2, c="k", label="AAPL")
plot(t, kama(cl), c="b", label="Kaufman AMA")
plot(t, trima(cl), c="g", label="Triangula MA")
plot(t, hma(cl), c="r", label="Hull MA")
grid(ls="-", c=[0.8,0.8,0.8])
legend(loc="best", frameon=false)

subplot(412)
plot(t, kst(cl), c="m", label="KST")
plot(t, sma(kst(cl),n=9), c="c", label="Signal")
plot([t[1],t[end]], [0,0], ls="--", c=[0.4,0.4,0.4])
grid(ls="-", c=[0.8,0.8,0.8])
legend(loc="best", frameon=false)

subplot(413)
plot(t, wpr(hlc), c=[1,0.5,0], label="Williams %R")
plot([t[1],t[end]], [-20,-20], c="r", ls="--")
plot([t[1],t[end]], [-80,-80], c="g", ls="--")
grid(ls="-", c=[0.8,0.8,0.8])
legend(loc="best", frameon=false)

subplot(414)
plot(t, cci(hlc), c="c", label="CCI")
plot([t[1],t[end]], [-100,-100], c="g", ls="--")
plot([t[1],t[end]], [100,100], c="r", ls="--")
grid(ls="-", c=[0.8,0.8,0.8])
legend(loc="best", frameon=false)

tight_layout()


export NLS_ForwardDiff, create_NLS_problem_using_ForwardDiff

using ForwardDiff: jacobian

@doc raw"""

```julia
struct NLS_ForwardDiff <: AbstractNLS
    ...
end

```

A specialization that uses the `ForwardDiff` package to compute the Jacobian.

By comparison with [`AbstractNLS`](@ref) you only have to define these
functions:
- [`parameter_size`](@ref) : returns ``n_θ``
- [`residue_size`](@ref) : returns ``n_S``
- [`eval_r`](@ref) : computation of ``\mathbf{r}``

See: [`create_NLS_problem_using_ForwardDiff`](@ref)
"""
struct NLS_ForwardDiff <: AbstractNLS
    _eval_r_function::Function
    _residue_size::Int
    _parameter_size::Int
end

@doc raw"""
```julia
create_NLS_problem_using_ForwardDiff(r::Function;domain_image_dim::Pair{Int,Int})
```

Creates an [`AbstractNLS`](@ref) specialized instance where the
[`eval_r_J`](@ref) function is automatically defined using automatic
differentiation.

- `r` is a function that maps a parameter vector θ to its residue. The
  Jacobian matrix is computed using the `ForwardDiff` package.
- `domain_image_dim` is a pair of the form `θ length => r length` that
  defines domain and codomain dimensions.

# Usage example

An example defining the Rosenbrock function

```math
\frac{1}{2}\|r(\theta)\|^2\text{ where }r = \sqrt{2} \left( \begin{array}{c}  1-\theta_1 \\ 10(\theta_2-\theta_1^2) \end{array} \right)
```

```julia
nls = create_NLS_problem_using_ForwardDiff(2 => 2) do θ
     sqrt(2) sqrt(2)* [ 1-θ[1], 10*(θ[2]-θ[1]^2) ]* [ 1-θ[1], 10*(θ[2]-θ[1]^2) ]
end
```
"""
function create_NLS_problem_using_ForwardDiff(r::Function,domain_image_dim::Pair{Int,Int})
    NLS_ForwardDiff(r,last(domain_image_dim),first(domain_image_dim))
end

parameter_size(nls::NLS_ForwardDiff) = nls._parameter_size
residue_size(nls::NLS_ForwardDiff) = nls._residue_size

function eval_r(nls::NLS_ForwardDiff,θ::AbstractVector)
    nls._eval_r_function(θ)
end

function eval_r_J(nls::NLS_ForwardDiff, θ::AbstractVector{T}) where {T}

    r_evaluation = (r,θ)->(r.=eval_r(nls,θ))

    r=Vector{T}(undef,residue_size(nls))

    J = jacobian(r_evaluation, r, θ)

    r,J
end


module TestStateSelection

using TinyModia
using Test

function checkStateSelection(model, x_names, linearEquations=[])
    # Check names of the states
    for (i, state_i) in enumerate(model.equationInfo.x_info)
        @test state_i.x_name == x_names[i]
    end

    # Check linear equation systems
    if length(linearEquations) > 0
        @test model.equationInfo.linearEquations == linearEquations
    end
end

@testset "\nTest ModiaBase/StateSelection.jl" begin

    @testset "... Test FirstOrder" begin
        FirstOrder = Model(
            T = 0.2,
            x = Var(init=0.3),
            equations = :[
                    u = 1.0
                der(x) = (u - x)/T
                    y = 2*x]
        )
        firstOrder = @instantiateModel(FirstOrder)
        checkStateSelection(firstOrder, ["x"])
    end

    @testset "... Test TwoCoupledInertias" begin
        TwoCoupledInertias = Model(
            J1_J = 1,
            J2_J = 1,
            J1_phi = Var(start=1.0),
            J1_w   = Var(start=0.0),
            J2_phi = Var(start=0.0),
            J2_w   = Var(start=0.0),
            equations = :[
                J1_w = der(J1_phi)
                J1_J * der(J1_w) = J1_tau

                J2_w = der(J2_phi)
                J2_J * der(J2_w) = J2_tau

                J2_phi = J1_phi
                J2_tau + J1_tau = 0]
        )
        twoCoupledInertias = @instantiateModel(TwoCoupledInertias)
        checkStateSelection(twoCoupledInertias, ["J1_w", "J1_phi"], [(["J1_tau"], [1], 1, 1)])
    end

    @testset "... Test ODE with linear equations 1" begin
        ODEwithLinearEquations1 = Model(
            p1=4,
            p2=2,
            x6 = Var(start=1.0),
            equations = :[
                         x1 = sin(time)
                      p2*x2 = p1*x1
                2*x3 + 4*x4 = x1
                    x3 + x5 = 2*x2
                4*x4 + 5*x5 = x1 + x2
                    der(x6) = -x5*x6]
        )
        oDEwithLinearEquations1 = @instantiateModel(ODEwithLinearEquations1, unitless=true)
        checkStateSelection(oDEwithLinearEquations1, ["x6"], [(["x5"], [1], 1, 1)])
    end

    @testset "... Test ODE with linear equations 2" begin
        ODEwithLinearEquations2 = Model(
            p1=4,
            p2=2,
            x6 = Var(start=1.0),
            equations = :[
                        x1 = sin(time)
                     p2*x2 = p1*x1
                p2*x3 + p1*x4 + 5*der(x6) = x1
                    x3 + x5 = 2*x2
                4*x4 + 5*x5 = x1 + x2
                  3*der(x6) = -5*x3 - 2*x4]
        )
        oDEwithLinearEquations2 = @instantiateModel(ODEwithLinearEquations2, unitless=true)
        checkStateSelection(oDEwithLinearEquations2, ["x6"], [(["x5"], [1], 1, 1)])
    end

    @testset "... Test Multi-index DAE" begin
        MultiIndexDAE = Model(
            x2 = Var(init = 0.0),
            x2d = Var(init = 0.0),
            equations = :[
                u1 = sin(1*time)
                u2 = sin(2*time)
                u3 = sin(3*time)
                u4 = sin(4*time)
                u5 = sin(5*time)
                u6 = sin(6*time)
                u7 = sin(7*time)
                0 = u1 + x1 - x2
                0 = u2 + x1 + x2 - x3 + der(x6)
                0 = u3 + x1 + der(x3) - x4
                x1d  = der(x1)
                x2d  = der(x2)
                x3d  = der(x3)
                x6d  = der(x6)
                x6dd = der(x6d)
                0 = u4 + 2*der(x1d) + der(x2d) + der(x3d) + der(x4) + der(x6dd)
                0 = u5 + 3*der(x1d) + 2*der(x2d) + x5
                0 = u6 + 2*x6 + x7
                0 = u7 + 3*x6 + 4*x7]
        )
        multiIndexDAE = @instantiateModel(MultiIndexDAE, unitless=true)
        #@show multiIndexDAE.equationInfo.linearEquations
        checkStateSelection(multiIndexDAE, ["x2", "x2d"],
                            [(["x7"], [1], 1, 1),
                             (["der(x7)"], [1], 1, 1),
                             (["der(der(x7))"], [1], 1, 1),
                             (["der(der(der(x7)))"], [1], 1, 1),
                             (["der(x2d)"], [1], 1, 1)])
    end

    @testset "... Test free flying mass" begin
        FreeFlyingMass = Model(
            m =1.0,
            f=[1.0,2.0,3.0],
            r = Var(init=[0.1, 0.2, 0.3]),
            v = Var(init=[-0.1, -0.2, -0.3]),
            equations = :[
                v = der(r)
                m*der(v) = f]
        )
        freeFlyingMass = @instantiateModel(FreeFlyingMass)
        checkStateSelection(freeFlyingMass, ["v", "r"], [])
    end

#=

    @testset "... Test sliding mass" begin
        SlidingMass = Model(
            m = 1.0,
            c = 1e4,
            d = 10.0,
            g = 9.81,
            n = [1.0, 1.0, 1.0],
            s = Var(init=1.0),
            equations = :[
                r = n*s
                v = der(r)
                m*der(v) = f + m*g*[0,-1,0] + u
                0 = n*f
                u = -(c*s + d*der(s))*n]
        )
        slidingMass = @instantiateModel(SlidingMass, logStateSelection=true)
        @show slidingMass.equationInfo.linearEquations
        #checkStateSelection(slidingMass, ["x6"], [(["x5"], [1], 1, 1)])
    end
=#

end

using Unitful
include("../models/AllModels.jl")

#=
Gear = Model(
    flange_a = Flange,
    flange_b = Flange,
    gear     = IdealGear_withSupport | Map(ratio = 105.0),
    fixed    = Fixed,
    spring   = Spring | Map(c=5.84e5u"N*m/rad"),
    damper1  = Damper | Map(d=500.0u"N*m*s/rad"),
    damper2  = Damper | Map(d=100.0u"N*m*s/rad"),

    connect = :[
        (flange_a     , gear.flange_a)
        (fixed.flange , gear.support, damper2.flange_b)
        (gear.flange_b, damper2.flange_a, spring.flange_a, damper1.flange_a)
        (flange_b     , spring.flange_b, damper1.flange_b)]
)

Drive = Model(
    inertia1 = Inertia,
    inertia2 = Inertia,
    gear     = Gear,
    equations = :[
        inertia1.flange_a.tau = 0u"N*m"
        inertia2.flange_b.tau = 0u"N*m"],
    connect = :[
        (inertia1.flange_b, gear.flange_a),
        (gear.flange_b    , inertia2.flange_a)
    ]
)
=#

# Linear 1D rotational damper
AbsoluteDamper = Model(
    flange = Flange,
    d = 1.0u"N*m*s/rad", # (info = "Damping constant"),
    phi = Var(init=0.0u"rad"),
    equation = :[
        phi = flange.phi
        w   = der(phi)
        flange.tau = d * w ]
)

Drive = Model(
    inertia = Inertia,
    damper  = AbsoluteDamper,
    equations = :[
        inertia.flange_a.tau = 0u"N*m"],
    connect = :[
        (inertia.flange_b, damper.flange)]
)
Drive2 = Drive | Map(damper = Map(phi=Var(init=1.0u"rad")))

drive1 = @instantiateModel(Drive)
drive2 = @instantiateModel(Drive2)

using DifferentialEquations
simulate!(drive1, Tsit5(), stopTime = 4.0)

println("Next simulate! should result in an error:\n")
simulate!(drive2, Tsit5(), stopTime = 4.0)

end


using DataArrays,DataFrames,StatsBase

function main()

  path = "../ml-100k/u.data"
  df = readrel(path,'\t')

  indices = shuffle([1:size(df,1)])

  # number of unique users
  totalusers = size(aggregate(df,:user_id,length),1)
  # total number of ratings
  totalratings = length(indices)

  # training

  trainingsetlimit = int(totalratings*(8/10))
  trainingindices  = indices[1:trainingsetlimit]
  trainingset      = df[trainingindices,:]

  # aggregating by ITEMS
  avgratingsbyitem  = aggregate(trainingset,:item_id,mean)[ [:item_id,:rating_mean] ]
  # and by USERS
  avgratingsbyuser  = aggregate(trainingset,:user_id,mean)[ [:user_id,:rating_mean] ]

  # global average for all items
  globalavgbyitem = mean(avgratingsbyitem[:rating_mean])

  # global average for all users
  globalavgbyuser = mean(avgratingsbyuser[:rating_mean])

  # test

  testsetlimit   = trainingsetlimit + (totalratings - trainingsetlimit)
  testindices    = indices[trainingsetlimit:testsetlimit]
  testset        = df[testindices,:]

  global sumerror = 0.0

  testratings = testset[:,[:user_id,:item_id,:rating]]

  totaltestinstances = size(testratings,1)

  for row in eachrow(testratings)

    # guess based upon the item itself
    item_id = row[:item_id]

    if item_id in avgratingsbyitem[:item_id]
      guessbyitem = avgratingsbyitem[:rating_mean][avgratingsbyitem[:item_id] .== item_id]
    else
      guessbyitem = globalavgbyitem
    end

    # guess based upon the user
    user_id = row[:user_id]

    if user_id in avgratingsbyuser[:user_id]
      guessbyuser = avgratingsbyuser[:rating_mean][avgratingsbyuser[:user_id] .== user_id]
    else
      guessbyuser = globalavgbyuser
    end

    actual = row[:rating]

    guess = (guessbyitem + guessbyuser)/2

    error = abs(actual-guess)
    sumerror += error

#     print("user_id: ")
#     print(user_id)
#     print(" guess: ")
#     print(guess)
#     print(" actual: ")
#     print(actual)
#     print(" error: ")
#     println(abs(actual-guess))

  end

  # dividing the error by the number of instances in the test set
  mae = sumerror/totaltestinstances

  mae

end

function readrel(relpath,delimiter=',')

  base = dirname(@__FILE__)

  if relpath[1] == '/'
    relpathwslash = relpath
  else
    relpathwslash = string("/",relpath)
  end

  path = string(base,relpathwslash)

  df = readtable(path,separator = delimiter)

  names!(df,[:user_id,:item_id,:rating,:timestamp])

  df

end

mae = main()
println(mae)


# # @__NAME__

# PREAMBLE

# PKG_SETUP

# ## Setup

using Pkg # hide
Pkg.activate(@__DIR__) # hide
Pkg.instantiate() # hide
using RigidBodyDynamics, StaticArrays, ForwardDiff
using Test, Random
Random.seed!(1); # to get repeatable results

# ## Jacobians with respect to $q$ and $v$ - the naive way

# First, we'll load our trusty double pendulum from a URDF:

srcdir = dirname(pathof(RigidBodyDynamics))
urdf = joinpath(srcdir, "..", "test", "urdf", "Acrobot.urdf")
mechanism = parse_urdf(urdf)

# Of course, we can create a `MechanismState` for the double pendulum, and compute its momentum in some random state:

float64state = MechanismState(mechanism)
rand!(float64state)
momentum(float64state)

# But now suppose we want the Jacobian of momentum with respect to the joint velocity vector $v$.
# We can do this using the `ForwardDiff.Dual` type and the `ForwardDiff.jacobian` function.
# The ForwardDiff package implements forward-mode [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation).

# To use `ForwardDiff.jacobian` we'll create a function that maps `v` (as a `Vector`) to momentum (as a `Vector`):

q = configuration(float64state)
function momentum_vec(v::AbstractVector{T}) where T
    ## create a `MechanismState` that can handle the element type of `v` (which will be some `ForwardDiff.Dual`):
    state = MechanismState{T}(mechanism)

    ## set the state variables:
    set_configuration!(state, q)
    set_velocity!(state, v)

    ## return momentum converted to an `SVector` (as ForwardDiff expects an `AbstractVector`)
    Vector(SVector(momentum(state)))
end

# Let's first check that the function returns the same thing we got from `float64state`:

v = velocity(float64state)
@test momentum_vec(v) == SVector(momentum(float64state))

# That works, so now let's compute the Jacobian with `ForwardDiff`:

J = ForwardDiff.jacobian(momentum_vec, v)

# At this point we note that the matrix `J` is simply the momentum matrix (in world frame) of the `Mechanism`. In this case, RigidBodyDynamics.jl has a specialized algorithm for computing this matrix, so let's verify the results:

A = momentum_matrix(float64state)
@test J ≈ Array(A) atol = 1e-12

# Gradients with respect to $q$ can be computed in similar fashion.

# ## Improving performance

# Ignoring the fact that we have a specialized method available, let's look at the performance of using `ForwardDiff.jacobian`.

using BenchmarkTools
@benchmark ForwardDiff.jacobian($momentum_vec, $v)

# That's not great. Note all the allocations. We can do better by making the following modifications:
#
# 1. use an in-place version of the `jacobian` function, `ForwardDiff.jacobian!`
# 2. reimplement our `momentum_vec` function to be in-place as well
# 3. don't create a new `MechanismState` every time
#
# The third point is especially important; creating a `MechanismState` is expensive!
#
# Regarding the second point, we could also just stop converting momentum from a `StaticArrays.SVector` to a `Vector` to avoid allocations. However, the solution of making the function in-place also applies when the size of the output vector is not known statically (e.g., for `dynamics_bias!`).

# To facillitate reuse of `MechanismState`s while keeping the code nice and generic, we can use a `StateCache` object.
# `StateCache` is a container that stores `MechanismState`s of various types (associated with one `Mechanism`), and will ease the process of using `ForwardDiff`.
# Creating one is easy:

const statecache = StateCache(mechanism)

# `MechanismState`s of a given type can be accessed as follows (note that if a `MechanismState` of a certain type is already available, it will be reused):

float32state = statecache[Float32]
@test float32state === statecache[Float32]

# Now we'll use the `StateCache` to reimplement `momentum_vec`, making it in-place as well:

function momentum_vec!(out::AbstractVector, v::AbstractVector{T}) where T
    ## retrieve a `MechanismState` that can handle the element type of `v`:
    state = statecache[T]

    ## set the state variables:
    set_configuration!(state, q)
    set_velocity!(state, v)

    ## compute momentum and store it in `out`
    m = momentum(state)
    copyto!(out, SVector(m))
end

# Check that the in-place version works as expected on `Float64` inputs:

const out = zeros(6) # where we'll be storing our results
momentum_vec!(out, v)
@test out == SVector(momentum(float64state))

# And use `ForwardDiff.jacobian!` to compute the Jacobian:

const result = DiffResults.JacobianResult(out, v)
const config = ForwardDiff.JacobianConfig(momentum_vec!, out, v)
ForwardDiff.jacobian!(result, momentum_vec!, out, v, config)
J = DiffResults.jacobian(result)
@test J ≈ Array(A) atol = 1e-12

# Let's check the performance again:

@benchmark ForwardDiff.jacobian!($result, $momentum_vec!, $out, $v, $config)

# That's much better. Do note that the specialized algorithm is still faster:

q = copy(configuration(float64state))
@benchmark begin
    set_configuration!($float64state, $q)
    momentum_matrix!($A, $float64state)
end

# ## Time derivatives

# We can also use ForwardDiff to compute time derivatives. Let's verify that energy is conserved for the double pendulum in the absence of nonconservative forces (like damping). That is, we expect that the time derivative of the pendulum's total energy is zero when its state evolves according to the passive dynamics.

# Let's first compute the joint acceleration vector $\dot{v}$ using the passive dynamics:

dynamicsresult = DynamicsResult(mechanism)
set_configuration!(float64state, q)
set_velocity!(float64state, v)
dynamics!(dynamicsresult, float64state)
v̇ = dynamicsresult.v̇

# Now for the time derivative of total energy. ForwardDiff has a `derivative` function that can be used to take derivatives of functions that map a scalar to a scalar. But in this example, we'll instead use ForwardDiff's `Dual` type directly. `ForwardDiff.Dual` represents a (potentially multidimensional) dual number, i.e., a type that stores both the value of a function evaluated at a certain point, as well as the partial derivatives of the function, again evaluated at the same point. See the [ForwardDiff documentation](http://www.juliadiff.org/ForwardDiff.jl/stable/dev/how_it_works.html) for more information.

# We'll create a vector of `Dual`s representing the value and derivative of $q(t)$:

q̇ = v
q_dual = ForwardDiff.Dual.(q, q̇)

# **Note**: for the double pendulum, $\dot{q} = v$, but this is not the case in general for `Mechanism`s created using RigidBodyDynamics.jl. For example, the `QuaternionSpherical` joint type uses a unit quaternion to represent the joint configuration, but angular velocity (in body frame) to represent velocity. In general $\dot{q}$ can be computed from the velocity vector $v$ stored in a `MechanismState` using
#
# ```julia
# configuration_derivative(::MechanismState)
# ```
#
# or its in-place variant, `configuration_derivative!`.

# We'll do the same thing for $v(t)$:

v_dual = ForwardDiff.Dual.(v, v̇)

# Now we're ready to compute the total energy (kinetic + potential) using these `ForwardDiff.Dual` inputs. We'll use our `StateCache` again:

T = eltype(q_dual)
state = statecache[T]
set_configuration!(state, q_dual)
set_velocity!(state, v_dual)
energy_dual = kinetic_energy(state) + gravitational_potential_energy(state)

# Note that the result type of `energy_dual` is again a `ForwardDiff.Dual`. We can extract the energy and its time derivative (mechanical power) from `energy_dual` as follows:

energy = ForwardDiff.value(energy_dual)
partials = ForwardDiff.partials(energy_dual)
power = partials[1];

# So the total energy in the system is:

energy

# **Note**: the total energy is negative because the origin of the world frame is used as a reference for computing gravitational potential energy, i.e., the center of mass of the double pendulum is somewhere below this origin.

# And we can verify that, indeed, there is no power flow into or out of the system:

@test power ≈ 0 atol = 1e-14

examples/5. Derivatives and gradients using ForwardDiff/5. Derivatives and gradients using ForwardDiff.jl


### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 7357f256-ead0-401c-9cff-1aa099b67aa0
begin
    import Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
	    Pkg.PackageSpec(name="DifferentialEquations", version="6"),
	    Pkg.PackageSpec(name="Parameters", version="0.12"),
        Pkg.PackageSpec(name="LabelledArrays", version="1"),
        Pkg.PackageSpec(name="Setfield", version="0.7"),
    ])
    using Plots, PlutoUI, LinearAlgebra, DifferentialEquations, Parameters, LabelledArrays, Setfield

	Plots.gr(fmt=:png, lw=2)

	PlutoUI.TableOfContents()
end

# ╔═╡ 0e455e50-bdcb-11eb-08b9-8d2d26788d14
md"""
# Fig 1.07 Collins toggle switch

For Figures 1.7, 7.13, 7.14, 7.15
"""

# ╔═╡ 30cdc452-ff17-4b08-939a-98c7dbc0c674
# Initial conditions
u0 = LVector(s1=0.075, s2=2.5)

# ╔═╡ 520f36c3-e164-459e-8869-99a8421774f6
# Simulation time
tend = 50.0

# ╔═╡ f47e84f9-7ff5-4c32-996a-47c600b1a71a
# Parameters
p = (a1 = 3.0, a2 = 2.5, β = 4.0, γ = 4.0, i1 = 0.0, i2 = 0.0)

# ╔═╡ 463ef6e7-73db-4b55-b97f-71e90c2f14f4
begin
	# Events and callbacks
    on_i1!(integrator) = integrator.p = @set p.i1 = 10.0
    off_i1!(integrator) = integrator.p = @set p.i1 = 0.0
    on_i2!(integrator) = integrator.p = @set p.i2 = 10.0
    off_i2!(integrator) = integrator.p = @set p.i2 = 0.0

    cbs = CallbackSet(PresetTimeCallback([10.0], on_i2!),
                  PresetTimeCallback([20.0], off_i2!),
                  PresetTimeCallback([30.0], on_i1!),
                  PresetTimeCallback([40.0], off_i1!))
end

# ╔═╡ d29832bb-74f9-443a-af7f-a12e7ae652c3
md"""
# Appendix
Running environment and some auxillary functions.

**You can use [this helper tool](https://fonsp.com/article-test-3/pkghelper.html) to generate these commands!**
"""

# ╔═╡ 796ff12c-b74d-4fcf-b228-4e3c6fed2a89
begin
	# Convenience functions
    hill(x, k) = x / (x + k)
    hill(x, k, n) = hill(x^n, k^n)
    exprel(x) = ifelse(x≈zero(x), one(x), x / expm1(x))
end

# ╔═╡ 34178e95-cff5-400e-a5a6-bdbcf8326214
# The model
	function collins!(du, u, p, t)
		@unpack a1, a2, β, γ, i1, i2 = p
		@unpack s1, s2 = u
		du.s1 = a1 * hill(1 + i2, s2, β) - s1
		du.s2 = a2 * hill(1 + i1, s1, γ) - s2
    	return du
	end

# ╔═╡ 759580e6-0910-43a2-b81a-cdaa3bf34441
sol = solve(ODEProblem(collins!, u0, tend, p), callback=cbs)

# ╔═╡ b922fac1-e43c-45eb-a6db-0a9b18f541aa
plot(sol, legend=:right, xlabel="Time", ylabel="Concentration", title="Figure 1.7")

# ╔═╡ Cell order:
# ╠═0e455e50-bdcb-11eb-08b9-8d2d26788d14
# ╠═34178e95-cff5-400e-a5a6-bdbcf8326214
# ╠═463ef6e7-73db-4b55-b97f-71e90c2f14f4
# ╠═30cdc452-ff17-4b08-939a-98c7dbc0c674
# ╠═520f36c3-e164-459e-8869-99a8421774f6
# ╠═f47e84f9-7ff5-4c32-996a-47c600b1a71a
# ╠═759580e6-0910-43a2-b81a-cdaa3bf34441
# ╠═b922fac1-e43c-45eb-a6db-0a9b18f541aa
# ╠═d29832bb-74f9-443a-af7f-a12e7ae652c3
# ╠═796ff12c-b74d-4fcf-b228-4e3c6fed2a89
# ╠═7357f256-ead0-401c-9cff-1aa099b67aa0


import Combinatorics

"""
For every right kernel vector of M with nz nonzeros, finds
a vector [i1, ..., inz] of the  indices of the nonzero
entries. Returns an array with at most 'count' index vectors.
If 'count' is negative, finds all (default).
"""
function kernelnz(M::AbstractMatrix, nz::Integer; count::Integer = -1)
    @assert nz > 0
    M = full(M)
    ker = Vector{Vector{Int}}()
    for c in Combinatorics.combinations(1 : size(M, 2), nz)
        @assert length(c) == nz
        0 ≤ count ≤ length(ker) && break
        if rank(M[:,c]) < nz
            push!(ker, c)
        end
    end
    return ker
end

"Returns true if there is a vector v with `nz` non-zeros such that M⋅v = 0."
haskernelnz(M::AbstractMatrix, nz::Integer) = !isempty(kernelnz(M, nz; count = 1))

"""
Matrix flip around a diagonal:
   (a b     (d b
    c d) ->  c a)
"""
function transpoflip(A::Matrix)
    @assert size(A, 1) == size(A, 2)
    N = size(A, 1)
    B = [A[N + 1 - j, N + 1 - i] for i = 1 : N, j = 1 : N]
end

"""
Given the inverse A of a matrix M, returns the inverse matrix
of the matrix M', which differs from M only in n'th diagonal
entry, for which M'[n,n] = M[n,n] + δ. Uses the Sherman-Morrison-Woodbury
formula. See D.S. Bernestein, Matrix Mathematics 2009, Fact 2.16.3.
"""
function invupdate end

function invupdate(A::AbstractMatrix, n::Integer, δ::Real)
    B = copy(A)
    for i = 1 : size(A, 1), j = 1 : size(A, 2)
        B[i,j] = A[i,j] - δ * A[i,n] * A[n,j] / (1 + δ * A[n,n])
    end
    return B
end

function invupdate(A::Hermitian, n::Integer, δ::Real)
    B = copy(A)
    for i = 1 : size(A, 1), j = 1 : i
        Bij = A[i,j] - δ * A[i,n] * A[n,j] / (1 + δ * A[n,n])
        B.data[i,j] = B.data[j,i]= Bij
    end
    return B
end


#plotmaker.jl
#I'm doing this in julia because it's easier than in Python
using PGFPlots, DataFrames, CSV, RollingFunctions, LinearAlgebra

df = DataFrame(CSV.File("current_data/dqn.csv"))
name_array = names(df)
# temp = Matrix(df)
# samples = I*temp'
rollouts = Matrix(df)
episodes = (rollouts[:,1]*I)[:]
num_episodes = size(episodes)[1]
reward = (rollouts[:,4]*I)[:] #col 4 for dqn
# reward_std = (rollouts[:,5]*I)[:]
# threshold =
moving_avg = 50
smoothed_reward = rolling(sum, reward, moving_avg)./moving_avg

p = Axis([
    Plots.Scatter(episodes, reward, legendentry="Rollout Data", markSize=0.1),
    Plots.Linear(episodes[moving_avg:end], smoothed_reward, legendentry="Average Reward", mark="none"),
    Plots.Linear(episodes, 300*ones(num_episodes), legendentry="Reward Threshold", mark="none", style="black,dashed")
    ], xlabel="Episode", ylabel="Reward", title="Accumulated Reward vs Episode", xmin=0, xmax=num_episodes
    )
# savefig("plots/biped_ars001.pdf")
# a = Axis(Plots.Linear(x, y, legendentry="My Plot"), xlabel="X", ylabel="Y", title="My Title")
# p = Axis(Plots.Linear(episodes, rollouts, legendentry="reward"))#, markSize=0.1, onlyMarks=true)
# p = Axis(Plots.Linear(episodes, rollouts, legendentry="reward", markSize=0.1, onlyMarks=true), xlabel="Episode", ylabel="Reward", title="Episode Reward vs Episode Number")
p.legendStyle = "at={(1.05,1.0)}, anchor=north west"
save("plots/dqn.pdf", p)


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 29443685-beb0-4102-a46d-63c48b008742
begin
	using LinearAlgebra
	using Plots
	plotly()
end

# ╔═╡ bdeabaa3-8a91-4f9d-8cb7-4527a6385438
md"""
This notebook is a Julia translation of Josh Bevan's [intro to DG YouTube video](https://www.youtube.com/watch?v=sMfPHJUKfaI).
"""

# ╔═╡ dbbb3fea-9dea-11eb-1f7b-afb4a0c5114e
md"""
Here we attempt to solve the 1D scalar conservation eqn of the form:

$ \displaystyle \frac{\partial u}{\partial t} + \frac{\partial f(u)}{\partial x} = 0 $

where $f(u)$ is some flux function describing the "flow" of a conserved
quantity. In this simplified case $f(u) = u$ giving us a linear PDE
We will use periodic boundary conditions to examine numerical dissipation
effects
"""

# ╔═╡ 35708c02-2ef3-432f-963b-62b77f072b19
md"""
Let $u(x,t)$ be the exact solution on the domain $0 \leq x \leq 1$.

Let $u(x,0) = u_0(x) = \sin(2\pi x)$.

Let $V_h$ be the finite vector space of linear polynomials.

Let $u_h$ be the approximate numerical solution consisting of a linear
combination of basis functions ($\psi_n$ for the nth basis) in $V_h$ with
scalar coefficients $a_n$.

Let $\phi_n$ be the nth test function in the same vector space as the basis
functions ($V_h$), in this simplified case N=2.
"""

# ╔═╡ 14c235cf-7b61-45b2-85a8-3d8af45b4edb
N = 2

# ╔═╡ 6be6e442-6983-4629-b734-e6487ac8001d
md"""
Discretize the domain into $K$ elements with $K+1$ nodes.
"""

# ╔═╡ ec7aeb56-230b-4d05-a561-632f71c09e02
struct GridElement{T}
	xl :: T  # Left endpoint
	xr :: T  # Right endpoint
end

# ╔═╡ 34a301b0-58b4-4617-ad47-bd7883b09d8b
struct Grid{T, E}
	K :: Int
	Δ :: T
	elements :: E
end

# ╔═╡ d7c33092-ca55-4e1f-92be-d62aef900393
begin
	K = 32
	Δx = 1/K
	elements = [GridElement(i*Δx, (i+1)*Δx) for i in 0:K-1]
	grid = Grid(K, Δx, elements)
end

# ╔═╡ 28c5d715-c14a-4573-81b5-1e835bc4926c
md"""
Compute and store the initial condition $u_0(x) = \sin(2\pi x)$.
"""

# ╔═╡ d75f11b0-e29b-41d6-87bb-108b0d3cb468
struct FieldElement{D}
	data :: D
end

# ╔═╡ 3f8100e7-dd79-44eb-864c-45ae08c6445a
begin
	struct Field{G, E}
		grid :: G
		elements :: E
	end

	function Field(grid, f::Function)
		field_elements = []
		for grid_element in grid.elements
			field_data = [f(grid_element.xl), f(grid_element.xr)]
			field_element = FieldElement(field_data)
			push!(field_elements, field_element)
		end
		return Field(grid, field_elements)
	end
end

# ╔═╡ b52a8325-a0e0-4697-81b4-26abfbe70a98
begin
	u₀(x) = sin(2π*x)
	u = Field(grid, u₀)
end

# ╔═╡ 45016410-7e25-4f5b-82c5-aa665aaa3c7f
md"""
Let's plot ``u_0(x)`` now!
"""

# ╔═╡ 2f617585-62c1-4f6c-9088-aed874dc43e0
begin
	p = plot()
	for (ge, ue) in zip(grid.elements, u.elements)
		plot!([ge.xl, ge.xr], ue.data, linewidth=2, label="")
	end
	p
end

# ╔═╡ 4de78270-1f76-4f66-9271-34d0f148bc0c
md"""
According to Cockburn & Shu (2001, eq 2.2) let $u_h(x,0)$ be computed by

`` \displaystyle \int u_h(x,0) \phi(x) dx = \int u_0(x) \phi(x) dx ``

for each element (``x_{j-1/2} < x < x_{j+1/2}``).

We can explicitly define a formula for the value of the RHS

`` \displaystyle \int u_0(x) \phi_N(x) dx ``
"""

# ╔═╡ 98e97ab0-6623-4a9e-a9b9-fb80ca182f53
md"""
For each element in the domain $a \le x \le b$, the cell coordinate is $0 \le \xi \le 1$ where

`` x = a + (b - a) \xi  = a + \Delta \xi ``

and

`` \displaystyle \xi = \frac{x - a}{\Delta} ``

so the basis ramp function can be written as

`` \displaystyle \phi_0(\xi) = 1 - \xi \implies \phi_0(x) = \frac{b - x}{\Delta} ``

`` \displaystyle \phi_1(\xi) = \xi \implies \phi_1(x) = \frac{x - a}{\Delta} ``
"""

# ╔═╡ c22c283e-615e-4964-9ff1-f9102f6862b2
md"""
For the left endpoint

``
\displaystyle
\int_a^b u_0(x) \phi_0(x) \; dx
= \int_a^b \sin(2\pi x) \frac{b-x}{\Delta} \; dx
= \frac{2\pi\Delta \cos(2\pi a) + \sin(2\pi a) - \sin(2\pi b)}{4\pi^2\Delta}
``
"""

# ╔═╡ 520fcb4f-bffc-4d53-97b1-703dfaf0b8ca
md"""
For the right endpoint

``
\displaystyle
\int_a^b u_0(x) \phi_1(x) \; dx
= \int_a^b \sin(2\pi x) \frac{x-a}{\Delta} \; dx
= \frac{-2\pi\Delta \cos(2\pi b) + \sin(2\pi b) - \sin(2\pi a)}{4\pi^2\Delta}
``
"""

# ╔═╡ 1096111c-9281-46bc-b5a4-4178a9855cad
begin
	∫u₀ϕ₀dx(a, b, Δ) = (2π*Δ*cos(2π*a) + sin(2π*a) - sin(2π*b)) / (4π^2*Δ)
	∫u₀ϕ₁dx(a, b, Δ) = (-2π*Δ*cos(2π*b) + sin(2π*b) - sin(2π*a)) / (4π^2*Δ)

	∫u₀ϕ_elements = []
	for grid_element in grid.elements
		xₗ = grid_element.xl
		xᵣ = grid_element.xr

		e0 = ∫u₀ϕ₀dx(xₗ, xᵣ, grid.Δ)
		e1 = ∫u₀ϕ₁dx(xₗ, xᵣ, grid.Δ)

		field_element = FieldElement([e0, e1])
		push!(∫u₀ϕ_elements, field_element)
	end

	∫u₀ϕ = Field(grid, ∫u₀ϕ_elements)
end

# ╔═╡ 16e1782b-3948-4d9b-8bfd-fb3134211977
md"""
To find the initial basis weights we need to solve for $a_i$ from

``
\displaystyle
\int_a^b u_0(x) \phi_j(x) \; dx
= \sum_{i=0}^1 a_i \int_a^b \phi_i(x) \phi_j(x) \; dx
``

or from the system of linear equations

``
\begin{pmatrix}
\int u_0 \phi_0
\int u_0 \phi_1
\end{pmatrix}
``
"""

# ╔═╡ 7bc02a02-947b-4526-8bed-49007f3fa8a4
mass_matrix = [1/3 1/6;
               1/6 1/3]

# ╔═╡ 92bdd228-617f-42da-b3d9-b9e9e18a7e4b
begin
	u_basis_elements = []
	for (k, grid_element) in enumerate(grid.elements)
		e = mass_matrix \ ∫u₀ϕ.elements[k].data / grid.Δ
		u_basis_element = FieldElement(e)
		push!(u_basis_elements, u_basis_element)
	end
	u_basis = Field(grid, u_basis_elements)
end

# ╔═╡ 50142bc9-3dcd-4817-b9a6-3c4f37b292b1
begin
	p2 = plot()
	for (ge, ue) in zip(grid.elements, u_basis.elements)
		plot!([ge.xl, ge.xr], ue.data, linewidth=2, label="")
	end
	p2
end

# ╔═╡ 4a103139-e2f9-4182-89b0-9b0b3b4b2503
stiffness_matrix = [0 -1/2 -1/2;
                    0  1/2  1/2]

# ╔═╡ 5c7f8b83-1830-4466-82a2-dd2a6da4115b
upwind_flux = [1 0  0;
               0 0 -1]

# ╔═╡ 2d955948-bf0e-4f89-8bfa-2ffa6b9ac328
stencil = mass_matrix \ (stiffness_matrix + upwind_flux)

# ╔═╡ aed204c6-d68d-449a-aee7-b5fcc537b81a
Δt = 1e-3

# ╔═╡ b7d972b4-82de-4c81-8a9a-80b98850763f
linear_operator = [0 1 0; 0 0 1] + Δt * stencil / grid.Δ

# ╔═╡ 741398fd-cfb6-46af-9bd0-ad047f8f4187
gr()

# ╔═╡ ca064da4-ab88-468e-a2cb-5b76debf6e95
begin
	iterations = 1000
	anim = @animate for i in 1:iterations
		@info "Iteration $i"
		for k in 1:length(grid.elements)
			left_index = k == 1 ? grid.K : k-1

			left_element = u_basis.elements[left_index]
			current_element = u_basis.elements[k]

			u_basis_local = [left_element.data[end], current_element.data...]
			current_element.data .= linear_operator * u_basis_local
		end

		pp = plot()
		for (ge, ue) in zip(grid.elements, u_basis.elements)
			plot!(pp, [ge.xl, ge.xr], ue.data, linewidth=2, label="")
		end
	end every 5
	mp4(anim)
end

# ╔═╡ Cell order:
# ╠═29443685-beb0-4102-a46d-63c48b008742
# ╟─bdeabaa3-8a91-4f9d-8cb7-4527a6385438
# ╟─dbbb3fea-9dea-11eb-1f7b-afb4a0c5114e
# ╟─35708c02-2ef3-432f-963b-62b77f072b19
# ╟─14c235cf-7b61-45b2-85a8-3d8af45b4edb
# ╟─6be6e442-6983-4629-b734-e6487ac8001d
# ╠═ec7aeb56-230b-4d05-a561-632f71c09e02
# ╠═34a301b0-58b4-4617-ad47-bd7883b09d8b
# ╠═d7c33092-ca55-4e1f-92be-d62aef900393
# ╟─28c5d715-c14a-4573-81b5-1e835bc4926c
# ╠═d75f11b0-e29b-41d6-87bb-108b0d3cb468
# ╠═3f8100e7-dd79-44eb-864c-45ae08c6445a
# ╠═b52a8325-a0e0-4697-81b4-26abfbe70a98
# ╟─45016410-7e25-4f5b-82c5-aa665aaa3c7f
# ╠═2f617585-62c1-4f6c-9088-aed874dc43e0
# ╟─4de78270-1f76-4f66-9271-34d0f148bc0c
# ╟─98e97ab0-6623-4a9e-a9b9-fb80ca182f53
# ╟─c22c283e-615e-4964-9ff1-f9102f6862b2
# ╟─520fcb4f-bffc-4d53-97b1-703dfaf0b8ca
# ╠═1096111c-9281-46bc-b5a4-4178a9855cad
# ╟─16e1782b-3948-4d9b-8bfd-fb3134211977
# ╠═7bc02a02-947b-4526-8bed-49007f3fa8a4
# ╠═92bdd228-617f-42da-b3d9-b9e9e18a7e4b
# ╠═50142bc9-3dcd-4817-b9a6-3c4f37b292b1
# ╠═4a103139-e2f9-4182-89b0-9b0b3b4b2503
# ╠═5c7f8b83-1830-4466-82a2-dd2a6da4115b
# ╠═2d955948-bf0e-4f89-8bfa-2ffa6b9ac328
# ╠═aed204c6-d68d-449a-aee7-b5fcc537b81a
# ╠═b7d972b4-82de-4c81-8a9a-80b98850763f
# ╠═741398fd-cfb6-46af-9bd0-ad047f8f4187
# ╠═ca064da4-ab88-468e-a2cb-5b76debf6e95


using ApproxFun, Compat, Base.Test

## Rectangle PDE

dx=dy=Interval()
d=dx*dy
x=Fun(identity,dx);y=Fun(identity,dy)

#dirichlet(d) is u[-1,:],u[1,:],u[:,-1],u[:,1]

G=[real(exp(-1+1.0im*y));
                        real(exp(1+1im*y));
                        real(exp(x-1im));
                        real(exp(x+1im));0.];

A=[dirichlet(d);lap(d)]
u=A\G
@test_approx_eq u(.1,.2) real(exp(0.1+0.2im))

A=[dirichlet(d);lap(d)+0.0I]
u=A\G
@test_approx_eq_eps u(.1,.2) real(exp(0.1+0.2im)) 1E-11

println("    Poisson tests")

## Poisson

f=Fun((x,y)->exp(-10(x+.2)^2-20(y-.1)^2))  #default is [-1,1]^2
d=domain(f)
OS=S=schurfact([dirichlet(d);lap(d)],10)
u=OS\[zeros(∂(d));f]
@test_approx_eq u(.1,.2) -0.042393137972085826

d=PeriodicInterval()^2
f=ProductFun((x,y)->exp(-10(sin(x/2)^2+sin(y/2)^2)),d)
A=lap(d)+.1I
u=A\f
@test (lap(u)+.1u-f)|>coefficients|>norm < 1000000eps()

println("    Kron tests")

@static if is_apple()
    ## Kron

    dx=dy=Interval()
    d=dx*dy
    x=Fun(identity,dx);y=Fun(identity,dy)

    #dirichlet(d) is u[-1,:],u[1,:],u[:,-1],u[:,1]

    G=[real(exp(-1+1.0im*y));
                            real(exp(1+1im*y));
                            real(exp(x-1im));
                            real(exp(x+1im));0.];

    A=[dirichlet(d);lap(d)]

    S=schurfact(A,40)

    uex=A\G

    nx=ny=40
    K=kronfact(A,nx,ny)

    uex2=K\G

    @test (uex-uex2|>coefficients|>norm)<10000eps()

    # dirichlet bcs

    import ApproxFun.ChebyshevDirichlet

    S=ChebyshevDirichlet()⊗ChebyshevDirichlet();
    A=[dirichlet(S);lap(S)]
    nx=ny=20;
    KD=kronfact(A,nx,ny);

    #dirichlet(d) is u[-1,:],u[1,:],u[:,-1],u[:,1]
    x=Fun(identity);y=Fun(identity);
    G=[Fun(real(exp(-1+1.0im*y)),S[2]);
        Fun(real(exp(1+1im*y)),S[2]);
        Fun(real(exp(x-1im)),S[1]);
                            Fun(real(exp(x+1im)),S[1]);0.];

    uD=KD\G;

    @test_approx_eq uD(.1,.2) real(exp(.1+.2im))

    # fourth order
    dx=dy=Interval()
    d=dx*dy
    Dx=Derivative(dx);Dy=Derivative(dy)
    L=Dx^4⊗I+2*Dx^2⊗Dy^2+I⊗Dy^4

    K=kronfact([dirichlet(d);
         neumann(d);
         L],100,100)

    x=Fun(identity,dx);y=Fun(identity,dy)

    G=[real(exp(-1+1.0im*y));
                    real(exp(1+1im*y));
                    real(exp(x-1im));
                    real(exp(x+1im));
                    real(exp(-1+1.0im*y));
                    real(exp(1+1im*y));
                    -imag(exp(x-1im));
                    -imag(exp(x+1im))
       ]
    u=K\G
    @test_approx_eq u(.1,.2) real(exp(.1+.2im))

    # mixed

    K=kronfact([(ldirichlet(dx)+lneumann(dx))⊗I;
            (rdirichlet(dx)+rneumann(dx))⊗I;
            I⊗(ldirichlet(dy)+lneumann(dy));
            I⊗(rdirichlet(dy)+rneumann(dy));
            (ldirichlet(dx)-lneumann(dx))⊗I;
            (rdirichlet(dx)-rneumann(dx))⊗I;
            I⊗(ldirichlet(dy)-lneumann(dy));
            I⊗(rdirichlet(dy)-rneumann(dy));
             L],100,100)
    G=[2real(exp(-1+1.0im*y));
                    2real(exp(1+1im*y));
                    real(exp(x-1im))-imag(exp(x-1im));
                    real(exp(x+1im))-imag(exp(x+1im));
                    0;
                    0;
                    real(exp(x-1im))+imag(exp(x-1im));
                    real(exp(x+1im))+imag(exp(x+1im))
       ]
    u=K\G

    @test_approx_eq u(.1,.2) real(exp(.1+.2im))
end

## Test periodic x interval

d=PeriodicInterval()*Interval()
g=Fun(z->real(cos(z)),∂(d))  # boundary data
u=[dirichlet(d);lap(d)]\g

@test_approx_eq u(.1,.2) real(cos(.1+.2im))

dθ=PeriodicInterval(-2.,2.);dt=Interval(0,3.)
d=dθ*dt
Dθ=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])
u=[I⊗ldirichlet(dt);Dt+Dθ]\Fun(θ->exp(-20θ^2),dθ)

@test_approx_eq u(.1,.2) exp(-20(0.1-0.2)^2)

println("   Domain Decompositon tests")

## Domain Decomposition
d=Interval(0,1)^2
n,m=20,80
A=discretize([dirichlet(d);lap(d)],n)
∂d=∂(d)
g=Fun(z->real(exp(z)),∂d)
f=[Fun([zeros(k-1);1.0],∂d) for k=1:m].'
U=A\f
@test_approx_eq dot(real(g.coefficients),U[1:ncoefficients(g)])(.1,.2) real(exp(.1+.2im))

Rectangle(a,b,c,d)=Interval(a,b)*Interval(c,d)
Γ=Rectangle(0,1,0,1)∪Rectangle(1,2,0,1)
Fun(identity,∂(Γ))|>values

## Small diffusoion

dx=Interval();dt=Interval(0,1.)
d=dx*dt
Dx=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])
x=Fun(identity,dx)
B=0.0
C=0.0
V=B+C*x
ε=0.001
f=Fun(x->exp(-20x^2),dx)
u=[timedirichlet(d);Dt-ε*Dx^2-V*Dx]\f
@test_approx_eq u(.1,.2) 0.8148207991358946
B=0.1
C=0.2
V=B+C*x
u=[timedirichlet(d);Dt-ε*Dx^2-V*Dx]\f
@test_approx_eq u(.1,.2) 0.7311625132209619

## Schrodinger

dx=Interval(0.,1.);dt=Interval(0.0,.1)
d=dx*dt

V=Fun(x->x^2,dx)

Dt=Derivative(d,[0,1]);Dx=Derivative(d,[1,0])

ϵ=1.
u0=Fun(x->exp(-100*(x-.5)^2)*exp(-1./(5*ϵ)*log(2cosh(5*(x-.5)))),dx)
L=ϵ*Dt+(.5im*ϵ^2*Dx^2)
ny=200;u=pdesolve([timedirichlet(d);L],u0,ny)
@test_approx_eq_eps u(.2,.1) (0.2937741918470843 + 0.22130344715160255im )  0.000001

## Periodic

d=PeriodicInterval()^2
f=Fun((θ,ϕ)->exp(-10(sin(θ/2)^2+sin(ϕ/2)^2)),d)
A=lap(d)+.1I
u=A\f
@test_approx_eq u(.1,.2) u(.2,.1)

d=PeriodicInterval()*Interval()
g=Fun(z->real(cos(z)),∂(d))  # boundary data
u=[dirichlet(d);lap(d)]\g
@test_approx_eq u(.1,.2) real(cos(.1+.2im))

dθ=PeriodicInterval(-2.,2.);dt=Interval(0,3.)
d=dθ*dt
Dθ=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])
u=[I⊗ldirichlet(dt);Dt+Dθ]\Fun(θ->exp(-20θ^2),dθ)

# Check bug in cache
CO=cache(ldirichlet(dt))
ApproxFun.resizedata!(CO,:,2)
ApproxFun.resizedata!(CO,:,4)
@test_approx_eq CO*Fun(exp,dt) 1.0

dθ=PeriodicInterval(-2.,2.);dt=Interval(0,3.)
d=dt*dθ
Dt=Derivative(d,[1,0]);Dθ=Derivative(d,[0,1])
A=[ldirichlet(dt)⊗I;Dt+Dθ]
f=Fun(θ->exp(-20θ^2),dθ)
ut=A\f

@test_approx_eq u(.1,.2) ut(.2,.1)

# Beam

dθ=PeriodicInterval(0.0,1.0);dt=Interval(0,0.03)
d=dθ*dt
Dθ=Derivative(d,[1,0]);Dt=Derivative(d,[0,1]);

B=[I⊗ldirichlet(dt),I⊗lneumann(dt)]
u=pdesolve([B;Dt^2+Dθ^4],Fun(θ->exp(-200(θ-.5).^2),dθ),200)

@test_approx_eq_eps u(.1,.01) -0.2479768394633227  1E-8 #empirical

## Rectangle PDEs

# Screened Poisson

d=Interval()^2
u=[neumann(d);lap(d)-100.0I]\ones(∂(d))
@test_approx_eq u(.1,.9) 0.03679861429138079

# PiecewisePDE

a=Fun([1,0.5,1],[-1.,0.,0.5,1.])
s=space(a)
dt=Interval(0,2.)
Dx=Derivative(s);Dt=Derivative(dt)
Bx=[ldirichlet(s);continuity(s,0)]

# test resize bug
CO=cache(Bx[2])
ApproxFun.resizedata!(CO,:,2)
ApproxFun.resizedata!(CO,:,4)
@test_approx_eq CO.data*collect(1:4) [3.,-1.]

u=pdesolve([I⊗ldirichlet(dt);Bx⊗I;I⊗Dt+(a*Dx)⊗I],Any[Fun(x->exp(-20(x+0.5)^2),s)],200)
@test_approx_eq_eps u(-.1,.2) exp(-20(-.2-.1+0.5)^2) 0.00001

## Test error

dx=Interval();dt=Interval(0,2.)
d=dx*dt
Dx=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])
x=Fun(identity,dx)
u=[I⊗ldirichlet(dt);Dt+x*Dx]\Fun(x->exp(-20x^2),dx)

@test_approx_eq u(0.1,0.2) 0.8745340845783758  # empirical

dθ=PeriodicInterval();dt=Interval(0,10.)
d=dθ*dt
ε=.01
Dθ=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])

# Parentheses are a hack to get rank 2 PDE
u=[I⊗ldirichlet(dt);Dt-ε*Dθ^2-Dθ]\Fun(θ->exp(-20θ^2),dθ)

@test_approx_eq_eps u(0.1,0.2) 0.1967278179230314 1000eps()


using Test
using SymbolicRegression

include("test_params.jl")

## Test Base.print
options = Options(;
    default_params..., binary_operators=(+, *, /, -), unary_operators=(cos, sin)
)

f = (x1, x2, x3) -> (sin(cos(sin(cos(x1) * x3) * 3.0) * -0.5) + 2.0) * 5.0

tree = f(Node("x1"), Node("x2"), Node("x3"))

s = repr(tree)
true_s = "((sin(cos(sin(cos(x1) * x3) * 3.0) * -0.5) + 2.0) * 5.0)"

@test s == true_s

# TODO: Next, we test that custom varMaps work:

EquationSearch(
    randn(Float32, 3, 10),
    randn(Float32, 10);
    options=options,
    varMap=["v1", "v2", "v3"],
    niterations=0,
    multithreading=true,
)

s = repr(tree)
true_s = "((sin(cos(sin(cos(v1) * v3) * 3.0) * -0.5) + 2.0) * 5.0)"
@test s == true_s


using Plots
function grid_test(name; xlim=(-5,5),ylim=(-5,5), dims=[1,2], ϵ=-1e-1,order=5,geom=true,condition=100)

	hyperparameters = getParameters(X)
	parameters = copy(θ)

	function ∇cost(x::Number,y::Number)
		parameters[dims] .= x,y
		return ∇loss(F,parameters,X,hyperparameters)[2]
	end

	function cost(x::Number,y::Number)
		parameters[dims] .= x,y
		return loss(F,parameters,X,hyperparameters)
	end

	function finite_differences(x,y)
		parameters[dims] .= x,y
		return finite_difference_gradient(θ->loss(F,θ,X),parameters)
	end

	plot(size=(600,600), xlim=xlim, ylim=ylim, xlabel="parameters, θ")
	x,y = range(xlim...,length=50), range(ylim...,length=50)
	contour!( x, y, cost, alpha=0.5 )

	x,y = range(xlim...,length=12), range(ylim...,length=12)
	grid = collect(Iterators.product(x[2:end-1],y[2:end-1]))

	xGrid = vcat(map(x->x[1], grid)...)
	yGrid = vcat(map(x->x[2], grid)...)

	differences = map( finite_differences, xGrid, yGrid )
	quiver!( xGrid, yGrid, quiver=(ϵ*step(x)*getindex.(differences,dims[1]),ϵ*step(y)*getindex.(differences,dims[2])),
		color=:darkblue, lw=3 )

	gradients = map( ∇cost, xGrid, yGrid )
	quiver!( xGrid, yGrid, quiver=(ϵ*step(x)*getindex.(gradients,dims[1]),ϵ*step(y)*getindex.(gradients,dims[2])),
		color=:gold, lw=2 )

	errors = filter(x->~isnan(x), map( (x,y) -> acos(x'y/(norm(x)*norm(y)))/π, differences, gradients ))
	plot!([],[],color=:gold, lw=3, label="ForwardDiff", title="Mean Error $(100*round(sum(errors)/length(errors),digits=2))%")
	plot!([],[],color=:darkblue, lw=3, label="Central Differences")
	savefig(name)
end


# Create example point data
x = 8 .* rand(100) .- 4
y = 8 .* rand(100) .- 4
z = sin.(x) + cos.(y)
# Triangular surface plot
trisurf(x, y, z)


## Ziyi Yin converted the Compass velocity model to permeability and porosity

using DrWatson
@quickactivate "CompassTimeLapseCCS"

using Random, PyPlot
using LinearAlgebra, JLD2, MAT, Images, FFTW

n = (637, 82)
d = (25.0f0, 25.0f0)

JLD2.@load "../model/CompassSlice.jld2" v rho d # original Vp in 25m, 6m grid spacing

function find_water_bottom(m::AbstractArray{avDT,2};eps = 1e-4) where {avDT}
  #return the indices of the water bottom of a seismic image
  ### function from JUDI package
  n = size(m)
  idx = zeros(Integer, n[1])
  for j=1:n[1]
      k=1
      while true
          if abs(m[j,k]) > eps
              idx[j] = k
              break
          end
          k += 1
      end
  end
  return idx
end

function down_sample(temp1)
    ## downsample the model to 25m grid spacing w/ low-pass filter
    fft1 = fftshift(fft(temp1))
    mask = zeros(size(fft1))
    taper_length = 8
    half_taper_length = Int(taper_length/2)
    taper = (1f0 .+ sin.((Float32(pi)*(0:taper_length-1))/(taper_length - 1) .- Float32(pi)/2f0))/2f0
    mask[:,171-41:171+41] .= 1f0
    mask[:,171-41-taper_length+1+half_taper_length:171-41+half_taper_length] .= reshape(taper,1,taper_length)
    mask[:,171+41-half_taper_length:171+41+taper_length-1-half_taper_length] .= reshape(reverse(taper),1,taper_length)
    fft2 = mask .* fft1
    temp2 = real.(ifft(ifftshift(fft2)))
    temp2[:,1:35] .= temp1[:,1:35]  # fix water layer
    v = imresize(temp2,n)
    v[v.<=minimum(temp1)] .= minimum(temp1)
    v[v.>=maximum(temp1)] .= maximum(temp1)
    return v
end

v = down_sample(deepcopy(v))

extentx = (n[1]-1)*d[1]
extentz = (n[2]-1)*d[2]

idx_wb = find_water_bottom(v.-minimum(v))[1]

perm = 5f0*(v.-minimum(v))/(maximum(v)-minimum(v))

for i = 1:n[1]
  perm[i,1:idx_wb] .= 1f-4
  unconformaty = findall(v[i,:].>=3500f0)[1]
  perm[i,idx_wb+1:unconformaty] = perm[i,idx_wb+1:unconformaty] .+ 15f0
  perm[i,unconformaty+1:unconformaty+2] = perm[i,unconformaty+1:unconformaty+2]*1f-3
  perm[i,unconformaty+3:end] = perm[i,unconformaty+3:end] .+ 170f0
end

function poro_to_perm(phi)
    # Kozeny-carman equation
  c = (1.527/0.0314)^2f0
  perm = c*phi.^3f0./(1f0.-phi).^2f0
  return perm
end

phi = (perm*(1-0.22)^2/(1.527/0.0314)^2f0).^(1/3f0) # a proxy

JLD2.@save "../model/PermPoro.jld2" perm phi d


module Day13
export day13

TESTINPUT = """6,10
0,14
9,10
0,3
10,4
4,11
6,0
6,12
4,1
0,13
10,12
3,4
3,0
8,4
1,10
2,14
8,10
9,0

fold along y=7
fold along x=5"""

function readInput(inputfile="inputs/input13.txt")
    lines = Iterators.Stateful(eachline(inputfile))

    # parse dots
    inds = Vector{Int64}[]
    for line in lines
        length(line) == 0 && break # empty line separates coordinates from folds
        push!(inds, parse.(Int64, split(line, ",")))
    end

    # parse folds
    folds = []
    for line in lines
        axis = line[12] == 'x' ? 1 : 2
        values = parse(Int64, line[14:end])
        push!(folds, (axis, values)) # +1 to account for Julia arrays
    end
    return inds, folds
end

function day13((indices, folds) = readInput("inputs/input13.txt"); output="inputs/output13.txt")
    part1 = length(_fold!(indices, folds[1]))
    [_fold!(indices, fold) for fold in folds[2:end]]

    # visualize
    maxX = maximum(x->x[1], indices)
    maxY = maximum(x->x[2], indices)
    arr = fill('.', maxX+1, maxY+1)

    arr[[CartesianIndex(ind...) + CartesianIndex(1,1) for ind in indices]] .= '#'
    open(output, "w") do out
        println.(Ref(out), String.(eachcol(arr)))
    end

    return (part1, "See $output")
end

function _fold!(inds, (axis, coordinate))
    for ind in inds
        if ind[axis] > coordinate
            ind[axis] = 2*coordinate - ind[axis]
        end
    end
    unique!(inds)
end

end #module


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Julia 1.4.0
#     language: julia
#     name: julia-1.4
# ---

using Plots
using Statistics
using LinearAlgebra

# $N$ 個の sample $X_1,X_2,\dots, X_N \in \mathbb{R}^n$ が与えられているときにその性質をよく反映する部分空間を求めたい.
# 部分空間を求めるということはその空間を生成する正規直交基底を取れば良い. sample は平均がゼロベクトルになるように正規化されているとする.
#
# $$
# X=
# \begin{bmatrix}
# X_1\ X_2 \dots X_N
# \end{bmatrix}
# $$
#
# で定める行列 $X$ を定義する. これは縦ベクトルであるサンプルを横にN列並べて得られる行列である. $X$ を特異値分解(SVD)することで
#
# $$
#         X=U\Sigma V^\top
# $$
#
# のように記述できたとする. ただし $\Sigma$ は特異値 $\sigma_1,\dots,\sigma_{\mathrm{rank}(\Sigma)}$をの大きい順に並べた対角行列となっているとする.
# 部分空間の次元が $r$ であれば求めたい部分空間は $\sigma_1,\dots,\sigma_r$ に対応する左特異ベクトル $u_1,\dots,u_r$ であることが知られている.

# # Line Fitting

# +

f(x) = 2x + 1
N = 10
x = 2 .* rand(N) .- 1
plot(x, f.(x), label="ground truth")
y = f.(x) + 0.1 * randn(N)
scatter!(x, y, label="sample")
gx = mean(x)
gy = mean(y)

X = hcat(x .- gx, y .- gy)'
U, Σ, Vt = svd(X)
(u11, u21, u12, u22) = U

x = range(-1,1,length=100)
y = @. -u12 * (x - gx) / u22 + gy
plot!(x, y, label="predict")

# -

# # Surface Fitting

# +
using Statistics
using Plots
using LinearAlgebra

f(x, y) = x + y
N = 50

plot(xlim = [-1, 1], ylim = [-1, 1])
x = range(-1.0, 1.0, length = 30)
y = range(-1.0, 1.0, length = 30)
plot!(x, y, f, st = :wireframe, label="ground truth")

x = 2 .* rand(N) .- 1
y = 2 .* rand(N) .- 1
z = f.(x, y) .+ 0.5 .* randn(N)

scatter!(x, y, z, label="sample")
gx = mean(x)
gy = mean(y)
gz = mean(z)

X = hcat(x .- gx, y .- gy, z .- gz)'
U, Σ, Vt = svd(X)
ux, uy, uz = U[:, 3]
x = range(-0.5, 0.5, length = 30)
y = range(-0.5, 0.5, length = 30)
pred(x, y) = -(ux * (x - gx) + uy * (y - gy)) / uz + gz
plot!(x, y, pred, st = :surface, alpha = 0.5, label="predict")


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

# # Assembling stiffness matrix and force vector for 2d heat problem

using FEMBase, HeatTransfer, Test, SparseArrays

X = Dict(
    1 => [0.0,0.0],
    2 => [1.0,0.0],
    3 => [1.0,1.0],
    4 => [0.0,1.0])

T = Dict(
    1 => 0.0,
    2 => 0.0,
    3 => 0.0,
    4 => 0.0)

element = Element(Quad4, [1, 2, 3, 4])
update!(element, "geometry", X)
update!(element, "temperature", 0.0 => T)
update!(element, "thermal conductivity", 6.0)
update!(element, "heat source", 4.0)

problem = Problem(PlaneHeat, "test problem", 1)
add_elements!(problem, [element])
assemble!(problem, 0.0)
K = Matrix(sparse(problem.assembly.K))
f = Vector(sparse(problem.assembly.f)[:])
K_expected = [
               4.0 -1.0 -2.0 -1.0
              -1.0  4.0 -1.0 -2.0
              -2.0 -1.0  4.0 -1.0
              -1.0 -2.0 -1.0  4.0
             ]
f_expected = [1.0; 1.0; 1.0; 1.0]
@test isapprox(K, K_expected)
@test isapprox(f, f_expected)
@test get_unknown_field_name(problem) == "temperature"


using DiffEqBase, DiffEqJump, DataStructures
using Test, Statistics

using StableRNGs
rng = StableRNG(12345)

# using Plots; plotlyjs()
doplot = false

# using BenchmarkTools
# dobenchmark = false

dotestmean   = true
doprintmeans = true

# SSAs to test
SSAalgs = (RDirect(), RSSACR(),Direct(), DirectFW(), FRM(), FRMFW(), SortingDirect(), NRM(), RSSA(), DirectCR(), MNRM())

Nsims        = 32000
tf           = .01
u0           = [200, 100, 150]
expected_avg = 84.876015624999994
reltol       = .01

# MODEL SETUP

# using DiffEqBiological
# rs = @reaction_network dtype begin
#     k1, 2A --> B
#     k2, B --> 2A
#     k3, A + B --> C
#     k4, C --> A + B
#     k5, 3C --> 3A
# end k1 k2 k3 k4

# model using mass action jumps
# ids: A = 1, B = 2, C = 3
reactstoch = [
    [1 => 2],
    [2 => 1],
    [1 => 1, 2 => 1],
    [3 => 1],
    [3 => 3]
]
netstoch = [
    [1 => -2, 2 => 1],
    [1 => 2, 2 => -1],
    [1 => -1, 2 => -1, 3 => 1],
    [1 => 1, 2 => 1, 3 => -1],
    [1 => 3, 3 => -3]
]
rates = [1., 2., .5, .75, .25]
spec_to_dep_jumps = [[1,3],[2,3],[4,5]]
jump_to_dep_specs = [[1,2],[1,2],[1,2,3],[1,2,3],[1,3]]
majumps = MassActionJump(rates, reactstoch, netstoch)

# average number of proteins in a simulation
function runSSAs(jump_prob)
    Psamp = zeros(Int, Nsims)
    for i in 1:Nsims
        sol = solve(jump_prob, SSAStepper())
        Psamp[i] = sol[1,end]
    end
    mean(Psamp)
end

# TESTING:
prob = DiscreteProblem(u0, (0.0, tf), rates)

# plotting one full trajectory
if doplot
    for alg in SSAalgs
        local jump_prob = JumpProblem(prob, alg, majumps, vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs, rng=rng)
        local sol = solve(jump_prob, SSAStepper())
        local plothand = plot(sol, seriestype=:steppost, reuse=false)
        display(plothand)
    end
end

# test the means
if dotestmean
    means = zeros(Float64,length(SSAalgs))
    for (i,alg) in enumerate(SSAalgs)
        local jump_prob = JumpProblem(prob, alg, majumps, save_positions=(false,false), vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs, rng=rng)
        means[i]  = runSSAs(jump_prob)
        relerr = abs(means[i] - expected_avg) / expected_avg
        if doprintmeans
            println("Mean from method: ", typeof(alg), " is = ", means[i], ", rel err = ", relerr)
        end

        # if dobenchmark
        #      @btime (runSSAs($jump_prob);)
        # end

        @test abs(means[i] - expected_avg) < reltol*expected_avg
    end
end

# benchmark performance
# if dobenchmark
#     # exact methods
#     for alg in SSAalgs
#         println("Solving with method: ", typeof(alg), ", using SSAStepper")
#         jump_prob = JumpProblem(prob, alg, majumps, vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs, rng=rng)
#         @btime solve($jump_prob, SSAStepper())
#     end
#     println()
# end

# add a test for passing MassActionJumps individually (tests combining)
if dotestmean
    majump_vec = Vector{MassActionJump{Float64,Vector{Pair{Int,Int}}}}()
    for i = 1:length(rates)
        push!(majump_vec, MassActionJump(rates[i], reactstoch[i], netstoch[i]))
    end
    jset = JumpSet((),(),nothing,majump_vec)
    jump_prob = JumpProblem(prob, Direct(), jset, save_positions=(false,false), vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs, rng=rng)
    meanval = runSSAs(jump_prob)
    relerr = abs(meanval - expected_avg) / expected_avg
    if doprintmeans
        println("Using individual MassActionJumps; Mean from method: ", typeof(Direct()), " is = ", meanval, ", rel err = ", relerr)
    end
    @test abs(meanval - expected_avg) < reltol*expected_avg
end