module EnsembleKalmanInversions

using Distributions
using ProgressBars
using Random
using Suppressor
using LinearAlgebra
using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
using EnsembleKalmanProcesses.ParameterDistributionStorage

using ..InverseProblems: n_ensemble, observation_map, forward_map

function lognormal_with_mean_std(mean, std)
    k = std^2 / mean^2 + 1
    μ = log(mean / sqrt(k))
    σ = sqrt(log(k))
    return LogNormal(μ, σ)
end

struct ConstrainedNormal{FT}
    # θ is the original constrained paramter, θ̃ is the unconstrained parameter ~ N(μ, σ)
    # θ = lower_bound + (upper_bound - lower_bound)/（1 + exp(θ̃)）
    μ::FT
    σ::FT
    lower_bound::FT
    upper_bound::FT
end

# Model priors are sometimes constrained; EKI deals with unconstrained, Normal priors.
convert_prior(prior::LogNormal) = Normal(prior.μ, prior.σ)
convert_prior(prior::Normal) = prior
convert_prior(prior::ConstrainedNormal) = Normal(prior.μ, prior.σ)

# Convert parameters to unconstrained for EKI
forward_parameter_transform(::LogNormal, parameter) = log(parameter)
forward_parameter_transform(::Normal, parameter) = parameter
forward_parameter_transform(cn::ConstrainedNormal, parameter) = log((cn.upper_bound - parameter)/(cn.upper_bound - cn.lower_bound))

# Convert parameters from unconstrained (EKI) to constrained
inverse_parameter_transform(::LogNormal, parameter) = exp(parameter)
inverse_parameter_transform(::Normal, parameter) = parameter
inverse_parameter_transform(cn::ConstrainedNormal, parameter) = cn.lower_bound+(cn.upper_bound - cn.lower_bound)/(1 + exp(parameter))

# Convert covariance from unconstrained (EKI) to constrained
inverse_covariance_transform(::Tuple{Vararg{LogNormal}}, parameters, covariance) = Diagonal(exp.(parameters)) * covariance * Diagonal(exp.(parameters))
inverse_covariance_transform(::Tuple{Vararg{Normal}}, parameters, covariance) = covariance
function inverse_covariance_transform(cn::Tuple{Vararg{ConstrainedNormal}}, parameters, covariance)
    upper_bound = [cn[i].upper_bound for i = 1:length(cn)]
    lower_bound = [cn[i].lower_bound for i = 1:length(cn)]
    dT = Diagonal(-(upper_bound - lower_bound) .* exp.(parameters)./(1.0 .+ exp.(parameters)).^2)
    return dT * covariance * dT'
end

mutable struct EnsembleKalmanInversion{I, P, E, M, O, F, S, D}
    inverse_problem :: I
    parameter_distribution :: P
    ensemble_kalman_process :: E
    mapped_observations :: M
    noise_covariance :: O
    inverting_forward_map :: F
    iteration :: Int
    iteration_summaries :: S
    dropped_ensemble_members :: D
end

Base.show(io::IO, eki::EnsembleKalmanInversion) =
    print(io, "EnsembleKalmanInversion", '\n',
              "├── inverse_problem: ", typeof(eki.inverse_problem).name.wrapper, '\n',
              "├── parameter_distribution: ", typeof(eki.parameter_distribution).name.wrapper, '\n',
              "├── ensemble_kalman_process: ", typeof(eki.ensemble_kalman_process), '\n',
              "├── mapped_observations: ", typeof(eki.mapped_observations), '\n',
              "├── noise_covariance: ", typeof(eki.noise_covariance), '\n',
              "├── inverting_forward_map: ", typeof(eki.inverting_forward_map).name.wrapper, '\n',
              "├── iteration: $(eki.iteration)", '\n',
              "├── iteration_summaries: $(eki.iteration_summaries)", '\n',
              "└── dropped_ensemble_members: $(eki.dropped_ensemble_members)", '\n'
              )

construct_noise_covariance(noise_covariance::AbstractMatrix, y) = noise_covariance

function construct_noise_covariance(noise_covariance::Number, y)
    # Independent noise for synthetic observations
    n_obs = length(y)
    return noise_covariance * Matrix(I, n_obs, n_obs)
end

"""
    EnsembleKalmanInversion(inverse_problem; noise_covariance = 1e-2)

Return an object that interfaces with EnsembleKalmanProcesses.jl and uses Ensemble
Kalman Inversion to iteratively "solve" the inverse problem

```math
y = G(θ) + η,
```

for the parameters ``θ``, where ``y`` is a "normalized" vector of observations,
``G(θ)`` is a forward map that predicts the observations, and ``η ∼ N(0, Γ_y)`` is zero-mean
random noise with covariance matrix ``Γ_y`` representing uncertainty in the observations.

By "solve", we mean that the iteration finds  ``θ`` that minimizes the distance
between ``y`` and ``G(θ)``.

(For more details on the Ensemble Kalman Inversion algorithm refer to the
[EnsembleKalmanProcesses.jl Documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/ensemble_kalman_inversion/).)

The "forward map output" `G` can have many interpretations. The specific statistics that `G`
computes have to be selected for each use case to provide a concise summary of the complex
model solution that contains the values that we would most like to match to the corresponding
truth values `y`. For example, in the ocean surface boundary layer context, this summary
could be, for example, a vector of concatenated `u`, `v`, `b`, `e` profiles at all or some time steps
of the CATKE solution.

Arguments
=========
- `inverse_problem :: InverseProblem`: an inverse problem representing the comparison between
                                       synthetic observations generated by Oceananigans and model
                                       predictions also generated by Oceananigans.

- `noise_covariance` (`AbstractMatrix` or `Number`): normalized covariance representing observational
                                                     uncertainty. If `noise_covariance isa Number` then
                                                     it's converted to an identity matrix scaled by
                                                     `noise_covariance`.
"""
function EnsembleKalmanInversion(inverse_problem; noise_covariance = 1e-2)
    free_parameters = inverse_problem.free_parameters
    original_priors = free_parameters.priors

    transformed_priors = [Parameterized(convert_prior(prior)) for prior in original_priors]
    no_constraints = [[no_constraint()] for _ in transformed_priors]
    parameter_distribution = ParameterDistribution(transformed_priors, no_constraints, collect(string.(free_parameters.names)))

    # Seed for pseudo-random number generator for reproducibility
    initial_ensemble = construct_initial_ensemble(parameter_distribution, n_ensemble(inverse_problem); rng_seed = Random.seed!(41))

    # Build EKP-friendly observations "y" and the covariance matrix of observational uncertainty "Γy"
    y = dropdims(observation_map(inverse_problem), dims=2) # length(forward_map_output) column vector
    Γy = construct_noise_covariance(noise_covariance, y)

    # The closure G(θ) maps (N_params, ensemble_size) array to (length(forward_map_output), ensemble_size)
    function G(θ)
        batch_size = size(θ, 2)
        inverted_parameters = [inverse_parameter_transform.(values(original_priors), θ[:, i]) for i in 1:batch_size]
        return forward_map(inverse_problem, inverted_parameters)
    end

    ensemble_kalman_process = EnsembleKalmanProcess(initial_ensemble, y, Γy, Inversion())

    return EnsembleKalmanInversion(inverse_problem, parameter_distribution, ensemble_kalman_process, y, Γy, G, 0, [], Set())
end

"""
    UnscentedKalmanInversion(inverse_problem, prior_mean, prior_cov;
                             noise_covariance = 1e-2, α_reg = 1, update_freq = 0)

Return an object that interfaces with EnsembleKalmanProcesses.jl and uses Unscented
Kalman Inversion to iteratively "solve" the inverse problem

```math
y = G(θ) + η,
```

for the parameters ``θ``, where ``y`` is a "normalized" vector of observations,
``G(θ)`` is a forward map that predicts the observations, and ``η ∼ N(0, Γ_y)`` is zero-mean
random noise with covariance matrix ``Γ_y`` representing uncertainty in the observations.

By "solve", we mean that the iteration finds  ``θ`` that minimizes the distance
between ``y`` and ``G(θ)``.

(For more details on the Unscented Kalman Inversion algorithm refer to the
[EnsembleKalmanProcesses.jl Documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/unscented_kalman_inversion/).)

Arguments
=========

- `inverse_problem :: InverseProblem`: an inverse problem representing the comparison between
                                       synthetic observations generated by Oceananigans and model
                                       predictions also generated by Oceananigans.
- `prior_mean :: Vectorß{Float64}`: prior mean
- `prior_cov :: Array{Float64, 2}`: prior covariance
- `noise_covariance :: Float64`: observation error covariance
- `α_reg :: Float64`: regularization parameter toward the prior mean (0 < `α_reg` ≤ 1);
                      default `α_reg=1` implies no regularization
- `update_freq::IT`: set to 0 when the inverse problem is not identifiable (default), namely the
                     inverse problem has multiple solutions, the covariance matrix will represent
                     only the sensitivity of the parameters, instead of posterior covariance information;
                     set to 1 (or anything > 0) when the inverse problem is identifiable, and
                     the covariance matrix will converge to a good approximation of the
                     posterior covariance with an uninformative prior
"""
function UnscentedKalmanInversion(inverse_problem, prior_mean, prior_cov;
                                  noise_covariance = 1e-2, α_reg = 1, update_freq = 0)
    free_parameters = inverse_problem.free_parameters
    original_priors = free_parameters.priors

    transformed_priors = [Parameterized(convert_prior(prior)) for prior in original_priors]
    no_constraints = [[no_constraint()] for _ in transformed_priors]
    parameter_distribution = ParameterDistribution(transformed_priors, no_constraints, collect(string.(free_parameters.names)))

    # Build EKP-friendly observations "y" and the covariance matrix of observational uncertainty "Γy"
    y = dropdims(observation_map(inverse_problem), dims=2) # length(forward_map_output) column vector
    Γy = construct_noise_covariance(noise_covariance, y)

    # The closure G(θ) maps (N_params, ensemble_size) array to (length(forward_map_output), ensemble_size)
    function G(θ)
        batch_size = size(θ, 2)
        inverted_parameters = [inverse_parameter_transform.(values(original_priors), θ[:, i]) for i in 1:batch_size]
        return forward_map(inverse_problem, inverted_parameters)
    end

    ensemble_kalman_process = EnsembleKalmanProcess(y, Γy, Unscented(prior_mean, prior_cov, α_reg, update_freq))

    return EnsembleKalmanInversion(inverse_problem, parameter_distribution, ensemble_kalman_process, y, Γy, G, 0, [], Set())
end

# return:
# mean : N_iter × Nθ mean matrix
# cov  : N_iter vector of Nθ × Nθ covariance matrix
# std  : N_iter × Nθ standard deviation matrix
# err  : N_iter error array
function UnscentedKalmanInversionPostprocess(eki)

    original_priors = eki.inverse_problem.free_parameters.priors
    θ_mean_raw = hcat(eki.ensemble_kalman_process.process.u_mean...)
    θθ_cov_raw = eki.ensemble_kalman_process.process.uu_cov

    θ_mean = similar(θ_mean_raw)
    θθ_cov= similar(θθ_cov_raw)
    θθ_std_arr = similar(θ_mean_raw)
    for i = 1:size(θ_mean, 2)
        θ_mean[:, i] = inverse_parameter_transform.(values(original_priors), θ_mean_raw[:, i])
        θθ_cov[i] = inverse_covariance_transform(values(original_priors), θ_mean_raw[:, i], θθ_cov_raw[i])

        for j in 1:size(θ_mean, 1)
            θθ_std_arr[j, i] = sqrt(θθ_cov[i][j, j])
        end
    end

    return θ_mean, θθ_cov, θθ_std_arr, eki.ensemble_kalman_process.err
end

struct IterationSummary{P, E}
    parameters :: P
    mean_square_errors :: E
end

function IterationSummary(parameters, forward_map, observations)
    N_observations, N_ensemble = size(forward_map)

    mean_square_errors = [
        mapreduce((x, y) -> (x - y)^2, +, observations, view(forward_map, m, :)) / N_observations
        for m in 1:N_ensemble
    ]

    return IterationSummary(parameters, mean_square_errors)
end

function drop_ensemble_member!(eki, member)
    parameter_ensemble = eki.iteration_summary[end].parameters
    new_parameter_ensemble = vcat(parameter_ensemble[1:member-1], parameter_ensemble[member+1:end])
    new_ensemble_kalman_process = EnsembleKalmanProcess(new_parameter_ensemble, eki.mapped_observations, eki.noise_covariance, Inversion())

    push!(eki.dropped_ensemble_members, memeber)
    eki.ensemble_kalman_process = new_ensemble_kalman_process

    return nothing
end

"""
    iterate!(eki::EnsembleKalmanInversion; iterations=1)

Iterate the ensemble Kalman inversion problem `eki` forward by `iterations`.
"""
function iterate!(eki::EnsembleKalmanInversion; iterations = 1)

    first_iteration = eki.iteration + 1
    final_iteration = eki.iteration + 1 + iterations

    @suppress begin
        for iter = ProgressBar(first_iteration:final_iteration)
            θ = get_u_final(eki.ensemble_kalman_process) # (N_params, ensemble_size) array
            G = eki.inverting_forward_map(θ) # (len(G), ensemble_size)

            # Save the parameter values and mean square error between forward map
            # and observations at the current iteration
            summary = IterationSummary(θ, G, eki.mapped_observations)

            @show summary.mean_square_errors

            eki.iteration = iter
            push!(eki.iteration_summaries, summary)

            update_ensemble!(eki.ensemble_kalman_process, G)
        end
    end

    #=
    # All unconstrained
    params = mean(get_u_final(ensemble_kalman_process), dims=2) # ensemble mean

    # losses = [norm([G([mean(get_u(ensemble_kalman_process, i), dims=2)...])...]  - y) for i in 1:n_iterations] # particle loss
    mean_vars = [diag(cov(get_u(ensemble_kalman_process, i), dims=2)) for i in 1:iterations] # ensemble variance for each parameter

    original_priors = eki.inverse_problem.free_parameters.priors
    params = inverse_parameter_transform.(values(original_priors), params)

    # return params, mean_vars, mean_us_constrained
    mean_us_eki = [[mean(get_u(ensemble_kalman_process, i), dims=2)...] for i in 1:iterations]
    mean_us_constrained = [inverse_parameter_transform.(values(original_priors), mean_u) for mean_u in mean_us_eki]
    return [params...], mean_vars, mean_us_constrained
    =#

    return nothing
end

end # module


using CompScienceMeshes
using BEAST

