"""
    struct HomSpace{S<:ElementarySpace, P1<:CompositeSpace{S}, P2<:CompositeSpace{S}}
        codomain::P1
        domain::P2
    end

Represents the linear space of morphisms with codomain of type `P1` and domain of type `P2`.
Note that HomSpace is not a subtype of VectorSpace, i.e. we restrict the latter to denote
certain categories and their objects, and keep HomSpace distinct.
"""
struct HomSpace{S<:ElementarySpace, P1<:CompositeSpace{S}, P2<:CompositeSpace{S}}
    codomain::P1
    domain::P2
end

"""
    TensorSpace{S<:ElementarySpace}

A constant which is the union type of both elementary and product vector space and defined
as `HomSpace{S, ProductSpace{S, N₁}, ProductSpace{S, N₂}}`.
"""
const TensorSpace{S<:ElementarySpace} = Union{S, ProductSpace{S}}

"""
    TensorMapSpace{S<:ElementarySpace, N₁, N₂}

A constant which is an alias name of the HomSpace and defined as
`HomSpace{S, ProductSpace{S, N₁}, ProductSpace{S, N₂}}`.
"""
const TensorMapSpace{S<:ElementarySpace, N₁, N₂} =
    HomSpace{S, ProductSpace{S, N₁}, ProductSpace{S, N₂}}

"""
    →(dom::TensorSpace{S}, codom::TensorSpace{S}) where {S<:ElementarySpace}
    ←(codom::TensorSpace{S}, dom::TensorSpace{S}) where {S<:ElementarySpace}

Convenient constructor of HomSpace instance.
"""
→(dom::TensorSpace{S}, codom::TensorSpace{S}) where {S<:ElementarySpace} =
    HomSpace(ProductSpace(codom), ProductSpace(dom))

←(codom::TensorSpace{S}, dom::TensorSpace{S}) where {S<:ElementarySpace} =
    HomSpace(ProductSpace(codom), ProductSpace(dom))

# The index of a HomSpace is defined as the following: first, index spaces for codomain,
# then, index spaces for domain; dual is taken for each index space of domain, but the order
# are not reversed.
Base.getindex(W::TensorMapSpace{<:IndexSpace, N₁, N₂}, i) where {N₁, N₂} =
    i <= N₁ ? codomain(W)[i] : dual(domain(W)[i-N₁])

"""
    codomain(W::HomSpace)

Return the codomain of a HomSpace.
"""
codomain(W::HomSpace) = W.codomain

"""
    domain(W::HomSpace)

Return the domain of a HomSpace.
"""
domain(W::HomSpace) = W.domain

"""
    spacetype(W::HomSpace) -> Type{ElementarySpace}
    spacetype(::Type{<:HomSpace{S}})

Return the ElementarySpace type associated to a HomSpace instance or type.
"""
spacetype(W::HomSpace) = spacetype(typeof(W))
spacetype(::Type{<:HomSpace{S}}) where S = S

"""
    field(W::HomSpace) -> Field
    field(L::Type{<:HomSpace})

Return the field type over which a HomSpace instance or type is defined.
"""
field(W::HomSpace) = field(typeof(W))
field(L::Type{<:HomSpace}) = field(spacetype(L))

"""
    sectortype(W::HomSpace) -> Type{<:Sector}
    sectortype(L::Type{<:HomSpace)

Return the type of sector over which the HomSpace `W` is defined.
"""
sectortype(W::HomSpace) = sectortype(typeof(W))
sectortype(L::Type{<:HomSpace}) = sectortype(spacetype(L))

"""
    blocksectors(W::HomSpace)

Return an iterator over the different unique coupled sector labels, i.e. the intersection
of the different fusion outputs that can be obtained by fusing the sectors present in the
domain, as well as from the codomain.
"""
function blocksectors(W::HomSpace)
    sectortype(W) === Trivial &&
        return TrivialOrEmptyIterator(dim(domain(W)) == 0 || dim(codomain(W)) == 0)
    return intersect(blocksectors(codomain(W)), blocksectors(domain(W)))
end

"""
    dim(W::HomSpace)

Return the total dimension of a `HomSpace`, i.e. the number of linearly independent
morphisms that can be constructed within this space.
"""
function dim(W::HomSpace)
    d = 0
    for c in blocksectors(W)
        d += blockdim(codomain(W), c) * blockdim(domain(W), c)
    end
    return d
end

"""
    dual(W::HomSpace)

Return the dual of a HomSpace which contains the dual of morphisms in this space.
It corresponds to 180 degree rotation in the graphical representation.
"""
dual(W::HomSpace) = HomSpace(dual(W.domain), dual(W.codomain))

"""
    adjoint(W::HomSpace{<:EuclideanSpace})

Return the adjoint of a HomSpace which contains the dagger of morphisms in this space.
It corresponds to mirror operation and then reversing all arrows in the graphical
representation.
"""
Base.adjoint(W::HomSpace{<:EuclideanSpace}) = HomSpace(W.domain, W.codomain)

Base.hash(W::HomSpace, h::UInt) = hash(domain(W), hash(codomain(W), h))
Base.:(==)(W1::HomSpace, W2::HomSpace) =
    (W1.codomain == W2.codomain) && (W1.domain == W2.domain)

function Base.show(io::IO, W::HomSpace)
    if length(W.codomain) == 1
        print(io, W.codomain[1])
    else
        print(io, W.codomain)
    end
    print(io, " ← ")
    if length(W.domain) == 1
        print(io, W.domain[1])
    else
        print(io, W.domain)
    end
end


__precompile__()

struct Population
    law::Function
    fertilities::Array{Function, 1}
    initial::Float64
end

logistic(f::Number, p::Number) = f * p * (1-p)

mandelbrot(f::Number, p::Number) = f + p^2


# GCM Initial Perturbation
# This file contains helpers and lists currely avaiable options

using Distributions
using Random
using Distributions: Uniform
using Random: rand

abstract type AbstractPerturbation end
struct NoPerturbation <: AbstractPerturbation end
struct DeterministicPerturbation <: AbstractPerturbation end
struct RandomPerturbation <: AbstractPerturbation end

# Helper for parsing `--init-perturbation` command line argument
function parse_perturbation_arg(arg)
    if arg === nothing
        perturbation = nothing
    elseif arg == "deterministic"
        perturbation = DeterministicPerturbation()
    elseif arg == "zero"
        perturbation = NoPerturbation()
    elseif arg == "random"
        perturbation = RandomPerturbation()
    else
        error("unknown perturbation: " * arg)
    end

    return perturbation
end

function init_perturbation(::NoPerturbation, bl, state, aux, coords, t)
    FT = eltype(state)

    u′, v′, w′ = (FT(0), FT(0), FT(0))
    rand_pert = FT(1)

    return u′, v′, w′, rand_pert
end

# Velocity perturbation following
# Ullrich et al. (2016) Dynamical Core Model Intercomparison Project (DCMIP2016) Test Case Document
function init_perturbation(
    ::DeterministicPerturbation,
    bl,
    state,
    aux,
    coords,
    t,
)
    FT = eltype(state)

    # get parameters
    param_set = parameter_set(bl)
    _a::FT = planet_radius(param_set)

    φ = latitude(bl, aux)
    λ = longitude(bl, aux)
    z = altitude(bl, aux)

    # perturbation specific parameters
    z_t::FT = 15e3
    λ_c::FT = π / 9
    φ_c::FT = 2 * π / 9
    d_0::FT = _a / 6
    V_p::FT = 10

    F_z::FT = 1 - 3 * (z / z_t)^2 + 2 * (z / z_t)^3
    if z > z_t
        F_z = FT(0)
    end
    d::FT = _a * acos(sin(φ) * sin(φ_c) + cos(φ) * cos(φ_c) * cos(λ - λ_c))
    c3::FT = cos(π * d / 2 / d_0)^3
    s1::FT = sin(π * d / 2 / d_0)
    if 0 < d < d_0 && d != FT(_a * π)
        u′::FT =
            -16 * V_p / 3 / sqrt(3) *
            F_z *
            c3 *
            s1 *
            (-sin(φ_c) * cos(φ) + cos(φ_c) * sin(φ) * cos(λ - λ_c)) /
            sin(d / _a)
        v′::FT =
            16 * V_p / 3 / sqrt(3) * F_z * c3 * s1 * cos(φ_c) * sin(λ - λ_c) /
            sin(d / _a)
    else
        u′ = FT(0)
        v′ = FT(0)
    end
    w′ = FT(0)
    rand_pert = FT(1)

    return u′, v′, w′, rand_pert
end

function init_perturbation(::RandomPerturbation, bl, state, aux, coords, t)
    FT = eltype(state)
    u′, v′, w′ = (FT(0), FT(0), FT(0))
    rand_pert = FT(1.0 + rand(Uniform(-1e-3, 1e-3)))

    return u′, v′, w′, rand_pert
end