tetrs = CompScienceMeshes.tetmeshsphere(1.0, 0.10)
# tetrs = CompScienceMeshes.tetmeshsphere(1.0, 0.35)
@show numcells(tetrs)

bndry = boundary(tetrs)
edges = skeleton(tetrs, 1)

bndry_edges = [sort(c) for c in cells(skeleton(bndry, 1))]
function is_interior(edge)
    !(sort(edge) in bndry_edges)
end

interior_edges = submesh(is_interior, edges)
@assert numcells(interior_edges) + numcells(skeleton(bndry,1)) == numcells(edges)

X = BEAST.nedelecc3d(tetrs, interior_edges)
@assert numfunctions(X) == numcells(interior_edges)

Id = BEAST.Identity()
Y = curl(X)
A1 = assemble(Id, Y, Y)
A2 = assemble(Id, X, X)
A = A1 - A2

using LinearAlgebra
f = BEAST.ScalarTrace(x -> point(1,0,0) * exp(-norm(x)^2/4))
b = assemble(f, X)

u = A \ b

using Plots

pts = [point(0,t,0) for t in range(-2,2,length=200)];
vals = BEAST.grideval(pts, u, X)
vals2 = BEAST.grideval(pts, u, Y)
vals2 = [point(0,1,0) × val for val in vals2]
plot(getindex.(real(vals),1), m=2)
plot(norm.(vals2), m=2)

# Inspect the value of a single basis function
p = 120
tet = chart(tetrs, cells(tetrs)[p])
nbd = neighborhood(tet, [0.5, 0.5, 0.0])
edg = simplex(tet.vertices[1], tet.vertices[2])
tgt = normalize(edg.vertices[2] - edg.vertices[1])

vals = refspace(X)(nbd)
@assert dot(tgt,vals[1].value) ≈ 1/volume(edg) atol=1e-8

# check also the value of the raviart-thomas 3d in the center
# of their defining face
Z = BEAST.nedelecd3d(tetrs)
@assert numfunctions(Z) == numcells(skeleton(tetrs,2))

fce = simplex(tet.vertices[2], tet.vertices[3], tet.vertices[4])
nbd_rt = neighborhood(tet, [0.0, 1/3, 1/3])
vals_rt = refspace(Z)(nbd_rt)

# carttobary(fce, cartesian(nbd_rt))

#
dot(vals_rt[1].value, normal(fce)) * volume(fce)
o, x, y, z = euclidianbasis(3)
@assert volume(simplex(x,y,z,o)) * 6 ≈ 1

G = boundary(tetrs)
Z = BEAST.ttrace(curl(X), G)

fcr, geo = facecurrents(u, Z)
import Plotly
Plotly.plot(patch(geo, norm.(fcr)))

const CSM = CompScienceMeshes
hemi = submesh(tet -> cartesian(CSM.center(chart(tetrs,tet)))[3] < 0, tetrs)
# error()
bnd_hemi = boundary(hemi)

Xhemi = BEAST.restrict(X, hemi)
tXhemi = BEAST.ttrace(Xhemi, bnd_hemi)

fcr, geo = facecurrents(u, tXhemi)
Plotly.plot(patch(geo, norm.(fcr)))


using FractionalTransforms
using Test

@testset "FractionalTransforms.jl" begin
    include("frft.jl")
    include("frst.jl")
    include("frct.jl")
end


# https://wickedengine.net/2020/04/26/capsule-collision-detection/

struct CapsuleContact{T}
    position::Vector{T}
    orientation::Quaternion{T}
    height::T
    radius::T
end

p1 = [0.0, 0.0, 0.0]
# o1 = one(Quaternion)
o1 = Quaternion(RotX(0.6) * RotY(0.0 * π))
h1 = 1.0
r1 = 0.1
cap1 = CapsuleContact(p1, o1, h1, r1)

p2 = [1.0, 0.0, 0.0]
o2 = one(Quaternion)
o2 = Quaternion(RotZ(0.0) * RotY(-0.25 * π))
h2 = 1.0
r2 = 0.1
cap2 = CapsuleContact(p2, o2, h2, r2)

function get_tip_base(cap::CapsuleContact)
    k = [0.0, 0.0, 0.5 * cap.height + cap.radius]
    tip = cap.position + vector_rotate(k, cap.orientation)
    base = cap.position + vector_rotate(-k, cap.orientation)
    return Array(tip), Array(base)
end

t1, b1 = get_tip_base(cap1)
t2, b2 = get_tip_base(cap2)

function closest_point_on_line(pa::Vector{T}, pb::Vector{T}, p::Vector{T}) where T
    pab = pb - pa
    t = dot(p - pa, pab) / dot(pab, pab)
    return pa + min(max(t, 0.0), 1.0) * pab
end

function closest_segment_points_capsule_capsule(capa::CapsuleContact, capb::CapsuleContact)
    # capsule a
    a_tip, a_base = get_tip_base(capa)
    a_radius = capa.radius

    # capsule b
    b_tip, b_base = get_tip_base(capb)
    b_radius = capb.radius

    # capsule a
    a_normal = a_tip - a_base
    a_normal ./= norm(a_normal)
    a_line_end_offset = a_normal * a_radius
    a_A = a_base + a_line_end_offset
    a_B = a_tip - a_line_end_offset

    # capsule b
    b_normal = b_tip - b_base
    b_normal ./= norm(b_normal)
    b_line_end_offset = b_normal * b_radius
    b_A = b_base + b_line_end_offset
    b_B = b_tip - b_line_end_offset

    # vectors between line endpoints
    v0 = b_A - a_A
    v1 = b_B - a_A
    v2 = b_A - a_B
    v3 = b_B - a_B

    # squared distances
    d0 = dot(v0, v0)
    d1 = dot(v1, v1)
    d2 = dot(v2, v2)
    d3 = dot(v3, v3)

    # best potential endpoint on capsule a
    if (d2 < d0 || d2 < d1 || d3 < d0 || d3 < d1)
        bestA = a_B
    else
        bestA = a_A
    end

    # select point on capsule B line segment nearest best potential endpoint on A capsule
    bestB = closest_point_on_line(b_A, b_B, bestA)

    # same for capsule A segment
    bestA = closest_point_on_line(a_A, a_B, bestB)

    return bestA, bestB
    # # direction
    # direction = bestB - bestA
    # direction_magnitude = norm(direction)
    # direction ./= direction_magnitude

    # c1 = bestA + direction * a_radius
    # c2 = bestB - direction * b_radius
    # # distance = direction_magnitude - (a_radius + b_radius)
    # return c1, c2
end

function kinematics_capsule_capsule(capa::CapsuleContact, capb::CapsuleContact)
    bestA, bestB = closest_segment_points_capsule_capsule(capa, capb)

    kA = vector_rotate(bestA - capa.position, inv(capa.orientation))
    kB = vector_rotate(bestB - capb.position, inv(capb.orientation))

    return kA, kB
end

ka, kb = kinematics_capsule_capsule(cap1, cap2)

function closest_points_capsule_capsule(capa::CapsuleContact, capb::CapsuleContact)
    @show "hi"
    bestA, bestB = closest_segment_points_capsule_capsule(capa, capb)

    direction = bestB - bestA
    direction_magnitude = norm(direction)
    direction ./= direction_magnitude

    c1 = bestA + direction * capa.radius
    c2 = bestB - direction * capb.radius
    # distance = direction_magnitude - (a_radius + b_radius)
    return c1, c2
end

function direction_capsule_capsule(capa::CapsuleContact, capb::CapsuleContact)
    ca, cb = closest_segment_points_capsule_capsule(capa, capb)

    direction = cb - ca
    direction_magnitude = norm(direction)

    if direction_magnitude > 0.0
        direction ./= direction_magnitude
    else
        direction = [1.0, 0.0, 0.0]
    end

    return direction
end

function sdf_capsule_capsule(capa::CapsuleContact, capb::CapsuleContact)
    ca, cb = closest_points_capsule_capsule(capa, capb)

    direction = cb - ca
    direction_magnitude = norm(direction)

    return direction_magnitude
end

c1, c2 = closest_points_capsule_capsule(cap1, cap2)

# dir = direction_capsule_capsule(cap1, cap2)

# dist = distance_capsule_capsule(cap1, cap2)

# pa = [0.0, 0.0, 0.0]
# pb = [1.0, 1.0, 0.0]
# p = [1.0, 0.5, 0.0]

# closest_point_on_line(pa, pb, p)

# # capsule a
# a_tip = [1.0, 1.0, 0.0]
# a_base = [0.0, 0.0, 0.0]
# a_radius = 0.01

# # capsule b
# b_tip = [2.0, 1.0, 0.0]
# b_base = [2.0, 0.0, 0.0]
# b_radius = 0.01

# # capsule a
# a_normal = a_tip - a_base
# a_normal ./= norm(a_normal)
# a_line_end_offset = a_normal * a_radius
# a_A = a_base + a_line_end_offset
# a_B = a_tip - a_line_end_offset

# # capsule b
# b_normal = b_tip - b_base
# b_normal ./= norm(b_normal)
# b_line_end_offset = b_normal * b_radius
# b_A = b_base + b_line_end_offset
# b_B = b_tip - b_line_end_offset

# # vectors between line endpoints
# v0 = b_A - a_A
# v1 = b_B - a_A
# v2 = b_A - a_B
# v3 = b_B - a_B

# # squared distances
# d0 = dot(v0, v0)
# d1 = dot(v1, v1)
# d2 = dot(v2, v2)
# d3 = dot(v3, v3)

# # best potential endpoint on capsule a
# if (d2 < d0 || d2 < d1 || d3 < d0 || d3 < d1)
#     bestA = a_B
# else
#     bestA = a_A
# end
# # select point on capsule B line segment nearest best potential endpoint on A capsule
# bestB = closest_point_on_line(b_A, b_B, bestA)

# # same for capsule A segment
# bestA = closest_point_on_line(a_A, a_B, bestB)

# # direction
# direction = bestB - bestA
# direction_magnitude = norm(direction)
# direction ./= direction_magnitude

# distance = direction_magnitude - (a_radius + b_radius)

# visualize
vis = Visualizer()
open(vis)

# capsule a
cyl1 = GeometryBasics.Cylinder(Point(0.0, 0.0, -0.5 * cap1.height), Point(0.0, 0.0, 0.5 * cap1.height), cap1.radius)
sph1 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), cap1.radius)
color1 = Colors.RGBA(0.7, 0.7, 0.7, 0.5);
setobject!(vis[:cap1][:cyl], cyl1, MeshPhongMaterial(color=color1))
setobject!(vis[:cap1][:base], sph1, MeshPhongMaterial(color=color1))
settransform!(vis[:cap1][:base], Translation(0.0, 0.0, -0.5 * cap1.height))
setobject!(vis[:cap1][:tip], sph1, MeshPhongMaterial(color=color1))
settransform!(vis[:cap1][:tip], Translation(0.0, 0.0, 0.5 * cap1.height))

# capsule b
cyl2 = GeometryBasics.Cylinder(Point(0.0, 0.0, -0.5 * cap2.height), Point(0.0, 0.0, 0.5 * cap2.height), cap2.radius)
sph2 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), cap2.radius)
color2 = Colors.RGBA(0.7, 0.7, 0.7, 0.5);
setobject!(vis[:cap2][:cyl], cyl2, MeshPhongMaterial(color=color2))
setobject!(vis[:cap2][:base], sph2, MeshPhongMaterial(color=color2))
settransform!(vis[:cap2][:base], Translation(0.0, 0.0, -0.5 * cap2.height))
setobject!(vis[:cap2][:tip], sph2, MeshPhongMaterial(color=color2))
settransform!(vis[:cap2][:tip], Translation(0.0, 0.0, 0.5 * cap2.height))

# set configuration
settransform!(vis[:cap1], compose(Translation(cap1.position), LinearMap(cap1.orientation)))
settransform!(vis[:cap2], compose(Translation(cap2.position), LinearMap(cap2.orientation)))

# closest points
color_cp = Colors.RGBA(0.0, 0.0, 0.0, 1.0);
cs1 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), 0.025)
cs2 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), 0.025)
setobject!(vis[:cs1], cs1, MeshPhongMaterial(color=color_cp))
setobject!(vis[:cs2], cs2, MeshPhongMaterial(color=color_cp))
settransform!(vis[:cs1], Translation(c1))
settransform!(vis[:cs2], Translation(c2))


function testA5b()

    answer = matread("matlabtestdata/testA5b.mat")

    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 10 # Hz
    M = Magnetization(1.0, 0, 0)
    dt = 1 # ms
    time = 0:dt:1000 # ms
    spin = Spin(M, 1, T1, T2, Δf, Position(0,0,0))
    mag = zeros(3,length(time))
    mag[:,1] = Vector(M)
    A = FreePrecessionMatrix()
    B = Magnetization()
    BtoM = Magnetization()
    freeprecess!(A, B, spin, dt)
    for t = 2:length(time)
        applydynamics!(spin, BtoM, A, B)
        mag[:,t] = Vector(spin.M)
    end

    return mag ≈ answer["M"]

end