module RayTracing
    export Vec
    export Sphere, Rectangle, Sketch
    export ConstantTexture, ImageTexture, Lambertian, Metal, Dielectric, DiffuseLight
    export Object, Scene, Camera
    export x, y, z, addObject, render

    using StaticArrays
    using LinearAlgebra
    using Images
    using Random
    using Printf

    include("vector.jl")
    include("ray.jl")
    include("surface.jl")
    include("rectangle.jl")
    include("sphere.jl")
    include("sketch.jl")
    include("hitpoint.jl")
    include("texture.jl")
    include("constant_texture.jl")
    include("image_texture.jl")
    include("material.jl")
    include("lambertian.jl")
    include("metal.jl")
    include("dielectric.jl")
    include("diffuse_light.jl")
    include("object.jl")
    include("scene.jl")
    include("camera.jl")

    function color(scene::Scene, ray::Ray, bg_color::Union{AbstractVec, Function}, depth::Int = 0)::AbstractVec
        local succeeded::Bool
        succeeded, hp::Union{Nothing, HitPoint}, material::Union{Nothing, Material} = hit(scene, ray, (min = 0.001, max = Inf))
        if succeeded
            emitted::AbstractVec = emit(material, hp.u, hp.v, hp.normal)
            if depth < 50
                succeeded, attenuation::Union{Nothing, AbstractVec}, scattered::Union{Nothing, Ray} = scatter(material, ray, hp)
                if succeeded
                    return emitted + attenuation .* color(scene, scattered, bg_color, depth + 1)
                else
                    return emitted
                end
            end
        end

        if bg_color isa AbstractVec
            bg_color
        else
            bg_color(normalize(ray.direction))
        end
    end

    function render(scene::Scene, camera::Camera, size::NamedTuple{(:width, :height), Tuple{T, T}} where T <: Integer, samples::Integer, bg_color::Union{AbstractVec, Function})::Array{RGB, 2}
        image::Array{RGB, 2} = zeros(RGB, size.height, size.width)

        function render_row(j::Integer)::Nothing
            for i = 1:size.width
                col::AbstractVec = Vec(0.0, 0.0, 0.0)
                for s = 1:samples
                    u::AbstractFloat = (i - rand()) / size.width
                    v::AbstractFloat = (j - rand()) / size.height
                    ray::Ray = get_ray(camera, u, v)
                    col += color(scene, ray, bg_color)
                end

                col /= samples
                image[size.height + 1 - j, i] = RGB(sqrt(x(col)), sqrt(y(col)), sqrt(z(col)))
            end
        end

        rows::Vector{Integer} = collect(1:size.height)
        shuffle!(rows)

        m::Threads.Mutex = Threads.Mutex()
        count::Integer = 0

        Threads.@threads for row in rows
            render_row(row)
            lock(m)
                count += 1
                @printf("\r%.0f%%", 100.0 * count / size.height)
            unlock(m)
        end

        image
    end

end


using LinearAlgebra
using Statistics
using Zygote
using PaddedViews
using Flux
using Flux: @epochs
using Tullio
using Plots
using Infiltrator

using HDF5
using JLD

#### Parameters
nx, ny = 100, 100 # Size of the grid
Δx, Δy = 1, 1
Δt = 0.01
t₁ = 5

D₀ = 1
tolnl = 1e-4
itMax = 100
damp = 0.1
dτsc   = 1.0/3.0
ϵ     = 1e-4            # small number
cfl  = max(Δx^2,Δy^2)/4.1

A₀ = 1
ρ = 9
g = 9.81
n = 3
p = (Δx, Δy, Δt, t₁, ρ, g, n)  # we add extra parameters for the nonlinear diffusivity

#dataset = 'fake'
dataset = "Argentiere"

if dataset == "fake"

    ### Reference dataset for the heat Equations
    T₀ = [ 250 * exp( - ( (i - nx/2)^2 + (j - ny/2)^2 ) / 300 ) for i in 1:nx, j in 1:ny ];
    T₁ = copy(T₀);

elseif dataset == "Argentiere"

    root_dir = cd(pwd, ".")
    argentiere_f = h5open(joinpath(root_dir, "data/Argentiere_2003-2100_aflow2e-16_50mres_rcp2.6.h5"), "r")

    mutable struct Glacier
        bed::Array{Float64}    # bedrock height
        thick::Array{Float64}  # ice thickness
        vel::Array{Float64}    # surface velocities
        MB::Array{Float64}     # surface mass balance
        lat::Float64
        lon::Float64
    end

    # Fill the Glacier structure with the retrieved data
    argentiere = Glacier(HDF5.read(argentiere_f["bed"])[begin:end-2,:],
                         HDF5.read(argentiere_f["thick_hist"])[begin:end-2,:,2:end],
                         HDF5.read(argentiere_f["vel_hist"])[begin:end-2,:,2:end],
                         HDF5.read(argentiere_f["s_apply_hist"])[begin:end-2,:,2:end],
                         0, 0);

    nx = size(argentiere.bed)[1]
    ny = size(argentiere.bed)[2]

    T₀ = argentiere.thick[:,:,1]
    # B  = copy(argentiere.bed)

end

#######   FUNCTIONS   ############