function testF1a()

    answer = matread("matlabtestdata/testF1a.mat")

    t = 2.3 # ms
    α = π/4 # rad
    θ = 0 # rad
    Δf = -500:5:500 # Hz
    T1 = 600 # ms
    T2 = 100 # ms
    spins = map(Δf -> Spin(1, T1, T2, Δf), Δf)
    rf = InstantaneousRF(α, θ)
    Ae = ExcitationMatrix()
    Af = FreePrecessionMatrix()
    Bf = Magnetization()
    BtoM = Magnetization()
    for f = 1:length(Δf)
        excite!(Ae, spins[f], rf)
        freeprecess!(Af, Bf, spins[f], t)
        applydynamics!(spins[f], BtoM, Ae)
        applydynamics!(spins[f], BtoM, Af, Bf)
        applydynamics!(spins[f], BtoM, Ae)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF1b()

    answer = matread("matlabtestdata/testF1b.mat")

    t = 2.3 # ms
    α = π/4 # rad
    θ = 0 # rad
    Δf = 0 # Hz
    T1 = 600 # ms
    T2 = 100 # ms
    grad = Gradient(0.1, 0, 0) # G/cm
    xpos = -2:0.01:2 # cm
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> Spin(1, T1, T2, Δf, pos), pos)
    rf = InstantaneousRF(α, θ)
    Ae = ExcitationMatrix()
    Af = FreePrecessionMatrix()
    Bf = Magnetization()
    BtoM = Magnetization()
    for i = 1:length(xpos)
        excite!(Ae, spins[i], rf)
        freeprecess!(Af, Bf, spins[i], t, grad)
        applydynamics!(spins[i], BtoM, Ae)
        applydynamics!(spins[i], BtoM, Af, Bf)
        applydynamics!(spins[i], BtoM, Ae)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF1c()

    answer = matread("matlabtestdata/testF1c.mat")

    t = 2.3 # ms
    α = π/4 # rad
    θ = 0 # rad
    Δf = 0 # Hz
    T1 = 600 # ms
    T2 = 100 # ms
    grad = Gradient(0.1, 0, 0) # G/cm
    grad2 = Gradient(-0.05, 0, 0) # G/cm
    xpos = -2:0.01:2 # cm
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> Spin(1, T1, T2, Δf, pos), pos)
    rf = InstantaneousRF(α, θ)
    Ae = ExcitationMatrix()
    Af = FreePrecessionMatrix()
    Bf = Magnetization()
    Ar = FreePrecessionMatrix()
    Br = Magnetization()
    BtoM = Magnetization()
    for i = 1:length(xpos)
        excite!(Ae, spins[i], rf)
        freeprecess!(Af, Bf, spins[i], t, grad)
        freeprecess!(Ar, Br, spins[i], t, grad2)
        applydynamics!(spins[i], BtoM, Ae)
        applydynamics!(spins[i], BtoM, Af, Bf)
        applydynamics!(spins[i], BtoM, Ae)
        applydynamics!(spins[i], BtoM, Ar, Br)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF2c()

    answer = matread("matlabtestdata/testF2c.mat")

    dt = 0.04 # ms
    t = 0:dt:6 # ms
    rf = 0.05 * sinc.(t .- 3) # G
    Δf = -1000:20:1000 # Hz
    spins = map(Δf -> Spin(1, 600, 100, Δf), Δf)
    A = BlochMatrix()
    B = Magnetization()
    BtoM = Magnetization()
    rf = RF(rf, dt)
    map(spins) do spin
        excite!(A, B, spin, rf)
        applydynamics!(spin, BtoM, A, B)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF3a()

    answer = matread("matlabtestdata/testF3a.mat")

    dt = 0.1 # ms
    t = 0:dt:6 # ms
    T = length(t)
    xpos = -2:0.01:2 # cm
    rf = 0.05 * sinc.(t .- 3) # G
    grad = Gradient(0.1, 0, 0)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 0 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> Spin(1, T1, T2, Δf, pos), pos)
    rf = RF(rf, dt, grad)
    A = BlochMatrix()
    B = Magnetization()
    BtoM = Magnetization()
    for i = 1:length(xpos)
        excite!(A, B, spins[i], rf)
        applydynamics!(spins[i], BtoM, A, B)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF3c()

    answer = matread("matlabtestdata/testF3c.mat")

    dt = 0.05 # ms
    t = 0.1:dt:6 # ms
    T = length(t)
    xpos = -2:0.1:2 # cm
    rf = 0.05 * sinc.(t .- 3) # G
    grad = [Gradient(0.1, 0, 0) for i = 1:T]
    grad2 = Gradient(-0.052, 0, 0)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 0 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> Spin(1, T1, T2, Δf, pos), pos)
    rf = RF(rf, dt, grad)
    Ae = BlochMatrix()
    Be = Magnetization()
    Af = FreePrecessionMatrix()
    Bf = Magnetization()
    BtoM = Magnetization()
    for i = 1:length(xpos)
        excite!(Ae, Be, spins[i], rf)
        freeprecess!(Af, Bf, spins[i], t[end] - dt, grad2)
        applydynamics!(spins[i], BtoM, Ae, Be)
        applydynamics!(spins[i], BtoM, Af, Bf)
    end
    sig = map(signal, spins)
    Mz = map(spin -> spin.M.z, spins)

    @test sig ≈ vec(answer["sig"])
    return Mz ≈ vec(answer["mz"])

end

function testF3d()

    answer = matread("matlabtestdata/testF3d.mat")

    dt = 0.05 # ms
    t = 0.1:dt:6 # ms
    T = length(t)
    xpos = -2:0.1:2 # cm
    rf = 0.05 * sinc.(t .- 3) # G
    grad = [Gradient(0.1, 0, 0) for i = 1:T]
    grad2 = Gradient(-0.052, 0, 0)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 100 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> Spin(1, T1, T2, Δf, pos), pos)
    rf = RF(rf, dt, grad)
    Ae = BlochMatrix()
    Be = Magnetization()
    Af = FreePrecessionMatrix()
    Bf = Magnetization()
    BtoM = Magnetization()
    map(spins) do spin
        excite!(Ae, Be, spin, rf)
        freeprecess!(Af, Bf, spin, t[end] - dt, grad2)
        applydynamics!(spin, BtoM, Ae, Be)
        applydynamics!(spin, BtoM, Af, Bf)
    end
    sig = map(signal, spins)
    Mz = map(spin -> spin.M.z, spins)

    @test sig ≈ vec(answer["sig"])
    return Mz ≈ vec(answer["mz"])

end

function testF3f()

    answer = matread("matlabtestdata/testF3f.mat")

    dt = 0.05 # ms
    t = 0.1:dt:6 # ms
    T = length(t)
    xpos = -5:0.1:5 # cm
    rf = 0.05 * sinc.(t .- 3) .* (exp.(im * 2π * 900 * t/1000) + exp.(-im * 2π * 900 * t/1000)) # G
    grad = [Gradient(0.1, 0, 0) for i = 1:T]
    grad2 = Gradient(-0.052, 0, 0)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 0 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> Spin(1, T1, T2, Δf, pos), pos)
    rf = RF(rf, dt, grad)
    Ae = BlochMatrix()
    Be = Magnetization()
    Af = FreePrecessionMatrix()
    Bf = Magnetization()
    BtoM = Magnetization()
    map(spins) do spin
        excite!(Ae, Be, spin, rf)
        freeprecess!(Af, Bf, spin, t[end] - dt, grad2)
        applydynamics!(spin, BtoM, Ae, Be)
        applydynamics!(spin, BtoM, Af, Bf)
    end
    sig = map(signal, spins)
    Mz = map(spin -> spin.M.z, spins)

    @test sig ≈ vec(answer["sig"])
    return Mz ≈ vec(answer["mz"])

end

function testA5bMC()

    answer = matread("matlabtestdata/testA5b.mat")

    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 10 # Hz
    M = MagnetizationMC((1.0, 0, 0), (0, 0, 0))
    dt = 1 # ms
    time = 0:dt:1000 # ms
    spin = SpinMC(M, 1, [1, 0], [T1, T1], [T2, T2], [Δf, Δf], [Inf, Inf])
    A = BlochMcConnellMatrix(2)
    B = MagnetizationMC(2)
    BtoM = MagnetizationMC(2)
    mag = zeros(3,length(time))
    mag[:,1] = Vector(M[1])
    for t = 2:length(time)
        freeprecess!(A, B, spin, dt)
        applydynamics!(spin, BtoM, A, B)
        mag[:,t] = Vector(spin.M[1])
    end

    return mag ≈ answer["M"]

end

function testF1cMC()

    answer = matread("matlabtestdata/testF1c.mat")

    t = 2.3 # ms
    α = π/4 # rad
    θ = 0 # rad
    Δf = 0 # Hz
    T1 = 600 # ms
    T2 = 100 # ms
    grad = Gradient(0.1, 0, 0) # G/cm
    grad2 = Gradient(-0.05, 0, 0) # G/cm
    xpos = -2:0.01:2 # cm
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> SpinMC(1, [1, 0], [T1, T1], [T2, T2], [Δf, Δf], [Inf, Inf], pos), pos)
    rf = InstantaneousRF(α)
    Ae = ExcitationMatrix()
    Af = BlochMcConnellMatrix(2)
    Bf = MagnetizationMC(2)
    Ar = BlochMcConnellMatrix(2)
    Br = MagnetizationMC(2)
    BtoM = MagnetizationMC(2)
    map(spins) do spin
        excite!(Ae, spin, rf)
        freeprecess!(Af, Bf, spin, t, grad)
        freeprecess!(Ar, Br, spin, t, grad2)
        applydynamics!(spin, BtoM, Ae)
        applydynamics!(spin, BtoM, Af, Bf)
        applydynamics!(spin, BtoM, Ae)
        applydynamics!(spin, BtoM, Ar, Br)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF2cMC()

    answer = matread("matlabtestdata/testF2c.mat")

    dt = 0.04 # ms
    t = 0:dt:6 # ms
    rf = 0.05 * sinc.(t .- 3) # G
    Δf = -1000:20:1000 # Hz
    grad = [Gradient(0, 0, 0) for r in rf]
    spins = map(Δf -> SpinMC(1, [1, 0], [600, 10], [100, 10], [Δf, 0], [Inf, Inf]), Δf)
    rf = RF(rf, dt, grad)
    A = BlochMcConnellMatrix(2)
    B = MagnetizationMC(2)
    BtoM = MagnetizationMC(2)
    map(spins) do spin
        excite!(A, B, spin, rf)
        applydynamics!(spin, BtoM, A, B)
    end
    sig = map(signal, spins)

    return sig ≈ vec(answer["sig"])

end

function testF3cMC()

    answer = matread("matlabtestdata/testF3c.mat")

    dt = 0.05 # ms
    t = 0.1:dt:6 # ms
    T = length(t)
    xpos = -2:0.1:2 # cm
    rf = 0.05 * sinc.(t .- 3) # G
    grad1 = Gradient(0.1, 0, 0)
    grad2 = Gradient(-0.052, 0, 0)
    rf = RF(rf, dt, grad1)
    rf0 = RF(zeros(length(rf)), dt, grad2)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 0 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> SpinMC(1, [1, 0], [T1, T1], [T2, T2], [Δf, Δf], [Inf, Inf], pos), pos)
    A1 = BlochMcConnellMatrix(2)
    B1 = MagnetizationMC(2)
    A2 = BlochMcConnellMatrix(2)
    B2 = MagnetizationMC(2)
    BtoM = MagnetizationMC(2)
    map(spins) do spin
        excite!(A1, B1, spin, rf)
        excite!(A2, B2, spin, rf0)
        applydynamics!(spin, BtoM, A1, B1)
        applydynamics!(spin, BtoM, A2, B2)
    end
    sig = map(signal, spins)
    Mz = map(spin -> spin.M[1].z, spins)

    @test sig ≈ vec(answer["sig"])
    return Mz ≈ vec(answer["mz"])

end

function testF3dMC()

    answer = matread("matlabtestdata/testF3d.mat")

    dt = 0.05 # ms
    t = 0.1:dt:6 # ms
    T = length(t)
    xpos = -2:0.1:2 # cm
    rf = 0.05 * sinc.(t .- 3) # G
    grad1 = Gradient(0.1, 0, 0)
    grad2 = Gradient(-0.052, 0, 0)
    rf = RF(rf, dt, grad1)
    rf0 = RF(zeros(length(rf)), dt, grad2)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 100 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> SpinMC(1, [1, 0], [T1, T1], [T2, T2], [Δf, Δf], [Inf, Inf], pos), pos)
    A1 = BlochMcConnellMatrix(2)
    B1 = MagnetizationMC(2)
    A2 = BlochMcConnellMatrix(2)
    B2 = MagnetizationMC(2)
    BtoM = MagnetizationMC(2)
    map(spins) do spin
        excite!(A1, B1, spin, rf)
        excite!(A2, B2, spin, rf0)
        applydynamics!(spin, BtoM, A1, B1)
        applydynamics!(spin, BtoM, A2, B2)
    end
    sig = map(signal, spins)
    Mz = map(spin -> spin.M[1].z, spins)

    @test sig ≈ vec(answer["sig"])
    return Mz ≈ vec(answer["mz"])

end

function testF3fMC()

    answer = matread("matlabtestdata/testF3f.mat")

    dt = 0.05 # ms
    t = 0.1:dt:6 # ms
    T = length(t)
    xpos = -5:0.1:5 # cm
    rf = 0.05 * sinc.(t .- 3) .* (exp.(im * 2π * 900 * t/1000) + exp.(-im * 2π * 900 * t/1000)) # G
    grad1 = Gradient(0.1, 0, 0)
    grad2 = Gradient(-0.052, 0, 0)
    rf = RF(rf, dt, grad1)
    T1 = 600 # ms
    T2 = 100 # ms
    Δf = 0 # Hz
    pos = [Position(x, 0, 0) for x in xpos]
    spins = map(pos -> SpinMC(1, [1, 0], [T1, T1], [T2, T2], [Δf, Δf], [Inf, Inf], pos), pos)
    A1 = BlochMcConnellMatrix(2)
    B1 = MagnetizationMC(2)
    A2 = BlochMcConnellMatrix(2)
    B2 = MagnetizationMC(2)
    BtoM = MagnetizationMC(2)
    map(spins) do spin
        excite!(A1, B1, spin, rf)
        freeprecess!(A2, B2, spin, t[end] - dt, grad2)
        applydynamics!(spin, BtoM, A1, B1)
        applydynamics!(spin, BtoM, A2, B2)
    end
    sig = map(signal, spins)
    Mz = map(spin -> spin.M[1].z, spins)

    @test sig ≈ vec(answer["sig"])
    return Mz ≈ vec(answer["mz"])

end

function testB2c()

    answer = matread("matlabtestdata/testB2c.mat")

    spin = Spin(1, 600, 100, 0)
    mese! = MESEBlochSim(1000, 50, 1)
    mag = mese!(spin)[1]

    return signal(mag) ≈ answer["sig"]

end

function testB2d()

    answer = matread("matlabtestdata/testB2d.mat")

    spin = Spin(1, 600, 100, 0)
    mese! = MESEBlochSim(1000, 50, 8)
    mag = mese!(spin)

    return map(signal, mag) ≈ vec(answer["sig"])

end

function testB2dMC()

    answer = matread("matlabtestdata/testB2d.mat")

    spin = SpinMC(1, [0, 1], [Inf, 600], [Inf, 100], [0, 0], [Inf, Inf])
    mese! = MESEBlochSim(1000, 50, 8)
    mag = mese!(spin)

    return map(signal, mag) ≈ vec(answer["sig"])

end

function testB3a()

    answer = matread("matlabtestdata/testB3a.mat")

    gradz = 0.3 # G/cm
    Tg = 3 # ms
    z = π/2 / (GAMMA * gradz * Tg/1000) # cm
    spin = Spin(1, 600, 100, 0, Position(0,0,z))
    spgr! = SPGRBlochSim(10, 2, π/3, GradientSpoiling(0, 0, gradz, Tg))
    spgr!(spin)

    show(devnull, spgr!)
    show(devnull, "text/plain", spgr!)
    return Vector(spin.M) ≈ vec(answer["M"])