# Utility functions
@views avg(A) = 0.25 * ( A[1:end-1,1:end-1] .+ A[2:end,1:end-1] .+ A[1:end-1,2:end] .+ A[2:end,2:end] )

@views avg_x(A) = 0.5 .* ( A[1:end-1,:] .+ A[2:end,:] )

@views avg_y(A) = 0.5 .* ( A[:,1:end-1] .+ A[:,2:end] )

### Functions to generate reference dataset to train UDE

function Heat_nonlinear(T, A, p)

    Δx, Δy, Δt, t₁, ρ, g, n = p

    #### NEW CODE TO BREAK  ########

    dTdx = diff(T, dims=1) / Δx
    dTdy = diff(T, dims=2) / Δy
    ∇T = sqrt.(avg_y(dTdx).^2 .+ avg_x(dTdy).^2)

    # D = A * avg(T) .* ∇T # breaking

    Γ = 2 * A * (ρ * g)^n / (n+2)
    D = Γ .* avg(T).^(n + 2) .* ∇T.^(n - 1) # test with SIA
    # D = A

    dTdx_edges = diff(T[:,2:end - 1], dims=1) / Δx
    dTdy_edges = diff(T[2:end - 1,:], dims=2) / Δy

    Fx = -avg_y(D) .* dTdx_edges # breaking
    Fy = -avg_x(D) .* dTdy_edges

    # Fx = -D * dTdx_edges
    # Fy = -D * dTdy_edges

    F = .-(diff(Fx, dims=1) / Δx .+ diff(Fy, dims=2) / Δy)

    # dτ = dτsc * min( 10.0 , 1.0/(1.0/Δt + 1.0/(cfl/(ϵ + D))))

    dτ = dτsc * min.( 10.0 , 1.0./(1.0/Δt .+ 1.0./(cfl./(ϵ .+ avg(D))))) # breaking

    #########################

    return F, dτ

end

# Fake law to create reference dataset and to be learnt by the NN
# fakeA(t) = A₀ * exp(2t)
# fakeA(t) = 0.1 * (2 + t^5 / t₁^4)

# fakeA(t) = (t + t^2) .* 2e-16 # original

# predict_A(UA, t) = (0 .+ UA(t)).*2e-16 # For the base value I've been trying
                              # either the avg value or the initial (t=0) value

predict_A(UA, t) = (2 .+ UA(t)).*1e-16 # SIA

fakeA(t) = (1 + t^1.7)*1e-17

### Heat equation based on a fake A parameter function to compute the diffusivity
function heatflow_nonlinear(T, p, fake, tol=Inf)

    Δx, Δy, Δt, t₁, ρ, g, n = p

    total_iter = 0
    t = 0

    while t < t₁

        iter = 1
        err = 2 * tolnl
        Hold = copy(T)
        dTdt = zeros(nx, ny)
        # err = Inf

        while iter < itMax+1 && tolnl <= err

            Err = copy(T)

            if fake
                A = fakeA(t)  # compute the fake A value involved in the nonlinear diffusivity
            else
                # Compute A with the NN once per time step
                A = predict_A(UA, [t]')[1]  # compute A parameter involved in the diffusivity
            end

            F, dτ = Heat_nonlinear(T, A, p)

            @tullio ResT[i,j] := -(T[i,j] - Hold[i,j])/Δt + F[pad(i-1,1,1),pad(j-1,1,1)]

            dTdt_ = copy(dTdt)
            @tullio dTdt[i,j] := dTdt_[i,j]*damp + ResT[i,j]

            T_ = copy(T)
            # @tullio T[i,j] := max(0.0, T_[i,j] + dτ * dTdt[i,j])
            @tullio T[i,j] := max(0.0, T_[i,j] + dTdt[i,j]*dτ[pad(i-1,1,1),pad(j-1,1,1)])  # breaking

            Zygote.ignore() do
                Err .= Err .- T
                err = maximum(Err)
                # println("error at iter ", iter, ": ", err)

            end

            iter += 1
            total_iter += 1

        end

        # println("t: ", t)
        t += Δt

    end

    # if(!fake)
    #     println("Values of UA in heatflow_nonlinear: ", UA([0., .5, 1.]')) # Simulations here are correct
    # end

    return T

end

# Patch suggested by Michael Abbott needed in order to correctly retrieve gradients
Flux.Optimise.update!(opt, x::AbstractMatrix, Δ::AbstractVector) = Flux.Optimise.update!(opt, x, reshape(Δ, size(x)))

function train(loss, UA, T, p)

    @epochs 1 hybrid_train_NN!(loss, UA, T, p, opt, losses)

    # println("Values of UA in train(): ", UA([0., .5, 1.]'))

    return UA, losses

end

function hybrid_train_NN!(loss, UA, T, p, opt, losses)

    θ = Flux.params(UA)
    # println("Values of UA in hybrid_train BEFORE: ", UA([0., .5, 1.]'))
    loss_UA, back_UA = Zygote.pullback(() -> loss(T, p), θ)
    push!(losses, loss_UA)

    ∇_UA = back_UA(one(loss_UA))

    println("Loss: ", loss_UA)

    for ps in θ
        println("ps: ", ps)
       println("Gradients ∇_UA[ps]: ", ∇_UA[ps])
    end

    # println("θ: ", θ) # parameters are NOT NaNs
    println("Values of predict_A in hybrid_train AFTER: ", predict_A(UA, [1,2,3,4,5]')) # Simulations here are all NaNs

    Flux.Optimise.update!(opt, θ, ∇_UA)

end

function loss_NN(T, p, λ=1)

    T = heatflow_nonlinear(T, p, false)

    Zygote.ignore() do
        # println("Values of UA in loss_NN: ", UA([0., .5, 1.]')) # Simulations here are all NaNs
        display(heatmap(T - T_ref, clim=(0, maximum(T₀)), title="Error"))
    end
    l_cost = sqrt(Flux.Losses.mse(T, T_ref; agg=mean))

    return l_cost
end

#######################

########################################
#####  TRAIN 2D HEAT EQUATION PDE  #####
########################################

T₂ = copy(T₀)
# Reference temperature dataset
T_ref = heatflow_nonlinear(T₂, p, true, 1e-1)

# display(heatmap(T₀ - T_ref, clim=(0, maximum(T₀)), title="T₀"))
# display(heatmap(T_ref, clim=(0, maximum(T₀)), title="T_ref"))

leakyrelu(x, a=0.01) = max(a*x, x)
relu(x) = max(0, x)

UA = Chain(
    Dense(1,10),
    Dense(10,10, leakyrelu, init = Flux.glorot_normal),
    Dense(10,10, leakyrelu, init = Flux.glorot_normal),
    Dense(10,5, leakyrelu, init = Flux.glorot_normal),
    Dense(5,1)
)

# UA = Chain(
#     Dense(1,10),
#     Dense(10,10, leakyrelu, init = Flux.kaiming_normal(gain=5)),
#     Dense(10,10, leakyrelu, init = Flux.kaiming_normal(gain=5)),
#     Dense(10,5, leakyrelu, init = Flux.kaiming_normal(gain=5)),
#     Dense(5,1)
# )

opt = RMSProp(0.001)
losses = []

# Train heat equation UDE
T = copy(T₀)
# UA_trained, losses = train(loss_NN, UA, T, p)

hybrid_train_NN!(loss_NN, UA, T, p, opt, losses)

all_times = LinRange(0, t₁, 50)
# println("UA_trained(all_times')': ",  UA_trained(all_times')')
plot(all_times, predict_A(UA_trained, all_times')', title="Simulated A values by the NN", yaxis="A", xaxis="Time", label="NN")
plot!(fakeA, 0, t₁, label="fake")


using Plots
using SpikingNeuralNetworks
SNN.@load_units

ear = SNN.Rate(;N = 2)
ear.r = zeros(ear.N) # exp(- (0:ear.N-1).^2 / 10^2)
ppc = SNN.Rate(;N = 200)
ear_ppc = SNN.PINningSynapse(ear, ear; σ = 1.5, p = 1.0)
ppc_ppc = SNN.PINningSynapse(ppc, ppc; σ = 1.5, p = 1.0)
P, C = [ppc], [ear_ppc, ppc_ppc]

SNN.monitor(ppc_ppc, [(:g, [1])])

A = 1.3 / 1.5; fr = 1 / 60ms
f(t) = (A /1.0) * sin(1π * fr * t) +  (A / 2.0) * sin(2π * fr * t) +
(A / 6.0) * sin(3π * fr * t) +   (A / 3.0) * sin(4π * fr * t)

ts = 0:0.1ms:1440ms
for (i, t) in enumerate(ts)
    ppc_ppc.f .= [f(t); ppc_ppc.g[2:end]]
    SNN.train!(P, C, 0.1ms, t)
end

SNN.vecplot(ppc_ppc, :g); plot!(f.(ts))


module Day9

export day9

function issumoftwo(window, number)
    for a in window
        if in(number-a, window)
            return true
        end
    end
    false
end

function day9(::Val{:a}, filename, prelen)
    input = Vector{Int}()
    for line in readlines(filename)
        push!(input, parse(Int, line))
    end
    for i in 1:length(input)-prelen
        window = input[i:prelen+i-1]
        new = input[prelen+i]
        if !issumoftwo(window, new)
            return new
        end
    end
end

function findwindow(input, target)
    for i in 1:length(input)-1
        current_sum = input[i]
        for j in i+1:length(input)
            current_sum += input[j]
            if current_sum == target
                return input[i:j]
            end
        end
    end
end

function day9(::Val{:b}, filename, prelen)
    input = Vector{Int}()
    for line in readlines(filename)
        push!(input, parse(Int, line))
    end
    target = day9a(filename, prelen)
    window = findwindow(input, target)
    (min, max) = reduce(window, init = (window[1],window[1])) do (min, max), n
        (n < min ? n : min, n > max ? n : max)
    end
    min + max
end

end


include("../utils.jl")

@testset "Essential manifold" begin
    M = EssentialManifold()
    a = π / 6
    b = π / 6
    c = π / 6
    r1 = [1.0 0.0 0.0; 0.0 cos(a) -sin(a); 0.0 sin(a) cos(a)]
    r2 = [cos(b) 0.0 sin(b); 0.0 1.0 0.0; -sin(b) 0.0 cos(b)]
    r3 = [cos(c) -sin(c) 0.0; sin(c) cos(c) 0.0; 0.0 0.0 1.0]
    nr = [1.0 0.0 0.0; 0.0 0.0 -1.0; 0.0 -1.0 0.0]
    p1 = [r1, r2]
    p2 = [r1, r3]
    p3 = [r2, r2]
    @testset "Essential manifold Basics" begin
        @test M.manifold == Rotations(3)
        @test repr(M) == "EssentialManifold(true)"
        @test manifold_dimension(M) == 5
        np1 = [r1, nr]
        np2 = [nr, nr]
        np3 = [r1, r2, r3]
        @test !is_point(M, r1)
        # first two components of r1 are not rotations
        @test_throws DomainError is_point(M, r1, true)
        @test_throws DomainError is_point(M, np3, true)
        @test is_point(M, p1)
        @test_throws ComponentManifoldError is_point(M, np1, true)
        @test_throws CompositeManifoldError is_point(M, np2, true)
        @test !is_vector(M, p1, 0.0)
        @test_throws DomainError is_vector(
            M,
            p1,
            [0.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0],
            true,
        )
        @test !is_vector(M, np1, [0.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0])
        @test !is_vector(M, p1, p2)
        # projection test
        @test is_vector(M, p1, project(M, p1, log(M, p1, p2)))
        @test is_vector(M, p1, project(M, p1, log(M, p2, p1)))
    end
    @testset "Signed Essential" begin
        test_manifold(
            M,
            [p1, p2, p3],
            test_vector_spaces=true,
            test_project_point=true,
            projection_atol_multiplier=10,
            test_project_tangent=false, # since it includes vert_proj.
            test_musical_isomorphisms=false,
            test_default_vector_transport=true,
            test_representation_size=false,
            test_exp_log=true,
            mid_point12=nothing,
            exp_log_atol_multiplier=4,
            test_inplace=true,
            parallel_transport=true,
        )
    end
    @testset "Unsigned Essential" begin
        test_manifold(
            EssentialManifold(false),
            [p1, p2, p3],
            test_vector_spaces=true,
            test_project_point=true,
            projection_atol_multiplier=10,
            test_project_tangent=false, # since it includes vert_proj.
            test_musical_isomorphisms=false,
            test_default_vector_transport=true,
            test_representation_size=false,
            test_exp_log=true,
            mid_point12=nothing,
            exp_log_atol_multiplier=4,
            parallel_transport=true,
        )
    end
end


"The effective low frequency bulk modulus and density of a material filled with particles"
function effective_medium(medium::Acoustic{T,Dim}, species::Species{T,Dim};
        numberofparticles::Number = Inf
    ) where {T<:AbstractFloat,Dim}

    scale_number_density = one(T) - one(T) / numberofparticles
    φ = scale_number_density * sum(volume_fraction.(species))

    # calculate effective properties
    β = medium.ρ*medium.c^2 # medium bulk modulus
    if abs(β) == zero(T) || abs(prod(s -> s.particle.medium.ρ * s.particle.medium.c^2, species)) == zero(T)
        β_eff = zero(T)
    else
        β_eff = ((one(T)-φ) / β + scale_number_density * sum(s.volume_fraction / (s.particle.medium.ρ * s.particle.medium.c^2) for s in species))^(-one(T))
    end

    ρ_frac = scale_number_density * sum(
        map(species) do s
            if s.particle.medium.ρ == Inf
                - s.volume_fraction / T(Dim - 1)
            else
                s.volume_fraction * (medium.ρ - s.particle.medium.ρ)/(medium.ρ + T(Dim - 1) * s.particle.medium.ρ)
            end
        end
    )

    ρ_eff = medium.ρ*(1 - ρ_frac)/(1 + T(Dim - 1) * ρ_frac)

    return Acoustic(Dim; ρ = ρ_eff, c = sqrt(β_eff/ρ_eff))
end

"""
reflection_coefficient(PlaneSource, Acoustic[, Halfspace = Halfspace(-psource.direction)])

calculates the reflection coefficient from a homogenious halfspace (assumed to direct incidence if not given), which is also the low frequency reflection from a particulate material when using the effective_medium.
"""
reflection_coefficient(ω::T, source::PlaneSource{T}, reflect_medium::Acoustic{T}, halfspace::Halfspace{T} = Halfspace(-psource.direction)) where T<:AbstractFloat = reflection_transmission_coefficients(ω, psource, reflect_medium, halfspace)[1]

"""
reflection_transmission_coefficients(PlaneSource, Acoustic[, Halfspace = Halfspace(-psource.direction)])

calculates the reflection and transmission coefficient from a homogenious halfspace (assumed to direct incidence if not given), which is also the low frequency reflection from a particulate material when using the effective_medium.
"""
function reflection_transmission_coefficients(ω::T, psource::PlaneSource{T,Dim,1,Acoustic{T,Dim}}, reflect_medium::Acoustic{T,Dim}, halfspace::Halfspace{T,Dim} = Halfspace(-psource.direction)) where {T<:AbstractFloat,Dim}

    k_in = ω / psource.medium.c
    k_r = if abs(reflect_medium.c) == zero(T)
        T(Inf) + zero(T) * im
    else
        ω / reflect_medium.c
    end

    # make sure the normal outward pointing
    n = halfspace.normal / norm(halfspace.normal)
    if real(dot(n,psource.direction)) > 0
        n = -n
    end

    θin = transmission_angle(psource.direction, n)

    vtran = transmission_direction(k_r, k_in * psource.direction, n)
    θtran = transmission_angle(vtran, n)

    q = reflect_medium.c*reflect_medium.ρ / (psource.medium.c*psource.medium.ρ)

    Z = dot(n,halfspace.origin)

    R1 = (exp(2im * k_in * Z * cos(θin)) * field(psource)(zeros(T,Dim),ω)) * (q*cos(θin) - cos(θtran)) / (q*cos(θin) + cos(θtran))
    R1 = R1 * exp(im * k_in * Z * cos(θin)) # assuming a reflected wave  = R * exp(i kR.(X-X1)) where kR.X1 =  k_in * Z * cos(θin)

    T1 = exp(im * Z * (k_in * cos(θin) - k_r * cos(θtran))) * field(psource)(zeros(T,Dim),ω) * 2*q*cos(θin) / (q*cos(θin) + cos(θtran))
    T1 = T1 * exp(im * k_r * Z * cos(θtran)) # assuming a transmitter wave  = T * exp(i k1.(X-X1)) where k1.X1 =  k_r * Z * cos(θtran)

    return [R1,T1]
end

"""
    planewave_amplitudes(psource::PlaneSource, reflect_medium::Acoustic, plate::Plate)

Calculates the coefficients, or amplitudes, of the plane waves scattering from and transmitted inside the plate, for the plane wave source `psource` and the `plate`. The function returns `[R, T, P1, P2]` where `R` is the reflection coefficient, `T` is the coefficient of the transmitted wave, and `P1` (`P2`) are the amplitudes of the wave travling forward (backward) inside the plate.
"""
function planewave_amplitudes(ω::T, psource::PlaneSource{T,Dim,1,Acoustic{T,Dim}}, plate_medium::Acoustic{T,Dim}, plate::Plate{T,Dim}) where {T<:AbstractFloat,Dim}

    k0 = ω / psource.medium.c
    k1 = if abs(plate_medium.c) == zero(T)
        T(Inf) + zero(T) * im
    else
        ω / plate_medium.c
    end

    # make the normal outward pointing
    n = plate.normal / norm(plate.normal)
    if real(dot(n,psource.direction)) > 0
        n = -n
    end

    θ0 = transmission_angle(psource, plate)

    v1 = transmission_direction(k1, k0 .* psource.direction, n)
    θ1 = transmission_angle(v1, n)

    C0 = k0 * cos(θ0) / psource.medium.ρ
    C1 = k1 * cos(θ1) / plate_medium.ρ

    Z0 = dot(-n,plate.origin)
    Z1 = Z0 - plate.width / T(2)
    Z2 = Z0 + plate.width / T(2)

    denom = ((C0 + C1)^2 * exp(2im*k1*Z1*cos(θ1)) - (C0 - C1)^2 * exp(2im*k1*Z2*cos(θ1)))

    UR = (C0 - C1)*(C0 + C1) * exp(2im*k0*Z1*cos(θ0)) * (exp(2im*k1*Z1*cos(θ1)) - exp(2im*k1*Z2*cos(θ1))) / denom
    UR = UR * exp(im*k0*Z1*cos(θ0)) # we assume the reflected wave = UR * exp(im*KR.(X-X1)), where KR.X1 = im*k0*Z1*cos(θ0) = k0 .* [sin(θ0),0,-cos(θ0)], and X1 = [0.0,0.0,Z1], so that KR.X1 = im*k0*Z1*cos(θ0).

    UT = 4*C0*C1*exp(im*(k0*(Z1 - Z2)*cos(θ0) + k1*(Z1 + Z2)*cos(θ1))) / denom
    UT = UT * exp(im*k0*Z2*cos(θ0))

    UP1 = 2C0*(C0 + C1) * exp(im*Z1*(k0*cos(θ0) + k1*cos(θ1))) / denom
    UP1 = UP1 * exp(im*k1*Z1*cos(θ1))

    UP2 = - 2*C0*(C0 - C1)*exp(im*(k0*Z1*cos(θ0) + k1*(Z1 + 2*Z2)*cos(θ1))) / denom
    UP2 = UP2 * exp(im*k1*Z2*cos(θ1))

    return [UR, UT, UP1, UP2] * field(psource)(zeros(T,3),ω)
end


function randomPointsNearVertices(V, point, NumPoints, radioIncidencia)

    numVertices = size(V,1);

    V_rand = poly2D.randomPointInPolygon(V, NumPoints)
    flagMat = zeros(NumPoints,numVertices)
    dist = sqrt.(sum( (V_rand - repeat(point,NumPoints,1)).^2, dims=2))
    flagIn = dist .<= radioIncidencia

    V_out = V_rand[flagIn[:] .== true, :]

    return V_out

end


using Test
using SWIG

function simpletest()
  a = 1 ~ unif
  b = 2 ~ a
  sample(b)
end

function testlogpdf()
  x = 1 ~ ω -> Uniform(ω, 0, 1)
  y = 2 ~ ω -> Uniform(ω, 0, 1)
  z = 3 ~ (ω -> Normal(ω, x(ω), 1.0)) <| (x,)
  x_ = 0.2
  y_ = 0.4
  z_ = 0.5
  ω = Ω(Dict((1,) => x_, (2,) => y_, (3,) => z_))
  l = logpdf(rt(x, y, z), ω)
  @test l == Distributions.logpdf(Normal(x_, 1), z_)
end

function testconditioning()
  # model
  x = 1 ~ unif
  y = 2 ~ unif
  z = 3 ~ ω -> normal(ω, x(ω), 1)
  summed = x + y + z
  rt(x, y, z) | observe(normal(ω, summed, 0.1), 3.4)

  # Trace
  x_ = 0.2
  y_ = 0.4
  z_ = 0.5
  ω = Dict(1 => x_, 2 => y_, 3 => z_)
  logpdf(rt, ω)
end

function test()
  uniform(a, b) = ω -> unif(ω) * (b - a) + a

  x = 1 ~ uniform(0, 1)
  y = 2 ~ uniform(0, 1)
  d = 7 ~ y
  z = ω -> (x(ω), x(ω), y(ω), d(ω))
  @show sample(z)
  function a(ω)
    x_ = x(ω)
    d = 3 ~ uniform(0, 4)
    e = 4 ~ uniform(0, 4)
    x_ + d(ω) + e(ω)
  end

  @show sample(rt(x, y, z, a))

  # Conditional Independence
  x = 1 ~ uniform(0, 19)
  measure(ω) = x(ω) + uniform(0, 1)(ω)
  m1 = 3 ~ measure <| (x,)
  m2 = 4 ~ measure <| (x,)
  m3 = 5 ~ measure

  @show sample(rt(m1, m2, m3))

  # Linaer regression
  α = 1 ~ uniform(0, 1)
  β = 2 ~ uniform(0, 1)
  f(x, i) = i ~ (ω -> x*α(ω) + β(ω) + uniform(0.0, 1.0)(ω)) <| (α, β)
  y1 = f(0.3, 3)
  y2 = f(0.3, 4)
  @show sample(rt(y1, y2))

  y1b, y2b = pointwise() do
    f2(x, i) = i ~ (x * α + β + uniform(0.0, 1.0)) <| (α, β)
    y1b = f2(0.3, 3)
    y2b = f2(0.3, 4)
    y1b, y2b
  end

  @show sample(rt(y1, y2, y1b, y2b))

  # A bunch of random variables with a shared parents
  pointwisetest() do
    x = 1 ~ unifQD
    y = x * 100
    function a(ω)
      y(ω) + 5
    end
    function b(ω)
      y(ω) + 10
    end
    c = 2 ~ b <| (y,)
    d = 3 ~ b
    @show sample(rt(y, a, b, c, d))
  end
end