end

function testB3b()

    answer = matread("matlabtestdata/testB3b.mat")

    gradz = 0.3 # G/cm
    Tg = 3 # ms
    zmax = 4π / (GAMMA * gradz * Tg/1000) # cm
    z = ((1:100)/100 .- 0.5) * zmax
    spins = map(z -> Spin(1, 600, 100, 0, Position(0,0,z)), z)
    spgr! = SPGRBlochSim(10, 2, π/3, GradientSpoiling(0, 0, gradz, Tg))
    foreach(spgr!, spins)
    M = mean(spin.M for spin in spins)

    return Vector(M) ≈ vec(answer["M"])

end

function testB3c()

    answer = matread("matlabtestdata/testB3c.mat")

    spin = Spin(1, 600, 100, 0)
    spgr! = SPGRBlochSim(10, 2, π/3)
    spgr!(spin)

    return Vector(spin.M) ≈ vec(answer["M"])

end

function testB5a()

    answer = matread("matlabtestdata/testB5a.mat")

    gradz = 0.3 # G/cm
    Tg = 3 # ms
    zmax = 2π / (GAMMA * gradz * Tg/1000) # cm
    z = (1:100)/100 * zmax
    spins = map(z -> Spin(1, 600, 100, 0, Position(0,0,z)), z)
    spgr! = SPGRBlochSim(10, 2, π/6, RFandGradientSpoiling((0, 0, gradz), Tg), Val(100))
    foreach(spgr!, spins)
    sig = mean(signal(spin.M) for spin in spins)

    return sig ≈ answer["sig"]

end

function testB5b()

    answer = matread("matlabtestdata/testB5b.mat")

    α = LinRange(0, π/2, 51) # rad
    gradz = 0.3 # G/cm
    Tg = 3 # ms
    zmax = 2π / (GAMMA * gradz * Tg/1000) # cm
    z = (1:100)/100 * zmax
    srf = Vector{ComplexF64}(undef, 51)
    for i = 1:51
        spins = map(z -> Spin(1, 600, 100, 0, Position(0,0,z)), z)
        spgr! = SPGRBlochSim(10, 2, α[i], RFandGradientSpoiling((0, 0, gradz), Tg), Val(100))
        foreach(spgr!, spins)
        srf[i] = mean(signal(spin.M) for spin in spins)
    end
    spin = Spin(1, 600, 100, 0)
    sideal = map(α) do α
        spgr! = SPGRBlochSim(10, 2, α)
        spgr!(spin)
        signal(spin)
    end

    @test srf ≈ vec(answer["sig1"])
    return sideal ≈ vec(answer["sig2"])

end

@testset "Compare to MATLAB" begin

    @testset "Single Compartment" begin

        @test testA5b()
        @test testF1a()
        @test testF1b()
        @test testF1c()
        @test testF2c()
        @test testF3a()
        @test testF3c()
        @test testF3d()
        @test testF3f()

    end

    @testset "Multicompartment" begin

        @test testA5bMC()
        @test testF1cMC()
        @test testF2cMC()
        @test testF3cMC()
        @test testF3dMC()
        @test testF3fMC()

    end

    @testset "Sequences" begin

        @testset "MESE" begin

            @test testB2c()
            @test testB2d()
            @test testB2dMC()

        end

        @testset "SPGR" begin

            @test testB3a()
            @test testB3b()
            @test testB3c()
            @test testB5a()
            @test testB5b()

        end

    end

end


#############################################################################
# geom_mean_epicone.jl
# Constrains T to
#   A #_t B ⪯ T
# where:
#   * A #_t B is the t-weighted geometric mean of A and B:
#     A^{1/2} (A^{-1/2} B A^{-1/2})^t A^{1/2}
#     Parameter t should be in [-1, 0] or [1, 2].
#   * Constraints A ⪰ 0, B ⪰ 0 are added.
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#
# Note on fullhyp: GeomMeanEpiCone will always return a full epigraph cone
# (unlike GeomMeanHypoCone) and so this parameter is not really used.  It is
# here just for consistency with the GeomMeanHypoCone function.
#
#REFERENCE
#   Ported from CVXQUAD which is based on the paper: "Lieb's concavity
#   theorem, matrix geometric means and semidefinite optimization" by Hamza
#   Fawzi and James Saunderson (arXiv:1512.03401)
#############################################################################

struct GeomMeanEpiCone
    A::AbstractExpr
    B::AbstractExpr
    t::Rational
    size::Tuple{Int, Int}

    function GeomMeanEpiCone(A::AbstractExpr, B::AbstractExpr, t::Rational, fullhyp::Bool=true)
        if size(A) != size(B)
            throw(DimensionMismatch("A and B must be the same size"))
        end
        n = size(A)[1]
        if size(A) != (n, n)
            throw(DimensionMismatch("A and B must be square"))
        end
        if t < -1 || (t > 0 && t < 1) || t > 2
            throw(DomainError(t, "t must be in the range [-1, 0] or [1, 2]"))
        end
        return new(A, B, t, (n, n))
    end

    GeomMeanEpiCone(A::Value,        B::AbstractExpr, t::Rational, fullhyp::Bool=true) = GeomMeanEpiCone(Constant(A), B, t, fullhyp)
    GeomMeanEpiCone(A::AbstractExpr, B::Value,        t::Rational, fullhyp::Bool=true) = GeomMeanEpiCone(A, Constant(B), t, fullhyp)
    GeomMeanEpiCone(A::Value,        B::Value,        t::Rational, fullhyp::Bool=true) = GeomMeanEpiCone(Constant(A), Constant(B), t, fullhyp)

    GeomMeanEpiCone(A::AbstractExprOrValue, B::AbstractExprOrValue, t::Integer, fullhyp::Bool=true) = GeomMeanEpiCone(A, B, t//1, fullhyp)
end

struct GeomMeanEpiConeConstraint <: Constraint
    head::Symbol
    id_hash::UInt64
    T::AbstractExpr
    cone::GeomMeanEpiCone

    function GeomMeanEpiConeConstraint(T::AbstractExpr, cone::GeomMeanEpiCone)
        if size(T) != cone.size
            throw(DimensionMismatch("T must be size $(cone.size)"))
        end
        id_hash = hash((cone.A, cone.B, cone.t, :GeomMeanEpiCone))
        return new(:GeomMeanEpiCone, id_hash, T, cone)
    end

    GeomMeanEpiConeConstraint(T::Value, cone::GeomMeanEpiCone) = GeomMeanEpiConeConstraint(Constant(T), cone)
end

in(T, cone::GeomMeanEpiCone) = GeomMeanEpiConeConstraint(T, cone)

function AbstractTrees.children(constraint::GeomMeanEpiConeConstraint)
    return (constraint.T, constraint.cone.A, constraint.cone.B, "t=$(constraint.cone.t)")
end

# For t ∈ [-1, 0] ∪ [1, 2] the t-weighted matrix geometric mean is matrix convex (arxiv:1512.03401).
# So if A and B are convex sets, then T ⪰ A #_t B will be a convex set.
function vexity(constraint::GeomMeanEpiConeConstraint)
    A = vexity(constraint.cone.A)
    B = vexity(constraint.cone.B)
    T = vexity(constraint.T)

    # NOTE: can't say A == NotDcp() because the NotDcp constructor prints a warning message.
    if typeof(A) == ConcaveVexity || typeof(A) == NotDcp
        return NotDcp()
    end
    if typeof(B) == ConcaveVexity || typeof(B) == NotDcp
        return NotDcp()
    end
    # Copied from vexity(c::GtConstraint)
    vex = ConvexVexity() + (-T)
    if vex == ConcaveVexity()
        vex = NotDcp()
    end
    return vex
end

function conic_form!(constraint::GeomMeanEpiConeConstraint, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, constraint)
        A = constraint.cone.A
        B = constraint.cone.B
        t = constraint.cone.t
        T = constraint.T
        is_complex = sign(A) == ComplexSign() || sign(B) == ComplexSign() || sign(T) == ComplexSign()
        if is_complex
            make_temporary = () -> HermitianSemidefinite(size(A)[1])
        else
            make_temporary = () -> Semidefinite(size(A)[1])
        end

        Z = make_temporary()

        if t <= 0
            conic_form!([T A; A Z] ⪰ 0, unique_conic_forms)
            conic_form!(Z in GeomMeanHypoCone(A, B, -t, false), unique_conic_forms)
        else
            @assert t >= 1 # range of t checked in GeomMeanEpiCone constructor
            conic_form!([T B; B Z] ⪰ 0, unique_conic_forms)
            conic_form!(Z in GeomMeanHypoCone(A, B, 2-t, false), unique_conic_forms)
        end

        cache_conic_form!(unique_conic_forms, constraint, Array{Convex.ConicConstr,1}())
    end
    return get_conic_form(unique_conic_forms, constraint)
end


include("../funcitons/layer_backend/linear.jl")

mutable struct Linear <: Layer
    W :: Union{Nothing, Variable}
    b :: Union{Nothing, Variable}
    in_dim :: Union{Nothing, Int}
    out_dim :: Int
    initial_method :: Union{String, Function}
    function Linear(out_dim; in_dim=nothing, initial_method="xavier")
        if in_dim isa Nothing
            b = Variable(zeros(1, out_dim), name="b")
            layer = new(nothing, b, in_dim, out_dim, initial_method)
            return layer
        else
            if initial_method == "xavier"
                W = xavier(in_dim, out_dim)
            elseif initial_method == "he"
                W = he(in_dim, out_dim)
            else initial_method isa Function
                W = initial_method(in_dim, out_dim)
            end
            W = Variable(W, name="W")
            b = Variable(zeros(1, out_dim), name="b")
            layer = new(W, b, in_dim, out_dim, initial_method)
            return layer
        end
    end
end

parameters(linear::Linear) = (linear.W, linear.b)

function forward(layer::Linear, x)
    if layer.W isa Nothing
        layer.in_dim = size(x.values)[2]
        if layer.initial_method == "xavier"
            W = xavier(layer.in_dim, layer.out_dim)
        elseif layer.initial_method == "he"
            W = he(layer.in_dim, layer.out_dim)
        else layer.initial_method isa Function
            W = layer.initial_method(layer.in_dim, layer.out_dim)
        end
        layer.W = Variable(W, name="W")
    end
    return linear(x, layer.W, layer.b)
end


ex=sqrt(:a)*sqrt(:a)
@test length(Equations.unsqrt!(ex[1]))==1
ex=:b*sqrt(:a)*sqrt(:a)
@test length(Equations.unsqrt!(ex[1]))==2
ex=sqrt(:a)*sqrt(:b)*sqrt(:a)
@test length(Equations.unsqrt!(ex[1]))==2

eq=@equ fp=sqrt(a)
@test (eq^2).rhs==:a


using Test
using HDGElasticity

l = 2.0
m = 3.0

matrix = HDGElasticity.plane_strain_voigt_hooke_matrix_2d(l,m)
testmatrix = [(l+2m)  l     0.0
                l   (l+2m) 0.0
               0.0   0.0    m]
@test all(matrix .== testmatrix)

matrix = HDGElasticity.plane_strain_voigt_hooke_matrix(l,m,2)
@test all(matrix .== testmatrix)
@test_throws ArgumentError HDGElasticity.plane_strain_voigt_hooke_matrix(l,m,3)
@test_throws ArgumentError HDGElasticity.plane_strain_voigt_hooke_matrix(l,m,1)


using Combinatorics

term = "RCode"
i = 0
pcnt = factorial(length(term))
print("All the permutations of ", term, " (", pcnt, "):\n    ")
for p in permutations(split(term, ""))
    global i
    print(join(p), " ")
    i += 1
    i %= 12
    i != 0 || print("\n    ")
end
println()


module cartoonLike

using Plots
using LinearAlgebra
using Interpolations
using NearestNeighbors

export create_circle, given_points, distinguishing_curvature, cartoonLikeImage

# still to make: piecewise smooth curves
#                fill in the interior, exterior with something smooth

struct cartoonLikeImage{T}
  xSize::Int
  ySize::Int
  ninterp::Int
  ε::T
  function cartoonLikeImage{T}(; xSize=100, ySize=100, ninterp = xSize*ySize,
                               ε = 1/(xSize*ySize)^(3/4)) where T<: Real
    return new(xSize, ySize, ninterp, T(ε))
  end
end

"""
    create_circle(r, rate = 10)
create a circle of radius r, sampled at rate `rate`
"""
function create_circle(r, rate = 10)
    t = range(0,1,length = 10)
    x = r .* sin.(2π*t)
    y = r .* cos.(2π*t)
    A = hcat(x,y)
end

"""
    image, xlocs, ylocs = given_points(points; xSize=100, ySize=100, ninterp =
        xSize*ySize, ε = .01)

given a collection of nodes, i.e. row vectors of the form (x,y), make a
(xSize)×(ySize) image of path going through them, which has been blurred with a
gaussian of standard deviation of ε (\epsilon)
"""
function given_points(points::AbstractArray{T}; cart::cartoonLikeImage{<:Real} =
                      cartoonLikeImage{T}()) where T<:Real
    xSize = cart.xSize
    ySize = cart.ySize
    ninterp = cart.ninterp
    ε = cart.ε
    t= range(0, 1, length = size(points, 1))
    itp = scale(interpolate(points, (BSpline(Cubic(Natural(OnGrid()))), NoInterp())), t, 1:2)
    tfine = range(0, 1, length = ceil(Int, ninterp))
    finePoints = itp(tfine, 1:2)'
    tree = KDTree(Array{Float64,2}(finePoints))

    xRange = range(-1, 1, length=xSize)
    yRange = range(-1, 1, length=ySize)

    imagePoints = zeros(2, xSize, ySize)
    for (ix,x) =enumerate(xRange), (iy,y) = enumerate(yRange)
        imagePoints[1, ix, iy] = x
        imagePoints[2, ix, iy] = y
    end
    image = zeros(xSize, ySize)
    idxs, dists = knn(tree, reshape(imagePoints,(2,:)),1)
    dists = [d[1] for d in dists]       # the distances are given as lists, and we want
    # it as just an array
    dists = reshape(dists, size(image))
    image = exp.(-(dists).^2 ./ε)
    return (image, xRange, yRange)
end

"""
make a collection of `n` images that have a range of different curvatures at a
  central point, equally spaced
"""
function distinguishing_curvature(nPoints; cart::cartoonLikeImage{<:Real} =
                                  cartoonLikeImage{Float32}())
    yRange = range(-1,1,length=2*nPoints+1)
    points = make_points.(yRange)
    image_groups = given_points.(points, cart = cart)
    xLocs = image_groups[1][2]

    yLocs = image_groups[1][3]

    images = zeros(cart.xSize, cart.ySize, length(image_groups))
    for (i,tup) in enumerate(image_groups)
        images[:, :, i] = tup[1]
    end
    return images
end

function make_points(y)
  return [-.5 0 .5;
           y  0  y]'
end

end


include("../src/SFFM.jl")
using LinearAlgebra, Plots

## define the model(s)
@testset "$modelfile" for modelfile in ["testModel1.jl"; "testModel2.jl"; "testModel3.jl"]
    include(modelfile)

    ## section 4.3: the marginal stationary distribution of X
    ## mesh
    Δ = 0.4
    Nodes = collect(approxBounds[1, 1]:Δ:approxBounds[1, 2])
    NBases = 2
    Basis = "lagrange"
    mesh = SFFM.DGMesh(
        model,
        Nodes,
        NBases,
        Basis = Basis,
    )
    Ψₓ = SFFM.PsiFunX(model)
    ξₓ = SFFM.MakeXiX(model, Ψₓ)

    pₓ, πₓ, Πₓ, Kₓ = SFFM.StationaryDistributionX(model, Ψₓ, ξₓ)

    analyticX = (
            pm = [pₓ[:];0;0],
            distribution = πₓ(SFFM.CellNodes(mesh)),
            x = SFFM.CellNodes(mesh),
            type = "density"
        )

    @test sum(Πₓ(40))-1 < 1e-6

    @testset "DG witn NBases $NBases" for NBases in 1:2
            mesh = SFFM.DGMesh(
                model,
                Nodes,
                NBases,
                Basis = Basis,
            )

            # compute the marginal via DG
            All = SFFM.MakeAll( model, mesh, approxType = "projection")
            @test isapprox(sum(All.B.B,dims=2), zeros(size(All.B.B,1)), atol = sqrt(eps()))
            @test -sum(All.B.BDict["++"],dims=2) ≈ sum(All.B.BDict["+-"],dims=2) + sum(All.B.BDict["+0"],dims=2)
            @test -sum(All.B.BDict["--"],dims=2) ≈ sum(All.B.BDict["-+"],dims=2) + sum(All.B.BDict["-0"],dims=2)
            @test -sum(All.B.BDict["00"],dims=2) ≈ sum(All.B.BDict["0-"],dims=2) + sum(All.B.BDict["0+"],dims=2)
            @test -sum(All.D["++"](s=0),dims=2) ≈ sum(All.D["+-"](s=0),dims=2)
            @test sum(All.D["-+"](s=0),dims=2) ≈ -sum(All.D["--"](s=0),dims=2)

            Ψ = SFFM.PsiFun(All.D)
            @test sum(Ψ,dims=2) ≈ ones(size(Ψ,1))

            # the distribution of X when Y first returns to 0
            ξ = SFFM.MakeXi(All.B.BDict, Ψ)
            @test sum(ξ) ≈ 1

            marginalX, p, K = SFFM.MakeLimitDistMatrices(
                All.B.BDict,
                All.D,
                All.R.RDict,
                Ψ,
                ξ,
                mesh,
                model,
            )

            # convert marginalX to a distribution for plotting
            Dist = SFFM.Coeffs2Dist(
                    model,
                    mesh,
                    marginalX,
                    type="probability",
                )
            @test sum(Dist.pm) + sum(Dist.distribution) ≈ 1
    end
end


function prs_predict(x::Array{Float64, 2}, y::Array{Float64, 2},
  xx::Array{Float64, 2}; )

# Some radial basis functions
rbf_thin_plate(r::Float64, c::Float64 = 1.0) = (r == 0.0) ? 0.0 : (r^2 * log(c*r))

rbf_gaussian(r::Float64, c::Float64 = 1.0) = exp(-c*r^2)

rbf_cubic(r::Float64, c::Float64 = 1.0) = r^3

rbf_multiquadric(r::Float64, c::Float64 = 1.0) = sqrt(r^2 + c^2)

# Calculate the pairwise kernel matrix for the points in x taken to be the
# individual columns.
function calc_kernel_matrix(x::Array{Float64, 2}, centers::Array{Float64, 2},
  phifunc = rbf_gaussian, distfunc = norm, c = 1.0)
  num_points = size(x, 2)
  num_centers = size(centers, 2)
  k_matrix = zeros(Float64, num_points, num_centers)
  for i in 1:num_points
    xi = x[:,i]
    for j in 1:num_centers
      k_matrix[i, j] = phifunc(norm(xi - centers[:,j]), c)
    end
  end
  k_matrix
end

function rbf_predict(x::Array{Float64, 2}, y::Array{Float64, 2},
  xx::Array{Float64, 2}; c = nothing, rbf = rbf_gaussian)

  # Set c to ok values unless it was specified
  if c == nothing && rbf == rbf_gaussian
    D = size(x, 1)
    c = 1/(D*D)
  elseif c == nothing
    c = 1.0
  end

  # Scale x per row (i.e. covariate)
  x_row_means = mean(x, 2)
  x_row_std = std(x, 2)
  xt = broadcast(-, x, x_row_means)
  xsc = broadcast(/, xt, x_row_std)

  # Scale y per its mean and std
  ymean = mean(y)
  ystd = std(y)
  ysc = (y - ymean) / ystd

  # Dor LSR to get lambdas
  A = calc_kernel_matrix(xsc, xsc, rbf, norm, c)
  lambdas = A\(ysc')

  # Scale xx
  xxt = broadcast(-, xx, x_row_means)
  xxsc = broadcast(/, xxt, x_row_std)

  # Predict for the test points
  B = calc_kernel_matrix(xxsc, xsc, rbf, norm, c)
  ypredsc = (B * lambdas)'

  # Transform back to original range
  ypred = ymean + ystd * ypredsc

  return ypred

end

function r2(ypred, y)
  SSR = sum( (y .- ypred).^2  )
  SST = sum( (y - mean(y)).^2 )
  1.0 - SSR / SST
end

# Test function
tf(x) = sum(x.^4)

N = 100
D = 10
xtraining = randn(D, N)
ytraining = mapslices(tf, xtraining, 1)
xtest = randn(D, 5*D)
y = mapslices(tf, xtest, 1)

ypred_tp = rbf_predict(xtraining, ytraining, xtest; rbf = rbf_thin_plate)
r2_tp = r2(ypred_tp, y)

ypred_multiquadric = rbf_predict(xtraining, ytraining, xtest; rbf = rbf_multiquadric)
r2_multiquadric = r2(ypred_multiquadric, y)

ypred_cubic = rbf_predict(xtraining, ytraining, xtest; rbf = rbf_cubic)
r2_cubic = r2(ypred_cubic, y)

ypred_gaussiand = rbf_predict(xtraining, ytraining, xtest; rbf = rbf_gaussian, c = 1/D)
r2_gaussiand = r2(ypred_gaussiand, y)

ypred_gaussiand2 = rbf_predict(xtraining, ytraining, xtest; rbf = rbf_gaussian, c = 1/(D*D))
r2_gaussiand2 = r2(ypred_gaussiand2, y)


abstract type KrylovExponentiator{T} <: AbstractExponentiator{T} end

# This is of course only true if the Krylov subspace is large enough.
exact(::KrylovExponentiator) = true
order(::KrylovExponentiator) = Inf

mutable struct Lanczos{T,U,O<:AbstractOperator{T},scale} <: KrylovExponentiator{T}
    A::O
    μ::T
    Ks::KrylovSubspace{U,T,U}
    cache::StegrCache{T,U}
    atol::U
    rtol::U
end

function LanczosMap(::Type{T}, A::O, μ::T, m::Integer; scale::Bool=true, atol=1e-8, rtol=√atol) where {T,O<:AbstractOperator}
    n = size(A,2)
    Ks = KrylovSubspace{T,real(T)}(n,m)
    LanczosMap{T,real(T),O,scale}(A, μ, Ks, StegrCache(T, m), atol, rtol)
end

LanczosMap(::Type{T}, A, m::Integer; kwargs...) where T =
    LanczosMap(T, A, one(T), m; kwargs...)

operator(L::LanczosMap) = L.A

# Scale the operator by the time-step, to reduce the spectral radius.
expv!(v, L::LanczosMap{T,U,true}, Δt::U, u, _) where {T,U} =
    expv!(v, L.μ, Δt*L.A, u, L.Ks, L.cache, atol=L.atol, rtol=L.rtol)

# Apply the time-step when calculating the subspace exponential.
expv!(v, L::LanczosMap{T,U,false}, Δt::U, u, _) where {T,U} =
    expv!(v, L.μ*Δt, L.A, u, L.Ks, L.cache, atol=L.atol, rtol=L.rtol)


function slope_intercept(wn_dict::Dict{Any,Any}, wn::PyCall.PyObject, link_results::Dict{Any,Any}, node_results::Dict{Any,Any}, num_timesteps::Int64)
    slopes = Dict{String,Any}()
    intercepts = Dict{String, Any}()
    pump_off = Dict{String, Any}()
    slope_sum = 0
    intercept_sum = 0
    count = 0
    for (key, pump) in wn_dict["pumps"]
        name = pump["name"]
        flow = link_results["flowrate"][name][:values][1:num_timesteps]
        headloss = link_results["headloss"][name][:values][1:num_timesteps]
        efficiency = nothing
        try
            efficiency = pump["efficiency"][:points]
        catch
            wn_dict["options"]["energy"]["global_efficiency"] != nothing ? efficiency = wn_dict["options"]["energy"]["global_efficiency"]: efficiency = 0.65
        end
        if length(efficiency) > 1
            energy = [((1000.0 * 9.81 * headloss[i] * flow[i])/efficiency[1][1]) for i = 1:num_timesteps]
        else
            energy = [((1000.0 * 9.81 * headloss[i] * flow[i])/efficiency) for i = 1:num_timesteps]
        end
        flows = Array{Float64}(0)
        energies = Array{Float64}(0)
        for i= 1:length(energy)
            if energy[i] != 0.0 && flow[i] != 0.0
                flows = vcat(flows,flow[i])
                energies = vcat(energies, energy[i])
            end
        end
        if length(energies) != 0.0
            count = count +1
            intercept, slope = LeastSquares(flows, energies)
            slopes[key] = slope
            intercepts[key] = intercept
            slope_sum = slope_sum + slope
            intercept_sum = intercept_sum + intercept
        else
            pump_off[key] = key
        end
    end

    avg_slope = slope_sum/count
    avg_intercept = intercept_sum/count
    for (key, name) in pump_off
        slopes[key] = avg_slope
        intercepts[key] = avg_intercept
    end
    return slopes, intercepts
end

function LeastSquares(flows::Array{Float64}, energies::Array{Float64})
    #if pump never turns on, delete control that inhibits pump, calulation energy, then replace control
        avg_energy = mean(energies)
        a, b = linear_fit(flows, energies)
        error = Array{Float64}(0)
        for i = 1:length(energies)
            push!(error, energies[i] - (avg_energy +b *flows[i]))
        end
        avg_error = mean(error)
        intercept = avg_energy + avg_error
        return intercept, b
    # total_error = energies - (intercept + b * flows)
end


""" geometric predicates for 2vectors, cartesian points (px,py), homogeneuos points (pw,px,py), lines (lw,lx,ly), and segments x[1:2],y[1:2].
	Julia can inline these.
"""
Vector2 = Tuple{Real,Real}
Point2C = Tuple{Real,Real}
Line2 = Tuple{Real,Real,Real}
Point2H = Tuple{Real,Real,Real}
Point2 = Union{Point2C,Point2H}

toPoint2H((x,y)::Point2C)::Point2H = (one(x),x,y)

# some special handling for floats
dot((ux,uy)::Tuple{AbstractFloat,AbstractFloat}, (vx,vy)::Tuple{AbstractFloat,AbstractFloat}) = muladd(ux, vx, 1.0uy*vy)
dot((ux,uy)::Vector2, (vx,vy)::Vector2) = muladd(ux, vx, uy*vy) # dot product for 2d vectors. (consider muladd)
perpdot((ux,uy)::Tuple{AbstractFloat,AbstractFloat}, (vx,vy)::Tuple{AbstractFloat,AbstractFloat}) = muladd(ux, vy, -1.0uy*vx)
perpdot((ux,uy)::Vector2, (vx,vy)::Vector2) = muladd(ux, vy, -uy*vx) # dot product of ccw perp(u) with v. (consider muladd)

# basic cartesian and homogeneous orientation determinants
ccw(p::Point2C,q::Point2C,r::Point2C) = perpdot(q.-p,r.-p) # orientation determinant
ccw((pw,px,py)::Point2H,(qw,qx,qy)::Point2H,(rw,rx,ry)::Point2H) = pw*perpdot((qx,qy),(rx,ry)) -qw*perpdot((px,py),(rx,ry)) + rw*perpdot((px,py),(qx,qy))# orientation determinant
ccw(p::Point2C,q::Point2C,(rw,rx,ry)::Point2H) = perpdot(q.-p,(rx,ry))+rw*perpdot(p,q.-p) # orientation determinant

# mixed cartesian and homogeneuos
ccw(p::Point2C,q::Point2H,r::Point2H) = ccw(toPoint2H(p),q,r)
ccw(p::Point2H,q::Point2C,r::Point2H) = ccw(p,toPoint2H(q),r)
ccw(p::Point2H,q::Point2H,r::Point2C) = ccw(p,q,toPoint2H(r))
ccw(p::Point2H,q::Point2C,r::Point2C) = ccw(q,r,p)
ccw(p::Point2C,q::Point2H,r::Point2C) = ccw(r,p,q)

## Geometric predicates based on ccw
leftturn(p::Point2,q::Point2,r::Point2)::Bool = ccw(p,q,r)>zero(p[1]) # p to q to r is a left turn
rightturn(p::Point2,q::Point2,r::Point2)::Bool = ccw(p,q,r)<zero(p[1]) # p to q to r is a right turn

## leftAhead halps solve on-the-llne cases consistent with infinitesimal rotation
leftAhead(p::Point2C,q::Point2C,r::Point2C)::Bool = (d = ccw(p,q,r)) == zero(d) ? dot(p.-q,r.-q)<zero(d) : d>zero(d) # should assert ≉ 0 for precision checking
rightAhead(p::Point2C,q::Point2C,r::Point2C)::Bool = (d = ccw(p,q,r)) == zero(d) ? dot(p.-q,r.-q)<zero(d) : d<zero(d) # should assert ≉ 0 for precision checking
leftOnAhead(p::Point2C,q::Point2C,r::Point2C)::Bool = (d = ccw(p,q,r)) == zero(d) ? dot(q.-p,r.-p)>zero(d) : d>zero(d) # should assert ≉ 0 for precision checking
rightOnAhead(p::Point2C,q::Point2C,r::Point2C)::Bool = (d = ccw(p,q,r)) == zero(d) ? dot(q.-p,r.-p)>zero(d) : d<zero(d) # should assert ≉ 0 for precision checking

## orientation with drected segments
#  This should probably be expanded.
leftofseg(q::Point2, xs, ys)::Bool = @inbounds leftturn((xs[1],ys[1]),(xs[2],ys[2]), q) # point q left of segment defined by pairs of xs and ys
rightofseg(q::Point2, xs, ys)::Bool = @inbounds rightturn((xs[1],ys[1]),(xs[2],ys[2]), q) # point q right of segment defined by pairs of xs and ys

leftward(s0::Point2C,s1::Point2C, t0::Point2C,t1::Point2C)::Bool = perpdot(t1.-t0, s1.-s0)>zero(s0[1])   # vector s0->s1 is in the direction from right to left of the line t0->t1.
rightward(s0::Point2C,s1::Point2C, t0::Point2C,t1::Point2C)::Bool = perpdot(t1.-t0, s1.-s0)<zero(s0[1])   # vector s0->s1 is in the direction from left to right of the line t0->t1.

# computing homogeneous line equations
line(p::Point2C,q::Point2C)::Line2 = (perpdot(p,q), last(p)-last(q), first(q)-first(p)) # line equation
orientptln((x,y)::Point2C, (lw,lx,ly)::Line2) = muladd(y,ly,muladd(x,lx,lw))
orientptln((w,x,y)::Point2H, (lw,lx,ly)::Line2) = muladd(y,ly,muladd(x,lx,w*lw))
leftof(q::Point2, l::Line2)::Bool = orientptln(q,l)>zero(l[1]) # point left of (consistent with) oriented line
rightof(q::Point2, l::Line2)::Bool = orientptln(q,l)<zero(l[1])


clear()
@testset "check ScalePath op's gradient" begin
    # [1] prepare input data and its label
    input = randn(19,17,7);
    label = randn(19,17,7);
    x = Variable(input, type=Array{Float64});
    l = Variable(label, type=Array{Float64});

    # [2] prepare a learnabel scalar multiplier
    k = ScalePath(1.0, ndims=3, type=Array{Float64});

    # [3] forward and get gradient by backpropagation
    y  = forward(k, x);
    C₁ = mseCost(y, l);
    backward();
    ∇ = Delta.δ(k.scale)[1];  # backpropagation gradient

    # [4] with a samll change of weight, forward again
    Δ = 1e-6;
    Delta.ᵛ(k.scale)[1] += Δ;

    y  = forward(k, x);
    C₂ = mseCost(y, l);
    backward();

    # [5] check if the auto-grad is true or not
    ϵ    = eps(Float64)
    ∂C∂k = (C₂ - C₁)/Δ;            # numerical gradient
    err  = (∂C∂k - ∇)/(∇+ϵ);       # relative error
    err  = abs(err) * 100;         # relative error in %
    err  = err < 1e-1 ? 0.0 : err; # zero-error if under 0.1%
    @test err < 1.0
end


using Statistics, SpecialFunctions

"""
    xerfcx_pi(x)

Computes 1/√π - x * erfcx(x), without cancellation.
"""
function xerfcx_pi(x::Real)
    if xerfcx_asym_thresh(x, 3)
        return xerfcx_asym_pi(x, 3)
    else
        return 1/√π - x * erfcx(x)
    end
end

"""
    xerfcx_asym(x, m)

Evaluates x * erfcx(x) by its asymptotic series for large x > 0,
up to m'th order.
"""
function xerfcx_asym(x::Real, m::Integer)
    #= NIST Handbook of Mathematical Functions (2010) 1st Edition, 7.12.1 =#
    @assert m ≥ 0 && x > 0
    s = one(x / 1)
    for j = m : -1 : 0
        #s = 1 - (j + 1/2) / x^2 * s
        s = muladd((j + 1/2) / x / x, -s, 1)
    end
    return s / √π
end

"""
    xerfcx_asym(x, m)

Estimated reminder of xerfcx_asym(x, m).
"""
function xerfcx_asym_rem(x::Real, m::Integer)
    #= NIST Handbook of Mathematical Functions (2010) 1st Edition, 7.12.1. =#
    @assert m ≥ 0 && x > 0
    return 1 / √π * exp(lgamma(m + 1/2) - (2m + 1) * x) / erfcx(x)
end

"""
    xerfcx_asym_pi(x, m)

Computes 1/√π - x * erfcx(x) by its asymptotic series for large x > 0, up to
m'th order.
"""
function xerfcx_asym_pi(x::Real, m::Integer)
    #= NIST Handbook of Mathematical Functions (2010) 1st Edition, 7.12.1 =#
    @assert m ≥ 0 && x > 0
    s = one(x / 1)
    for j = m : -1 : 1
        #s = 1 - s * (j + 1/2) / x^2
        s = muladd((j + 1/2) / x / x, -s, 1)
    end
    s = s / 2 / x / x
    return s / √π
end

"""
    xerfcx_asym_thresh(x, m)

Returns true if the asymptotic series of x * erfcx(x) to order m can be used
for this value of x without appreciable error.
"""
xerfcx_asym_thresh(x::Real, m::Integer) = xerfcx_asym_thresh(x, Val(m))
function xerfcx_asym_thresh(x::T, ::Val{m}) where {T <: Real, m}
    @assert m isa Integer
    @assert m > 0
    return x > (√π * eps(float(T)))^(-1/2m) * exp(lgamma(m + 1/2) / 2m)
end
xerfcx_asym_thresh(x::Union{Float64,Integer}, ::Val{3}) = x > 452
xerfcx_asym_thresh(x::Union{Float32}, ::Val{3}) = x > 16
xerfcx_asym_thresh(x::Union{Float16}, ::Val{3}) = x > 4

"""
    two

Returns 2 of appropriate type. Like `one`, but `two`.
"""
two(::Union{Type{T}, T}) where {T} = convert(T, 2)


using PyPlot
using Polynomials
include("sturms_theorem.jl")

# Compute power-spectrum & number of positive zeros:

# Choose an arbitrary number of Lorentzians from 2-10:
n_lor = rand(2:10)

# Set up random coefficients in a vector:
x = zeros(4*n_lor)
# Choose the maximum time & frequency for plotting purposes:
max_time = 0
max_freq = 0

# Randomly choose coefficients.
# These are in the order:  x = {(ar_j,ai_j,br_j,bi_j); j=1,...,n_lor}
for i=1:n_lor
# Linearly distribute ar & ai from -1 to 1:
  x[(i-1)*4+1]=(rand()-.5)*2
  x[(i-1)*4+2]=(rand()-.5)*2
# Logarithmically distribute br & bi:
  x[(i-1)*4+3]=10.^(rand()*2-1)
  x[(i-1)*4+4]=10.^(rand()*2-1)
# Set range of plot to capture behavior of these:
  if 1/x[(i-1)*4+3] > max_time
    max_time = 1/x[(i-1)*4+3]
  end
  if x[(i-1)*4+3] > max_freq
    max_freq = x[(i-1)*4+3]
  end
  if 1/x[(i-1)*4+4] > max_time
    max_time = 1/x[(i-1)*4+4]
  end
  if x[(i-1)*4+4] > max_freq
    max_freq = x[(i-1)*4+4]
  end
end

# The order of the numerator polynomial in z = \omega^2 that we need to
# see if it has any zeros:
pord = 2*(n_lor-1) + 1

# Compute coefficients in the numerator & denominator of PSD:
# The PSD is given by:
#
# \sum_i (a_i z +b_i)/(z^2 + c_i*z + d_i)
#
# where z = \omega^2
#
# We are computing the number of positive real zeros of this equation.
#
a = zeros(n_lor)
b = zeros(n_lor)
c = zeros(n_lor)
d = zeros(n_lor)

for i=1:n_lor
  ar_j = x[(i-1)*4+1]
  ai_j = x[(i-1)*4+2]
  br_j = x[(i-1)*4+3]
  bi_j = x[(i-1)*4+4]
  a[i] = ar_j*br_j-ai_j*bi_j
  b[i] = (bi_j^2+br_j^2)*(ai_j*bi_j+ar_j*br_j)
  c[i] = 2*(br_j^2-bi_j^2)
  d[i] = (bi_j^2+br_j^2)^2
end

# Now, loop over coefficients:

# Initialize a polynomial:
p0 = Poly(zeros(pord+1))
for i=1:n_lor
# The polynomial for the current Lorentzian term in the common-denominator expression:
  pcur = Poly([b[i],a[i]])
  for j=1:n_lor
# Only multiply by the denominators from the other Lorentzians:
    if j != i
      pcur *= Poly([d[j],c[j],1])
    end
  end
  p0 += pcur
end

# Compute the roots of this polynomial:
poly_root = roots(p0)

# Now that we've computed coefficients of the polynomial, we just need
# to apply Sturm's theorem!

# Set up an array to hold signs of polynomial at zero & infinity:
f_of_0 = zeros(pord+1)
f_of_inf = zeros(pord+1)

# Take the derivative of the polynomial:
p1 = polyder(p0)
#println(p0)
#println(p1)
# Insert the coefficients of the z^0 term:
f_of_0[1] = p0(0)
f_of_0[2] = p1(0)
# Insert the coefficient of the z^(p_ord-i) term:
f_of_inf[1] = p0[pord]
f_of_inf[2] = p1[pord-1]

# Now, loop over the Sturm polynomial series:
for i=3:pord+1
  p2 = -rem(p0,p1)
# Check that round-off error hasn't left us with a polynomial
# that is the same order as p0 or p1, but with a small coefficient:
  if length(p2) >= length(p1)
    coeff = zeros(pord-i+2)
    for j=0:pord-i+1
      coeff[j+1]=p2[j]
    end
    p2 = Poly(coeff)
  end
# Insert the z^0 term:
  f_of_0[i]=p2(0)
# Insert the z -> \infty term:
  f_of_inf[i]=p2[pord-i+1]
# Now move promote these polynomials, readying them for recursion in the next step:
  p0=copy(p1)
  p1=copy(p2)
end

# Now we'll compute the number of sign changes at z=0:
sig_0 = 0
for i=1:pord
#  if f_of_0[i+1]*f_of_0[i] < 0 && abs(f_of_0[i+1]) > eps() && abs(f_of_0[i]) > eps()
  if sign(f_of_0[i+1]) != sign(f_of_0[i])
    sig_0 +=1
  end
end
# Next, compute the number of sign changes at z=\infty:
sig_inf = 0
for i=1:pord
#  if f_of_inf[i+1]*f_of_inf[i] < 0 && abs(f_of_inf[i+1]) > eps() && abs(f_of_inf[i]) > eps()
  if sign(f_of_inf[i+1]) != sign(f_of_inf[i])
    sig_inf +=1
  end
end

# The difference between z=0 & z=\infty gives the number of positive, real roots:
println("Number of positive, real roots: ",sig_0-sig_inf)
println("Result  of sturms_theorem.jl:   ",sturms_theorem(x))

# Print out the actual roots for inspection:
println("Roots: ",poly_root)

# Now, determine which roots are positive & real from the numerial solver:
n_pos_real = 0
real_root = Float64[]
for i=1:pord
  if real(poly_root[i]) > 0 && imag(poly_root[i]) == 0
    n_pos_real +=1
    poly_root0 = sqrt(real(poly_root[i]))
    real_root = [real_root;poly_root0]
    if poly_root0 >  max_freq
      max_freq = poly_root0
    end
  end
end
println("Count of positive, real roots: ",n_pos_real)

# Finally, plot the power spectrum & auto-correlation functions:
nom = 100000
omega = linspace(0,2max_freq,nom)

psd = zeros(nom)
ntau = 100000
tau = linspace(0,max_time*4*pi,ntau)
acf = zeros(ntau)

# Compute PSD & ACF:
for i=1:n_lor
  ar_j = x[(i-1)*4+1]
  ai_j = x[(i-1)*4+2]
  br_j = x[(i-1)*4+3]
  bi_j = x[(i-1)*4+4]
  psd += sqrt(2/pi)*((omega.^2.*(ar_j*br_j-ai_j*bi_j)+(ai_j*bi_j+ar_j*br_j)*(bi_j^2+br_j^2))./(omega.^4+2(br_j^2-bi_j^2).*omega.^2 + (bi_j^2+br_j^2)^2))
  acf += (exp(-br_j.*tau).*(ar_j.*cos(bi_j.*tau)+ai_j.*sin(bi_j.*tau)))
end

# Make the plots:
fig,axes = subplots(2,1)
ax = axes[1]
ax[:plot](omega,psd)
# Overplot the numerically computed roots:
if n_pos_real > 0
  for i=1:n_pos_real
    ax[:scatter](real_root[i],0)
  end
end
# Now plot ACF:
ax = axes[2]
ax[:plot](tau,acf)

#println(minimum(psd))


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

W = TropicalWeight{Float64}
sym=Dict("$(Char(i))"=>i for i=97:122)
fst = WFST(sym; W = W)
add_arc!(fst,1,3,"<eps>","<eps>",one(W))
add_arc!(fst,3,6,"f","f",one(W))
add_arc!(fst,6,4,"d","d",one(W))
add_arc!(fst,1,4,"a","a",one(W))
add_arc!(fst,4,5,"c","c",one(W))
add_arc!(fst,5,7,"b","b",one(W))
add_arc!(fst,7,2,"a","a",one(W))
add_arc!(fst,5,2,"b","b",one(W))
initial!(fst,1)
final!(fst,2)
println(fst)

perm = topsortperm(fst)
@test perm == [1,3,6,4,5,7,2]

sorted_fst = topsort(fst)
@test size(sorted_fst) == size(fst)
@test topsortperm(sorted_fst) == [1,2,3,4,5,6,7]
@test isfinal(sorted_fst,7)

# fails if fst has cycles
add_arc!(fst,2,1,"z","z",one(W))
@test_throws ErrorException topsort(fst)


immutable HomotheticHyperbolic{T} <: Model
    a::T
    α::T
end

function Base.show(io::IO, ::Type{HomotheticHyperbolic})
    return print(io, "Homothetic Hyperbolic")
end

function Base.show(io::IO, m::HomotheticHyperbolic)
    @printf(
        io,
        "Homothetic Hyperbolic (a = %.3f, α = %.3f)",
        m.a,
        m.α,
    )
end

function initial_parameters(::Type{HomotheticHyperbolic})
    return Float64[log(1.0), log(1.0)]
end

canonical_parameters(m::HomotheticHyperbolic) = (m.a, m.α)

function HomotheticHyperbolic(parameters::Vector)
    a = exp(parameters[1])
    α = exp(parameters[2])
    return HomotheticHyperbolic(a, α)
end

@inline function predict(
    m::HomotheticHyperbolic,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    a, α = m.a, m.α
    u1 = log(x1 * (1 / (1 + α * t1)))
    u2 = log(x2 * (1 / (1 + α * t2)))
    z = a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::HomotheticHyperbolic,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    a, α = m.a, m.α
    d_t2 = 1 / (1 + α * t2)
    d_t1 = 1 / (1 + α * t1)
    gr = (
        (log(x2 * d_t2) - log(x1 * d_t1)) * a,
        a * (t1 * d_t1 - t2 * d_t2) * α,
    )
    return gr
end


Pkg.instantiate()

using Distributed, StanSample, CSV, DataFrames
Pkg.instantiate()

ProjDir = @__DIR__
cd(ProjDir)

rcd = CSV.read(joinpath(ProjDir, "RedCardData.csv"), DataFrame; missingstring = "NA")
rcd = dropmissing(rcd, :rater1)
rcd[!, :rater1] = Float64.(rcd.rater1)

stan_logistic_0 ="
data {
  int<lower=0> N;
  int<lower=0> n_redcards[N];
  int<lower=0> n_games[N];
  vector[N] rating;
}
parameters {
  vector[2] beta;
}
model {
  beta[1] ~ normal(0, 10);
  beta[2] ~ normal(0, 1);
  n_redcards ~ binomial_logit(n_games, beta[1] + beta[2] * rating);
}
";

stan_logistic_1 = "
functions {
  real partial_sum_lpmf(int[] slice_n_redcards,
                        int start, int end,
                        int[] n_games,
                        vector rating,
                        vector beta) {
    return binomial_logit_lupmf(slice_n_redcards |
                               n_games[start:end],
                               beta[1] + beta[2] * rating[start:end]);
  }
}
data {
  int<lower=0> N;
  int<lower=0> n_redcards[N];
  int<lower=0> n_games[N];
  vector[N] rating;
  int<lower=1> grainsize;
}
parameters {
  vector[2] beta;
}
model {

  beta[1] ~ normal(0, 10);
  beta[2] ~ normal(0, 1);
  target += reduce_sum(partial_sum_lupmf, n_redcards, grainsize,
                       n_games, rating, beta);
}
";

#isdir("tmp0") && rm("tmp0", recursive=true);
if !isdir("tmp0")
  tmp0 = mkdir(joinpath(ProjDir, "tmp0"))
else
  tmp0 = joinpath(ProjDir, "tmp0")
end;

logistic_0 = SampleModel("logistic0", stan_logistic_0, [1]; tmpdir=tmp0)

#isdir("tmp1") && rm("tmp1", recursive=true);
if !isdir("tmp1")
  tmp1 = mkdir(joinpath(ProjDir, "tmp1"))
else
  tmp1 = joinpath(ProjDir, "tmp1")
end;

logistic_1 = SampleModel("logistic1", stan_logistic_1; tmpdir=tmp1)

data = Dict(
    :N => size(rcd, 1),
    :grainsize => 1,
    :n_redcards => rcd.redCards,
    :n_games => rcd.games,
    :rating => rcd.rater1
)

for i in 1:1
  #addprocs(1)
  println("\nUsing $(nprocs()) cores.\n")
  @time rc_0 = stan_sample(logistic_0; data);

  if success(rc_0)
      dfs_0 = read_summary(logistic_0)
      dfs_0[8:9, [1,2,4,8,9,10]] |> display
      println()
  end

  @time rc_1 = stan_sample(logistic_1; data);

  if success(rc_1)
      dfs_1 = read_summary(logistic_1)
      dfs_1[8:9, [1,2,4,8,9,10]] |> display
      println()
  end
end


using GeneralizedNullspaces
using GenericLinearAlgebra
using Test
using Random
using LinearAlgebra

Random.seed!(1234)

include("testutils.jl")

quiet = true

# example from appendix B of the Guglielmi et al. paper
let tol=1e-10
    a = 233/896
    b = 248/896
    c = 171/896
    d = 29/896
    e = 15/896
    z = 69/448
    s = 2101/9632
    t = 295/19264
    u = 1403/28896

    A = [z    s      s      s      t     u     t     u     t     u;
         a    b      c      c      e     d     0     0     0     d;
         a    c      b      c      0     d     e     d     0     0;
         a    c      c      b      0     0     0     d     e     d;
         3/32 7/16   3/32   3/32   3/32  3/32  0     0     0     3/32;
         9/64 39/128 39/128 3/64   3/128 9/64  3/128 1/128 0     1/128;
         3/32 3/32   7/16   3/32   0     3/32  3/32  3/32  0     0;
         9/64 3/64   39/128 39/128 0     1/128 3/128 9/64  3/128 1/128;
         3/32 3/32   3/32   7/16   0     0     0     3/32  3/32  3/32;
         9/64 39/128 3/64   39/128 3/128 1/128 0     1/128 3/128 9/64]

    B, V, ν, μ = gnsd(A,tol)

    # check destructuring
    F = gnsd(A,tol)
    @test F isa GNSD{Float64}
    @test F.B == B
    @test F.V == V
    @test F.ν == ν
    @test F.μ == μ

    @test ν == 2
    @test μ == [3,1]
    @test gnsdcheck(A,B,V,ν,μ)
end

let tol=1e-9
    ν=1
    μx = [3,2]
    AA = makegnb(1,μx)
    n = size(AA,1)
    Q,R = qr(randn(n,n))
    A = Q' * AA * Q
    B,V,ν1,μ = gnsd(A,tol)
    @test ν == 1
    @test μ == μx[1:ν]
    @test gnsdcheck(A,B,V,ν,μ)
end

let tol=1e-9
    ν = 2
    μx = [3,2,1]
    AA = makegnb(ν,μx)
    n = size(AA,1)
    Q,R = qr(randn(n,n))
    A = Q' * AA * Q
    B,V,ν1,μ = gnsd(A,tol)
    quiet || println("input ν=$ν μ'=$μx")
    quiet || println("ν=$ν1 μ=$μ")
    @test ν == ν1
    @test gnsdcheck(A,B,V,ν1,μ; quiet=quiet)
end

for TT in (Float64,Complex{Float64})
    let tol=1e-9
        ν=2
        μx = [3,2,1]
        AA = makegnb(ν,μx)
        n = size(AA,1)
        Q,R = qr(randn(n,n))
        A = Matrix{TT}(Q' * AA * Q)
        B,V,ν1,μ = gnsd(A,tol)
        quiet || println("input ν=$ν μ'=$μx")
        quiet || println("ν=$ν1 μ=$μ")
        @test ν == ν1
        @test gnsdcheck(A,B,V,ν1,μ; quiet=quiet)
    end
end

let tol=1e-9
    ν=2
    μx=[3,2,1]
    AA = makegnb(ν,μx;cplx=true)
    n = size(AA,1)
    Q,R = qr(randn(n,n))
    A = Q' * AA * Q
    B,V,ν1,μ = gnsd(A,tol)
    quiet || println("input ν=$ν μ'=$μx")
    quiet || println("ν=$ν1 μ=$μ")
    @test ν == ν1
    @test gnsdcheck(A,B,V,ν1,μ; quiet=quiet)
end

let tol=1e-29
    ν=1
    μx = [3,2]
    AA = BigFloat.(makegnb(1,μx))
    n = size(AA,1)
    Q,R = qr(BigFloat.(randn(n,n)))
    Ab = Q' * AA * Q
    B,V,ν1,μ = gnsd(Ab,tol)
    @test ν == 1
    @test μ == μx[1:ν]
    @test gnsdcheck(Ab,B,V,ν,μ)
end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# ReverseMcCormick.jl
# A Reverse McCormick operator library in Julia
# See https://github.com/PSORLab/ReverseMcCormick.jl
#############################################################################
# src/reverse_operators/hyperbolic.jl
# Contains definitions of reverse McCormick operators of sinh, cosh, tanh,
# sech, csch, coth
#############################################################################

"""
$(SIGNATURES)

Reverse McCormick operator for `sinh`.
"""
function sinh_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ asinh(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `cosh`.
"""
function cosh_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    y = y ∩ Interval{Float64}(1.0, Inf)
    x = x ∩ acosh(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `tanh`.
"""
function tanh_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    y = y ∩ Interval{Float64}(-1.0, 1.0)
    x = x ∩ atanh(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `sech`.
"""
function sech_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    y = y ∩ Interval{Float64}(0.0, 1.0)
    x = x ∩ asech(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `csch`.
"""
function csch_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ acsch(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `coth`.
"""
function coth_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ coth(y)
    y, x
end


# Raytracer.jl
# Raytracing for the generation of photorealistic images in Julia
# Copyright (c) 2021 Samuele Colombo, Paolo Galli

# Sphere

"""
    Sphere <: SimpleShape

A [`SimpleShape`](@ref) representing a sphere.

This is a unitary sphere centered in the origin. A generic sphere can be specified by applying a [`Transformation`](@ref).

# Members

- `transformation::Transformation`: the `Transformation` associated with the sphere.
- `material::Material`: the [`Material`](@ref) of the spere.
"""
Base.@kwdef struct Sphere <: SimpleShape
    transformation::Transformation = Transformation()
    material::Material = Material()
end

@doc """
    Sphere(transformation::Transformation, material::Material)

Constructor for a [`Sphere`](@ref) instance.
""" Sphere(::Transformation, ::Material)

@doc """
    Sphere(transformation::Transformation = Transformation(),
           material::Material = Material())

Constructor for a [`Sphere`](@ref) instance.
""" Sphere(; ::Transformation, ::Material)

function get_all_ts(::Type{Sphere}, ray::Ray)
    # compute intersection
    origin_vec = convert(Vec, ray.origin)
    a = norm²(ray.dir)
    b = 2f0 * origin_vec ⋅ ray.dir
    c = norm²(origin_vec) - 1f0
    Δ = b^2 - 4f0 * a * c
    Δ < 0 && return Vector{Float32}()
    sqrt_Δ = sqrt(Δ)
    t_1 = (-b - sqrt_Δ) / (2f0 * a)
    t_2 = (-b + sqrt_Δ) / (2f0 * a)

    # @assert isfinite(t_1)
    # @assert isfinite(t_2)

    return [t_1, t_2]
end

function get_t(::Type{Sphere}, ray::Ray)
    # compute intersection
    origin_vec = convert(Vec, ray.origin)
    a = norm²(ray.dir)
    b = 2f0 * origin_vec ⋅ ray.dir
    c = norm²(origin_vec) - 1f0
    Δ = b^2 - 4f0 * a * c
    Δ < 0 && return Inf32
    sqrt_Δ = sqrt(Δ)
    t_1 = (-b - sqrt_Δ) / (2f0 * a)
    t_2 = (-b + sqrt_Δ) / (2f0 * a)
    # nearest point
    if ray.tmin < t_1 < ray.tmax
        return t_1
    elseif ray.tmin < t_2 < ray.tmax
        return t_2
    else
        return Inf32
    end
end

function get_uv(::Type{Sphere}, point::Point)
    u = atan(point[2], point[1]) / (2f0 * π)
    u = u >= 0 ? u : u+1f0
    v = acos(clamp(point[3], -1f0, 1f0)) / π
    Vec2D(u, v)
end

function get_normal(::Type{Sphere}, point::Point, ray::Ray)
    normal = convert(Normal, point) |> normalize
    (normal ⋅ ray.dir < 0) ? normal : -normal
end


using Cairo
using LSystems
using DataStructures

c = CairoRGBSurface(1024, 896);
cr = CairoContext(c);

save(cr);

set_source_rgb(cr, 1.0, 1.0, 1.0);
rectangle(cr, 0.0, 0.0, 1024.0, 896.0);
fill(cr);

restore(cr);
save(cr);

x, y, a, d = 64., 896., -pi / 6, 0.8
set_line_width(cr, 0.6);
move_to(cr, x, y);

plant_start = "X"
plant_trans = Dict([('X', "F-[[X]+X]+F[+FX]-X"), ('F', "FF")])

function determine_new_position(x, y, a, d)
    x -= sin(a) * d
    y -= cos(a) * d
    return x, y
end

s = Stack(Vector{Float64}) # [x, y, a]

level = 9
drawing_instructions = @task lindenmayer(plant_start, plant_trans, level)

for v in drawing_instructions
    if v == 'F'
        x, y = determine_new_position(x, y, a, d)
        line_to(cr, x, y)
    elseif v == '+'
        a -= pi / 7
    elseif v == '-'
        a += pi / 7
    elseif v == '['
        push!(s, [x, y, a])
    elseif v == ']'
        x, y, a = pop!(s)
        move_to(cr, x, y)
    end
end

stroke(cr);
write_to_png(c, "plant_$level.png");


module IntervalLinearAlgebra

using StaticArrays, Requires, Reexport
using LinearAlgebra: checksquare

import Base: *
import CommonSolve: solve
import IntervalArithmetic: mid

function  __init__()
    # @require IntervalConstraintProgramming = "138f1668-1576-5ad7-91b9-7425abbf3153" include("linear_systems/oettli_nonlinear.jl")
    @require LazySets = "b4f0291d-fe17-52bc-9479-3d1a343d9043" include("linear_systems/oettli_linear.jl")

    set_multiplication_mode(config[:multiplication])
end

@reexport using LinearAlgebra, IntervalArithmetic

const IA = IntervalArithmetic

export
    set_multiplication_mode, get_multiplication_mode,
    LinearKrawczyk, Jacobi, GaussSeidel, GaussianElimination, HansenBliekRohn, NonLinearOettliPrager, LinearOettliPrager,
    NoPrecondition, InverseMidpoint, InverseDiagonalMidpoint,
    solve, enclose, epsilon_inflation,
    comparison_matrix, interval_norm, interval_isapprox, Orthants,
    is_H_matrix, is_strongly_regular, is_strictly_diagonally_dominant, is_Z_matrix, is_M_matrix,
    rref,
    eigenbox, Rohn, Hertz, verify_eigen, bound_perron_frobenius_eigenvalue

include("linear_systems/enclosures.jl")
include("linear_systems/precondition.jl")
include("linear_systems/solve.jl")
include("linear_systems/verify.jl")
include("linear_systems/oettli.jl")
include("multiplication.jl")
include("utils.jl")
include("classify.jl")
include("rref.jl")

include("eigenvalues/interval_eigenvalues.jl")
include("eigenvalues/verify_eigs.jl")
end


using OpenQuantumBase, Test

f(s) = (1 - s) * σx + s * σz
H = hamiltonian_from_function(f)

@test H(0.5) == 0.5 * (σx + σz)

cache = get_cache(H)
update_cache!(cache, H, 2.0, 0.5)
@test cache == -0.5im * (σx + σz)


module MyinterpolationA
export myinterpolationA

function myinterpolationA(grid, vals)
    function func(x)
        if x < grid[1] || x > grid[length(grid)]
            return "error"
        end

        num_2 = searchsortedfirst(grid, x)
        if num_2 == 1
            return vals[num_2]
        end

        num_1 = num_2 - 1
        x_1 = grid[num_1]
        x_2 = grid[num_2]
        y_1 = vals[num_1]
        y_2 = vals[num_2]

        y = ((y_2- y_1)/(x_2 - x_1))*(x - x_1) + y_1
        return y
    end
end

end


function xiny(x,y,opt::Int=1)
# Determine if each row of x is a member of y
# If row j of x is a member of y and x(j,:) = y(k,:) then in[j] = k
# Else in[j] = 0

if opt == 1
    m,dim = size(x);
else
    dim,m = size(x);
end

in = zeros(Int,m,1);

if opt==1
    if dim==1
        for j=1:m
            d2 = (y[:,1].-x[j,1]).^2;
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    elseif dim==2
        for j=1:m
            d2 = (y[:,1].-x[j,1]).^2 + (y[:,2].-x[j,2]).^2;
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    elseif dim==3
        for j=1:m
            d2 = (y[:,1].-x[j,1]).^2 + (y[:,2].-x[j,2]).^2 + (y[:,3].-x[j,3]).^2;
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    else
        n = size(y,1);
        for j=1:m
            d2 = sum((y - repmat(x[j,:],[n 1])).^2,2);
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    end
else
    if dim==1
        for j=1:m
            d2 = (y[1,:].-x[1,j]).^2;
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    elseif dim==2
        for j=1:m
            d2 = (y[1,:].-x[1,j]).^2 + (y[2,:].-x[2,j]).^2;
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    elseif dim==3
        for j=1:m
            d2 = (y[1,:].-x[1,j]).^2 + (y[2,:].-x[2,j]).^2 + (y[3,:].-x[3,j]).^2;
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    else
        n = size(y,2);
        for j=1:m
            d2 = sum((y - repmat(x[:,j],[1 n])).^2,1);
            md,id = findmin(d2);
            if md<1e-12
                in[j]=id;
            end
        end
    end
end

return in;

end


export qft_circuit

"""
    cphase(nbits, i, j, θ)

Control-phase gate.
"""
cphase(nbits::Int, i::Int, j::Int, θ::T) where T = control(nbits, i, j=>shift(θ))

"""
    qft_circuit(n)

The quantum Fourer transformation (QFT) circuit.

References
------------------------
* [Wiki](https://en.wikipedia.org/wiki/Quantum_Fourier_transform)
"""
qft_circuit(n::Int) = chain(n, hcphases(n, i) for i = 1:n)
hcphases(n, i) = chain(n, i==j ? put(i=>H) : cphase(n, j, i, 2π/(2^(j-i+1))) for j in i:n);


function hessian_matrix(poly, x)
    h = hessian(poly, x)
    return [
        h[1] h[2]
        h[2] h[3]
    ]
end

function saye_newton_iterate(
    xguess,
    xq,
    func,
    grad,
    hess,
    cellmap,
    tol,
    r;
    maxiter = 20,
    condtol = 1e5eps(),
)
    dim = length(xguess)
    jac = jacobian(cellmap)

    x0 = copy(xguess)
    gp = grad(x0)
    l0 = gp' * ((xq - cellmap(x0)) .* jac) / (gp' * gp)

    x1 = copy(x0)
    l1 = l0

    for counter = 1:maxiter
        vp = func(x0)
        gp = grad(x0)
        hp = hess(x0)

        gf = vcat(((cellmap(x0) - xq) .* jac) + l0 * gp, vp)
        hf = [
            diagm(jac .^ 2)+l0*hp gp
            gp' 0.0
        ]

        if inv(cond(hf)) > condtol
            δ = hf \ gf
            normδx = norm(δ[1:dim])
            if normδx > 0.5r
                δ *= 0.5r / normδx
            end
            x1 = x0 - δ[1:dim]
            l1 = l0 - δ[end]
        else
            error("Chopp method not implemented")
        end

        if norm(x1 - xguess) > r
            error("Did not converge in ball of radius $r")
        elseif norm(x1 - x0) < tol
            return x1
        else
            x0 = x1
            l0 = l1
        end
    end
    error("Did not converge in $maxiter iterations")
end

function project_on_zero_levelset(
    xguess,
    func,
    grad,
    tol,
    r;
    maxiter = 20,
    normtol = 1e-8,
    perturbation = 0.1,
)
    x0 = copy(xguess)
    x1 = copy(x0)
    dim = length(xguess)

    for counter = 1:maxiter
        vf = func(x0)
        gf = grad(x0)

        if norm(gf) < normtol
            x0 = x0 + perturbation * (rand(dim) .- 0.5)
        else
            δ = vf / (gf' * gf) * gf
            normδ = norm(δ)
            if normδ > 0.5r
                δ *= 0.5r / normδ
            end

            x1 = x0 - δ

            if abs(func(x1)) < tol
                flag = norm(x1 - xguess) < r
                return x1, flag
            else
                x0 = x1
            end
        end
    end
    error("Did not converge after $maxiter iterations")
end

function reference_seed_points(n)
    @assert n > 0
    xrange = range(-1.0, stop = 1.0, length = n + 2)
    points = ImplicitDomainQuadrature.tensor_product_points(xrange[2:n+1]', xrange[2:n+1]')
end

function seed_cell_zero_levelset(xguess, func, grad; tol = 1e-12, r = 2.5)
    dim, nump = size(xguess)
    pf = [project_on_zero_levelset(xguess[:, i], func, grad, tol, r) for i = 1:nump]
    flags = [p[2] for p in pf]
    valididx = findall(flags)
    validpoints = [p[1] for p in pf[valididx]]
    return hcat(validpoints...)
end

function seed_zero_levelset(nump, levelset, levelsetcoeffs, cutmesh)
    refpoints = reference_seed_points(nump)
    refseedpoints = []
    spatialseedpoints = []
    seedcellids = Int[]
    cellsign = cell_sign(cutmesh)
    cellids = findall(cellsign .== 0)
    for cellid in cellids
        cellmap = cell_map(cutmesh, cellid)
        nodeids = nodal_connectivity(cutmesh.mesh, cellid)
        update!(levelset, levelsetcoeffs[nodeids])

        xk = seed_cell_zero_levelset(refpoints, levelset, x -> vec(gradient(levelset, x)))

        numseedpoints = size(xk)[2]

        append!(seedcellids, repeat([cellid], numseedpoints))
        push!(refseedpoints, xk)
        push!(spatialseedpoints, cellmap(xk))
    end

    refseedpoints = hcat(refseedpoints...)
    spatialseedpoints = hcat(spatialseedpoints...)
    return refseedpoints, spatialseedpoints, seedcellids
end

function seed_zero_levelset_with_interfacequads(interfacequads, cutmesh)
    cellsign = cell_sign(cutmesh)
    cellids = findall(cellsign .== 0)

    totalnumqps = sum([length(interfacequads[1, cellid]) for cellid in cellids])

    refseedpoints = zeros(2, totalnumqps)
    spatialseedpoints = zeros(2, totalnumqps)
    seedcellids = zeros(Int, totalnumqps)

    start = 1
    for cellid in cellids
        cellmap = cell_map(cutmesh, cellid)
        refpoints = interfacequads[1, cellid].points
        spatialpoints = cellmap(refpoints)
        numqps = size(refpoints)[2]

        stop = start + numqps - 1

        refseedpoints[:, start:stop] = refpoints
        spatialseedpoints[:, start:stop] = spatialpoints
        seedcellids[start:stop] = repeat([cellid], numqps)

        start = stop + 1
    end
    return refseedpoints, spatialseedpoints, seedcellids
end

function collect_interface_normals(interfacequads,cutmesh)
    cellsign = cell_sign(cutmesh)
    cellids = findall(cellsign .== 0)

    totalnumqps = sum([length(interfacequads[1,cellid]) for cellid in cellids])

    interfacenormals = zeros(2,totalnumqps)
    start = 1

    for cellid in cellids
        normals = interface_normals(interfacequads,cellid)
        numqps = size(normals)[2]

        stop = start + numqps - 1

        interfacenormals[:,start:stop] = normals
        start = stop + 1
    end
    return interfacenormals
end

function closest_reference_points_on_zero_levelset(
    querypoints,
    refseedpoints,
    spatialseedpoints,
    seedcellids,
    levelset,
    levelsetcoeffs,
    mesh,
    tol,
    boundingradius,
)

    dim, numquerypoints = size(querypoints)
    refclosestpoints = zeros(dim, numquerypoints)
    refclosestcellids = zeros(Int, numquerypoints)
    refgradients = zeros(dim, numquerypoints)

    tree = KDTree(spatialseedpoints)
    seedidx, seeddists = nn(tree, querypoints)

    for (idx, sidx) in enumerate(seedidx)
        xguess = refseedpoints[:, sidx]
        xquery = querypoints[:, idx]
        guesscellid = seedcellids[sidx]
        cellmap = cell_map(mesh, guesscellid)
        update!(levelset, levelsetcoeffs[nodal_connectivity(mesh, guesscellid)])

        refcp = saye_newton_iterate(
            xguess,
            xquery,
            levelset,
            x -> vec(gradient(levelset, x)),
            x -> hessian_matrix(levelset, x),
            cellmap,
            tol,
            boundingradius,
        )

        refclosestpoints[:, idx] = refcp
        refclosestcellids[idx] = guesscellid
        refgradients[:, idx] = gradient(levelset, refcp)
    end
    return refclosestpoints, refclosestcellids, refgradients
end

function distance_to_zero_levelset(
    querypoints,
    refseedpoints,
    spatialseedpoints,
    seedcellids,
    levelset,
    levelsetcoeffs,
    mesh,
    tol,
    boundingradius,
)

    dim, numquerypoints = size(querypoints)
    signeddistance = zeros(numquerypoints)

    refclosestpoints, refclosestcellids, refgradients =
        closest_reference_points_on_zero_levelset(
            querypoints,
            refseedpoints,
            spatialseedpoints,
            seedcellids,
            levelset,
            levelsetcoeffs,
            mesh,
            tol,
            boundingradius,
        )

    for i = 1:numquerypoints
        refcp = refclosestpoints[:, i]
        cellmap = cell_map(mesh, refclosestcellids[i])

        spatialcp = cellmap(refcp)
        xquery = querypoints[:, i]

        g = vec(transform_gradient(refgradients[:, i]', jacobian(cellmap)))
        s = sign(g' * (xquery - spatialcp))

        signeddistance[i] = s * norm(spatialcp - xquery)
    end

    return signeddistance
end


module SFS

export sfs, sfs_TVS, sfs_CV

include("../data.jl")
using .DataUtil
using Statistics, LinearAlgebra
using InvertedIndices

# forward stepwise feature selection with step size k
function sfs(𝐗, 𝐲, k)
    @assert k ≤ size(𝐗, 2)

    𝐴 = Array{Int}(undef, k)
    𝐬 = Array{Float64}(undef, k)
    for t = 1:k
        𝐴_t = 𝐴[1:(t - 1)]
        𝐴ᶜ_t = collect(1:size(𝐗, 2))[Not(𝐴_t)]
        𝐗_𝐴 = 𝐗[:, 𝐴_t]
        𝐫 = 𝐲 - 𝐗_𝐴 / (𝐗_𝐴' * 𝐗_𝐴) * (𝐗_𝐴' * 𝐲)
        j_t = 𝐴ᶜ_t[argmax(abs.(𝐗[:, 𝐴ᶜ_t]' * 𝐫))] # new feature
        𝐴[t] = j_t
        𝐬[t] = float(sign(𝐗[:,j_t]' * 𝐫))
    end
    return (𝐴, 𝐬)
end

# train-val-split for determining the hyperparameter "k" in candidate list K
function sfs_TVS(𝐗, 𝐲, K; ratio=0.5)
    𝐗_train, 𝐗_val = train_val_split(𝐗; ratio=ratio, do_centering=true)
    𝐲_train, 𝐲_val = train_val_split(𝐲; ratio=ratio, do_centering=true)
    k = K[argmin((k -> validation_error(k, 𝐗_train, 𝐗_val, 𝐲_train, 𝐲_val)).(K))]
    return (sfs(𝐗, 𝐲, k)..., k)
end

# cross-validation for determining the hyperparameter "k" in candidate list K
function sfs_CV(𝐗, 𝐲, K; 𝑘=length(𝐲))
    k = K[argmin((k -> cv_error(k, 𝐗, 𝐲, 𝑘)).(K))]
    return (sfs(𝐗, 𝐲, k)..., k)
end

function validation_error(k, 𝐗_train, 𝐗_val, 𝐲_train, 𝐲_val)
    𝐴_k, _ = sfs(𝐗_train, 𝐲_train, k)
    𝐗_train_𝐴 = 𝐗_train[:, 𝐴_k] # active train input
    𝐗_val_𝐴 = 𝐗_val[:, 𝐴_k] # active validation input
    𝐗₊ = (𝐗_train_𝐴' * 𝐗_train_𝐴) \ 𝐗_train_𝐴' # pseudo inverse matrix of 𝐗_train_𝐴
    𝐞 = 𝐲_val - 𝐗_val_𝐴 * 𝐗₊ * 𝐲_train
    return 𝐞' * 𝐞
end

function cv_error(k, 𝐗, 𝐲, 𝑘)
    𝑘_fold = KFold(length(𝐲), 𝑘) # 𝑘-fold Cross Validation
    return mean(
        begin
        𝐗_train, 𝐗_val = train_val_split(𝐗, 𝑘_fold.parts[i], 𝑘_fold.parts[i + 1] - 1)
        𝐲_train, 𝐲_val = train_val_split(𝐲, 𝑘_fold.parts[i], 𝑘_fold.parts[i + 1] - 1)
        validation_error(k, 𝐗_train, 𝐗_val, 𝐲_train, 𝐲_val)
    end for i = 1:𝑘)
end

end # module SFS