# %%
include("CellularBase.jl")
using .CellularBase
println("using CB")
@time using Plots
println("using Plots")
@time using Images
println("using Images")
# %%
mutable struct LifeGrid <: AbstractGrid{Int}
    state::Array{Int}
    otherstate::Array{Int}
    bc::BoundaryCondition
    neighborhood::AbstractNeighborhood
    τi::Int64
    τr::Int64
    max::Float64
    min::Float64
    function LifeGrid(
            size::Tuple{Int,Int}, τi::Int, τr::Int;
            max::Float64, min::Float64, init=nothing)
        z = zeros(size)
        if isnothing(init) z[(size .÷ 2)...] = 1
        else z[rand(CartesianIndices(z), init)] .= 1
        end
        new(z, copy(z), FixedMax, VonNeumannNeighborhood(3, 2), τi, τr, max, min)
    end
end

CellularBase.possible_states(l::LifeGrid) = 0:(l.τi + l.τr)

CellularBase.newstate(grid::LifeGrid) = grid.otherstate;

function CellularBase.state!(grid::LifeGrid, newstate)
    grid.otherstate = grid.state
    grid.state = newstate
end
# %%
infect(fraction::Float64, g::LifeGrid) = rand() < max(((g.max - g.min) * fraction + g.min), 0)
susceptible(i::Int, g::LifeGrid) = i == 0
infectious(i::Int, g::LifeGrid) = i ∈ 1:g.τi
refactory(i::Int, g::LifeGrid) = i ∈ g.τi:(g.τr + g.τi)
recovered(i::Int, g::LifeGrid) = i == (g.τr + g.τi)
function (g::LifeGrid)(neighbors::Array{Int}; kwargs...)
	current = neighbors[1]
    if susceptible(current, g)
        fraction = count(i -> infectious(i, g), neighbors) / length(neighbors)
        if infect(fraction, g)
            return 1
        else
            return 0
        end
    else
        return (current + 1) % (g.τi + g.τr)
    end
end
function color(i, grid)
    if susceptible(i, grid)
        RGB(0.9, 0.9, 0.9)
    elseif infectious(i, grid)
        RGB(1., 0., 0.)
    else
        RGB(0., 0., 0.)
    end
end;
# %%
function lockdown(grid, step; step_dict, radii)
    push!(radii, radius(grid.neighborhood))
    if step in keys(step_dict)
        grid.neighborhood = VonNeumannNeighborhood(step_dict[step], 2);
    end
end
step_dict = Dict(50 => 2, 100 => 1, 200 => 2, 300 => 3)
radii = []
grid = LifeGrid((200, 200), 4, 9, min=0., max=0.5, init=5)
k = simulate!(grid, 400; postrunhook=lockdown, step_dict=step_dict, radii=radii);
# %%
begin
    gif = @gif for frame ∈ k
        plot(frame .|> p -> color(p, grid))
    end every 2
    display(gif)
end
# %%
@time begin
    p = plot(1:length(k), k .|> frame -> count(frame .|> p -> susceptible(p, grid)), color=:green, label="Sus")
    plot!(1:length(k), k .|> frame -> count(frame .|> p -> infectious(p, grid)), color=:red, label="Inf")
    plot!(1:length(k), k .|> frame -> count(frame .|> p -> refactory(p, grid)), color=:blue, label="Ref")
    q = plot(radii, label="Neighbourhood Radius", legend=:bottomright)
    plot(p, q, size=(600, 500), layout=(2, 1))
end
# %%


using Pkg;
Pkg.activate(@__DIR__)
using CairoMakie, AbstractPlotting;
using GeometryBasics
using FromFile
#%%
Pkg.activate(joinpath(@__DIR__, "..") )
@from "../src/GAIO.jl" using GAIO
#%%

lb = fill(-2.0,2)
ub = fill(2.0,2)

f( x :: Vector{ R } where R<:Real ) = [
    sum( ( x .- 1 ).^2 );
    sum( ( x .+ 1 ).^2 )
]

#%%
root = GAIO.init_box_tree(fill(-2.0, 2), fill(2.0,2));

box_collection = [ root, ];
function progress!( box_collection, f :: F where F<:Function )

    # Subdivision Step
    new_collection = BoxNode[];
    for box ∈ box_collection
        subdivide!(box; method=:cycle)
        push!( new_collection, children(box)... )
    end

    # Selection Step

    ## 1) generate test sites and evaluate
    for box ∈ new_collection
        generate_sites!( box; method = :edges_center, num_edge_points = 3 )
        eval!( box, f )
    end

    ## 2) Discrete Hit Test

    ## 2a) Find all currently nondominated points
    all_points = vcat([ results(box) for box ∈ new_collection ]...)
    nondom_points = Vector{Float64}[];
    for X ∈ all_points
        X_nondominated = true
        for Y ∈ all_points
            if Y ≺ X
                X_nondominated = false
                break;
            end
        end
        if X_nondominated
            push!(nondom_points, X)
        end
    end

    ## 2b) Select Boxes containing a Nondom Point
    good_box_indices = Int[];
    for (i,boxᵢ) ∈ enumerate(new_collection)
        for Y in results(boxᵢ)
            if Y ∈ nondom_points
                push!(good_box_indices, i)
                break;
            end
        end
    end

    # TODO: Make more robust by reintroducing boxes

    empty!( box_collection );
    push!( box_collection, new_collection[good_box_indices]... );
end

#%%
for i = 1 : 15
    progress!( box_collection, f )
end

#%%
# Plotting …
# should work in 3D and 2D

function box2rect( box )
    rect = Rect( Vec( GAIO.lb( box )... ), Vec( GAIO.width(box)... ) );
    rect_v = decompose( Point{ GAIO.n_vars( box), Float64 }, rect )
    rect_f = decompose( TriangleFace{Int}, rect );
    return rect_v, rect_f
end

scene, layout = layoutscene();
ax = layout[1,1] = Axis(scene);
for box ∈ box_collection
    mesh!(ax, box2rect(box)...; color=:blue, strokecolor=:lightblue, strokewidth = 5.0, shading=false)
end
ax.aspect = AxisAspect(1)
scene

#


mutable struct Gaussian{T, RT <: Real, O} <: AbstractBlock
    μ::T
    σ::RT
    grad_μ::T
    grad_σ::RT

    input::O
    output::O

    Gaussian{T, RT, O}(mu, sigma) where {T, RT, O} =
        new{T, RT, O}(mu, sigma, zero(T), zero(RT))
end

Gaussian(::Type{O}, mu::T, sigma::RT) where {T, RT, O} =
    Gaussian{T, real(T), O}(mu, sigma)

Gaussian(mu::T, sigma::RT; nbatch=1) where {T, RT} =
    Gaussian(nbatch > 1 ? Matrix{T} : Vector{T}, mu, sigma)

# TODO: check if generated function will be faster (probably not)
function forward(op::Gaussian{T, RT, O}, input::O) where {T, RT, O}
    op.input = input
    factor = 1 / (sqrt(2 * pi) * op.σ)
    op.output = factor .* exp.(- abs2.(input .- op.μ) ./ (2 * op.σ^2))
    return op.output
end

function backward(op::Gaussian{T, RT, O}, grad::O) where {T <: Complex, RT, O}
    op.grad_μ = sum(real(op.input .- op.μ) .* op.output ./ op.σ^2 .* grad) # Trust your compiler
    op.grad_σ = sum(((abs2(op.input .- op.μ)) .- op.σ^2) ./ op.σ^3 .* op.output .* grad)
    return -(conj(op.input .- op.μ) .* op.output) ./ op.σ^2 .* grad
end


using Test
using ParametricSurfaces
using ParametricSurfaces: WavePropBase.Geometry

@testset "Disk" begin
    disk = ParametricSurfaces.Circle() # abstract entity
    Γ = boundary(disk) |> Domain
    M = meshgen(Γ,(10,))
    # plot(M,Γ)
    @test entities(Γ) == boundary(disk)
    @test geometric_dimension(disk) == 2
end

@testset "Ball" begin
    ball = ParametricSurfaces.Sphere() # abstract entity
    Γ = boundary(ball) |> Domain
    M = meshgen(Γ,(1,1))
    # plot(M,Γ)
    @test entities(Γ) == boundary(ball)
    @test geometric_dimension(ball) == 3
end


struct SpectralNormalizer{U,V}
    left_singular_vector::U
    right_singular_vector::V
    iterations::Int
    norm_multiplier::Float32
end

function SpectralNormalizer(
    kernel_shape;
    iterations::Int = 1,
    norm_multiplier::Float32 = 0.95,
    init = glorot_uniform,
    device = cpu,
)
    # Necessary shapes
    # Expecting this to be of shape (out_dim, in_dim)
    # For Conv layers, we expect the shape to be (out_channels, :)
    # U and V vectors
    right_singular_vector = device(init((kernel_shape[1], 1)))
    left_singular_vector = device(init((prod(kernel_shape[2:end]), 1)))

    return SpectralNormalizer(
        left_singular_vector,
        right_singular_vector,
        iterations,
        norm_multiplier,
    )
end

function (snl::SpectralNormalizer)(kernel)
    # update
    v_hat = snl.left_singular_vector
    u_hat = snl.right_singular_vector
    for iteration = 1:snl.iterations
        # Update the left singular vector
        v_hat = u_hat' * kernel # shape - (1, in_dim)
        v_hat = v_hat ./ norm(v_hat) # shape - (1, in_dim)
        # Update the right singular vector
        u_hat = v_hat * kernel'
        u_hat = u_hat ./ norm(u_hat) # shape - (out_dim, 1)
    end

    # Update the left, right singular vectors in the struct
    Flux.Zygote.@ignore copyto!(snl.right_singular_vector, u_hat)
    Flux.Zygote.@ignore copyto!(snl.left_singular_vector, v_hat)

    sigma = u_hat * kernel * v_hat'
    # Bound the spectral norm
    bounded_kernel =
        all((snl.norm_multiplier ./ sigma) .< 1) ?
        (snl.norm_multiplier ./ sigma) .* kernel : kernel
    return bounded_kernel
end


"""
r = _create_grid_Greens_function(number_of_cells, resolution)\n
Computes the radius for the domain given by number_of_cells.\n
This radius is used to compute the weakened Greens functions.\n
"""
function _create_grid_Greens_function(number_of_cells, resolution)
    #get max dimensions
    dx,dy,dz = 2 .^(ceil.(Int, log2.(number_of_cells.+2).+1))
    lx = resolution[1].*collect(-dx/2f0:dx/2f0-1)
    ly = resolution[2].*collect(-dy/2f0:dy/2f0-1)
    lz = resolution[3].*collect(-dz/2f0:dz/2f0-1)

    # repeat the axes along the number corresponding dimensions
    X = repeat(lx, outer=(1, dy, dz))
    Y = repeat(reshape(ly,1,:), outer=(dx, 1, dz))
    Z = repeat(reshape(lz,1,1,:), outer=(dx, dy, 1))

    radius = @. sqrt(X^2 + Y^2 + Z^2)

    return radius
end

"""
Internal function that computes the greens function.
"""
function _create_Greens_function(radius, k_b, a)
    # sommerfeld radiation condition
    scalar_G = 3/(k_b*a)^2*(sin(k_b*a)/(k_b*a)-cos(k_b*a))
    # Green's function at r != 0
    G = @. exp.(-1im*k_b*radius)/(4*Float32(π)*radius)*scalar_G
    # Green's function at r = 0
    G₀ = 3/(4*Float32(π)*k_b^2*a^3)*((1+1im*k_b*a)*exp(-1im*k_b*a)-1)
    G = circshift(G, floor.(Int, size(G)./2))
    G[1,1,1] = G₀

    return fft(G)
end

"""
G = create_Greens_function(number_of_cells, resolution, k_b).\n
creates the green function on the larger grid.\n
"""
function create_Greens_function(number_of_cells, resolution, k_b)
    a = findmin(resolution)[1]/2
    radius = _create_grid_Greens_function(number_of_cells, resolution)
    G = _create_Greens_function(radius, k_b, a)

    return prod(resolution).*G
end

"""
_Greens_function_times_contrast_source!(vector_potential_1d, G, dielectric, E, w, vector_potential_3d, Ig, fft_tmp, planfft, planifft)\n
Using the Green's function and contrast source we compute the vector potential.\n
a = ∫GdielectricEdV is computed using an FFT.\n
or f = ∫GKdV is computed using an FFT.\n
"""
function _Greens_function_times_contrast_source!(
    vector_potential_1d,
    G,
    dielectric,
    electric_field,
    vector_potential_3d,
    Ig,
    fft_tmp,
    planfft,
    planifft
)
    # move E to larger domain for fft
    mul!(fft_tmp, Transpose(Ig), electric_field)
    # multiply with the contrast
    vector_potential_3d .= dielectric.*reshape(fft_tmp, size(G)..., 3)
    # do the fft of the contrast source (dielectric.*E)
    planfft*vector_potential_3d
    # multiply the green's function with the contrast source in freq domain
    @. vector_potential_3d = G*vector_potential_3d
    planifft*vector_potential_3d
    # map to smaller domain again
    mul!(vector_potential_1d, Ig, vec(vector_potential_3d))

    return nothing
end

"""
source_to_incident_field!(source_to_field, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, operator, fft_tmp, planfft, planifft, scaling)\n
Compute the incident field in a vacuum given a source distribution\n
source_to_field will contain the initial source distribution and will contain the incident field when the function is finished\n
Use when electric source to electric field or magnetic source to magnetic field.\n
For an electric current density distribution to incident electric field:\n
    source_to_field: J\n
    vector_potential_1d: a\n
    vector_potential_3d: A\n
    operator: AtoE\n
    scaling: 1/(1im*ω*ϵ₀)\n
For a magnetic source distribution to incident magnetic field:\n
    source_to_field: K\n
    vector_potential_1d: f\n
    vector_potential_3d: F\n
    operator: FtoH\n
    scaling: 1/(1im*ω*μ₀)\n
source_to_incident_field!(incident_field, source, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, operator, fft_tmp, planfft, planifft, scaling)\n
Use when electric source to magnetic field or magnetic source to electric field.\n
For a magnetic source distribution to incident electric field:\n
    source_to_field: K\n
    vector_potential_1d: f\n
    vector_potential_3d: F\n
    operator: FtoE\n
    scaling: 1/(1im*ω*μ₀)\n
For an electric current density distribution to incident magnetic field:\n
    source_to_field: J\n
    vector_potential_1d: a\n
    vector_potential_3d: A\n
    operator: AtoH\n
    scaling: 1/(1im*ω*ϵ₀)\n
"""
function source_to_incident_field!(
    source_to_field,
    vector_potential_1d,
    vector_potential_3d,
    G,
    dielectric,
    Ig,
    operator,
    fft_tmp,
    planfft,
    planifft,
    scaling
)
    _Greens_function_times_contrast_source!(vector_potential_1d, G, dielectric.vac, source_to_field, vector_potential_3d, Ig, fft_tmp, planfft, planifft)
    mul!(source_to_field, operator, vector_potential_1d)
    @. source_to_field *= scaling

    return nothing
end

function source_to_incident_field!(
    incident_field,
    source,
    vector_potential_1d,
    vector_potential_3d,
    G,
    dielectric,
    Ig,
    operator,
    fft_tmp,
    planfft,
    planifft,
    scaling
)
    _Greens_function_times_contrast_source!(vector_potential_1d, G, dielectric.vac, source, vector_potential_3d, Ig, fft_tmp, planfft, planifft)
    mul!(incident_field, operator, vector_potential_1d)
    @. incident_field *= scaling

    return nothing
end

"""
total_minus_scattered_electric_field!(v,total_electric_field,tmp,a,w,A,G,dielectric,Ig,AToE,fft_tmp,planfft,planifft)\n
calculate the incident electric field from the total electric field.\n
This can also be seen as Ax in the context of ||b-Ax||_2^2\n
"""
_total_minus_scattered_electric_field!(v, total_electric_field, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, planfft, planifft) = _total_minus_scattered_electric_field!(v, total_electric_field, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, planfft, planifft, 1)
function _total_minus_scattered_electric_field!(
    v,
    total_electric_field,
    vector_potential_1d,
    vector_potential_3d,
    G,
    dielectric,
    Ig,
    AToE,
    fft_tmp,
    planfft,
    planifft,
    scalar
)
    _Greens_function_times_contrast_source!(vector_potential_1d, G, dielectric.patient, total_electric_field, vector_potential_3d, Ig, fft_tmp, planfft, planifft)
    copyto!(v, total_electric_field)
    mul!(v, AToE, vector_potential_1d, -scalar, scalar)

    return nothing
end

"""
_total_minus_scattered_electric_field_for_residual!(v,total_electric_field,tmp,vector_potential_3d,w,vector_potential_3d,G,dielectric,Ig,AToE,fft_tmp,planfft,planifft,scalar)\n
similar to total_minus_scattered_electric_field!(), but now r-AΔx is computed\n
That is the change in the residual as a result of the update in x\n
"""
function _total_minus_scattered_electric_field_for_residual!(
    v,
    total_electric_field,
    vector_potential_1d,
    vector_potential_3d,
    G,
    dielectric,
    Ig,
    AToE,
    fft_tmp,
    planfft,
    planifft,
    scalar
)
    _Greens_function_times_contrast_source!(vector_potential_1d, G, dielectric.patient, total_electric_field, vector_potential_3d, Ig, fft_tmp, planfft, planifft)
    mul!(total_electric_field, AToE, vector_potential_1d, -scalar, scalar)
    axpy!(one(ComplexF32), total_electric_field, v)

    return nothing
end

"""
_set_source_values_in_domain!(sources, incident_electric_field, v, S)\n
will set the current density values to -source_values\n
Using the S matrix these are placed into the incident electric field.\n
"""
function _set_source_values_in_domain!(sources, incident_electric_field, source_values, S)
    # set source values to the vector we are multiplying A with
    sources .= -one(ComplexF32).*source_values
    # set the source values into the incident electric field (will be equal to Jinc before the function JIncToEinc!)
    mul!(incident_electric_field, S, sources)

    return nothing
end

"""
library_matrix_times_current_density!(current_density, nUpdates, sources, incident_electric_field, S, tmp, fft_tmp, G, vector_potential_3d, w, dielectric, vector_potential_3d, AToE, Ig, scaling, pfft, pifft, x, p, r, rt, u, v, q, uq).\n
calculate the scattered electric field given the current density.\n
"""
function library_matrix_times_current_density!(
    current_density,
    sources,
    incident_electric_field,
    S,
    fft_tmp,
    G,
    vector_potential_3d,
    dielectric,
    vector_potential_1d,
    AToE,
    Ig,
    scaling,
    pfft,
    pifft,
    x,
    p,
    r,
    rt,
    u,
    v,
    q,
    uq
)
    source_values_input = -current_density

    _set_source_values_in_domain!(sources, incident_electric_field, source_values_input, S)
    #from source in E incident to the actual incident electric field
    source_to_incident_field!(incident_electric_field, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, scaling)
    #compute total electric field
    cgs_efield!(incident_electric_field, fft_tmp, G, vector_potential_3d, dielectric, vector_potential_1d, AToE, Ig, pfft, pifft, x, p, r, rt, u, v, q, uq)

    return x
end

"""
library_matrix_times_current_density_gpu!(current_density, nUpdates, sources, incident_electric_field, S, tmp, fft_tmp, G, vector_potential_3d, w, dielectric, vector_potential_3d, AToE, Ig, scaling, pfft, pifft, x, p, r, rt, u, v, q, uq).\n
calculate the scattered electric field given the current density.\n
"""
function library_matrix_times_current_density_gpu!(
    current_density,
    sources,
    incident_electric_field,
    S,
    fft_tmp,
    G,
    vector_potential_3d,
    dielectric,
    vector_potential_1d,
    AToE,
    Ig,
    scaling,
    pfft,
    pifft,
    x,
    p,
    r,
    rt,
    u,
    v,
    q,
    uq
)
    source_values_input = -current_density |> cu

    _set_source_values_in_domain!(sources, incident_electric_field, source_values_input, S)
    #from source in E incident to the actual incident electric field
    source_to_incident_field!(incident_electric_field, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, scaling)
    #compute total electric field
    cgs_efield!(incident_electric_field, fft_tmp, G, vector_potential_3d, dielectric, vector_potential_1d, AToE, Ig, pfft, pifft, x, p, r, rt, u, v, q, uq)

    return x
end #TODO: check if cpu version/collect are required

"""
_volume_integral_equation_wrapper!(output, sources, incident_electric_field, C, input_vector, S, tmp, fft_tmp, G, vector_potential_3d, w, dielectric, vector_potential_3d, AToE, Ig, scaling, pfft, pifft, x, p, r, rt, u, v, q, uq)\n
In place version to calculate Av = (I-CSZ)v = v - CSZv.\n
-CSZv is calculated using the VIE method by setting the sources to -v and multiplying the edges with the corresponding C value.
"""
function _volume_integral_equation_wrapper!(
    output,
    sources,
    incident_electric_field,
    C,
    input_vector,
    S,
    fft_tmp,
    G,
    vector_potential_3d,
    dielectric,
    vector_potential_1d,
    AToE,
    Ig,
    scaling,
    pfft,
    pifft,
    x,
    p,
    r,
    rt,
    u,
    v,
    q,
    uq
)
    # set sources in the vacuum to the correct values
    _set_source_values_in_domain!(sources, incident_electric_field, input_vector, S)
    # calculate the incident electric field inside the vacuum
    source_to_incident_field!(incident_electric_field, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, scaling)
    # solve for the field inside the patient anatomy
    cgs_efield!(incident_electric_field, fft_tmp, G, vector_potential_3d, dielectric, vector_potential_1d, AToE, Ig, pfft, pifft, x, p, r, rt, u, v, q, uq)
    # calculate the (v - CSZv) vector for the minimization process
    mul!(output, Transpose(S), x)
    @. output *= C
    # output = input_vector - CSZinput_vector
    axpy!(one(ComplexF32), input_vector, output)

    return nothing
end

"""
total_minus_scattered_electric_field_map!(v, total_electric_field, vector_potential_3d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, planfft, planifft, scalar).\n
used as a map to solve for the total electric field.\n
"""
function total_minus_scattered_electric_field_map!(
    v,
    total_electric_field,
    vector_potential_1d,
    vector_potential_3d,
    G,
    dielectric,
    Ig,
    AToE,
    fft_tmp,
    planfft,
    planifft,
    scalar
)
    copyto!(v, total_electric_field)
    _Greens_function_times_contrast_source!(vector_potential_1d, G, dielectric.patient, v, vector_potential_3d, Ig, fft_tmp, planfft, planifft)
    mul!(v, AToE, vector_potential_1d, -scalar, scalar)

    return v
end


using Distributions
using Random

function sample_infection_time(;omicron=false)
    """
    Sample infection times for num individuals based on the generation
    interval distribution, Gamma(shape_gen, scale_gen).
    """

    (shape_gen, scale_gen) = (2.75, 1.00)
    (shape_gen_omicron, scale_gen_omicron) = (1.58, 1.32)

    shape = (1-omicron)*shape_gen + omicron*shape_gen_omicron
    scale = (1-omicron)*scale_gen + omicron*scale_gen_omicron

    infection_time = ceil(Int, rand(Gamma(shape, scale)))

    return infection_time
end

function sample_onset_time(;omicron=false)
    """
    Sample incubation times for num individuals based on incubation period
    distribution, Gamma(shape_inc, scale_inc).
    """

    (shape_inc, scale_inc) = (5.807, 0.948)
    (shape_inc_omicron, scale_inc_omicron) = (3.33, 1.34)

    shape = (1-omicron)*shape_inc + omicron*shape_inc_omicron
    scale = (1-omicron)*scale_inc + omicron*scale_inc_omicron

    onset_time = ceil(Int, rand(Gamma(shape, scale)))

    return onset_time
end

function set_simulation_constants(state)
    """
    Contains the assumptions for simulation parameters. This includes all the dynamical constants:
        - k = heterogeneity parameter
        - p_symp = probability of symptoms
        - γ = relative infectiousness of asymptomatic individuals
        - p_symp_given_detect = probability of symptoms given detection
        - p_asymp_given_detect = probability of being asymptomatic given detection
        - consistency_multiplier = chosen such that sim_cases < consistency_multiplier*actual cases
            results in cases being injected into the simulation. This is used to account
            for superspreading events after periods of low incidence.
    These values are stored in sim_constants which is a dictionary indexed by the parameter name and
    ultimately stored on the stack in a SimulationParameters object.
    """

    # overdispersion parameter
    k = 0.15
    # assumptions surrouding the probability of symptomatic,
    # relative infectiousness γ and the ratio of Reff (α's)
    p_symp = 0.4

    p_detect_given_symp_dict = Dict{String, Float64}(
        "NSW" => 0.7,
        "QLD" => 0.7,
        "SA" => 0.7,
        "TAS" => 0.7,
        "WA" => 0.7,
        "ACT" => 0.7,
        "NT" => 0.7,
        "VIC" => 0.7,
    )

    p_detect_given_asymp_dict = Dict{String, Float64}(
        "NSW" => 0.467,
        "QLD" => 0.467,
        "SA" => 0.467,
        "TAS" => 0.467,
        "WA" => 0.467,
        "ACT" => 0.467,
        "NT" => 0.467,
        "VIC" => 0.467,
    )

    γ = 0.5     # relative infectiousness of asymptomatic
    # solve the system α_s*ps + α_a(1-ps) = 1 with α_a = γ*α_s
    p_detect_given_symp = p_detect_given_symp_dict[state]
    p_detect_given_asymp = p_detect_given_asymp_dict[state]
    # prob of detection
    p_detect = p_symp*p_detect_given_symp + (1-p_symp)*p_detect_given_asymp
    # prob symptomatic given detect
    p_symp_given_detect = p_detect_given_symp*p_symp/p_detect

    α_s = 1/(p_symp + γ*(1-p_symp))
    α_a = γ * α_s
    consistency_multiplier = 5.0

    # prob of detecting an international import
    qi = 0.98
    # prior parametes for the import model
    prior_alpha = 0.5
    prior_beta = 0.2
    # ema smoothing factor
    ϕ = 0.1

    simulation_constants = Constants(
        k,
        p_symp,
        p_detect_given_symp,
        p_detect_given_asymp,
        p_detect,
        p_symp_given_detect,
        γ,
        α_s,
        α_a,
        qi,
        prior_alpha,
        prior_beta,
        ϕ,
        consistency_multiplier,
    )

    # mapping between types
    individual_type_map = IndividualTypeMap(1, 2, 3)

    return (simulation_constants, individual_type_map)
end


using PyPlot, JUDI, LinearAlgebra, FFTW, JOLI, DSP
using IterativeSolvers

using Shampoo

n = (201, 201)   # (x,y,z) or (x,z)
d = (5f0, 5f0)
o = (0f0, 0f0)

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

v0 = 2f0*ones(Float32,n)
v = deepcopy(v0)
v[101,101] = 2.1f0  # dirac perturbation

# Slowness squared [s^2/km^2]
m = convert(Array{Float32,2},(1f0 ./ v).^2)	# true model
m0 = convert(Array{Float32,2},(1f0 ./ v0).^2)
dm = vec(m-m0)

model = Model(n, d, o, m)
model0 = Model(n, d, o, m0)

# Model structure

dtS = 2f0
timeS = 2000f0
nt = Int64(timeS/dtS)+1

fmin = 10f0
fmax = 60f0
cfreqs = (fmin, fmin+5f0, fmax-5f0, fmax) # corner frequencies
wavelet = reshape(cfreq_wavelet(500, nt, dtS/1f3, cfreqs; edge=hanning),:,1)
#wavelet = ricker_wavelet(timeS, dtS, 0.03f0)

################################### Source/receiver geometry ######################################

# Set up receiver geometry

nrec = 512

xrec, zrec = circleShape(extentx / 2,extentz/2, extentz/2-5*d[2],nrec)
yrec = 0f0 #2d so always 0

# receiver sampling and recording time
timeR = timeS   # receiver recording time [ms]
dtR   = dtS
# Set up receiver structure
nsrc = 32

xsrc, zsrc = circleShape(extentx / 2,extentz/2, extentz/2-5*d[2],nsrc)
ysrc = range(0f0,stop=0f0,length=nsrc) #2d so always 0

recGeometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR, nsrc=nsrc)

srcGeometry = Geometry(convertToCell(xsrc), convertToCell(ysrc), convertToCell(zsrc); dt=dtS, t=timeS)

q = judiVector(srcGeometry, wavelet)

mute_matrix = zeros(Float32,n)

for i = 1:n[1]
	for j = 1:n[2]
		mute_matrix[i,j] = ((i-201)^2+(j-201)^2<190^2)
	end
end

M = MaskOp(n,mute_matrix)

# Set up info structure for linear operators
ntComp = get_computational_nt(recGeometry, model0)
info = Info(prod(n), nsrc, ntComp)

opt = Options(free_surface=false,isic=false)

Pr  = judiProjection(info, recGeometry)
F0  = judiModeling(info, model0; options=opt) # modeling operator background model
Ps = judiProjection(info, srcGeometry)
# Combined operators
F0 = Pr*F0*adjoint(Ps)

J = judiJacobian(F0, q)
P = DiffOp(nt,nsrc,nrec)

PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=8); PyPlot.rc("ytick", labelsize=8)
figure(); imshow(reshape(sqrt.(1f0./m), n[1], n[2])',extent=(0,extentx,extentz,0));
PyPlot.scatter(xrec,zrec,marker=".",color="red",label="receivers",s=1);
PyPlot.scatter(xsrc,zsrc,marker=".",color="white",label="sources",s=1);
legend(loc=3,fontsize=8);
xlabel("Lateral position [m]", fontsize=8);ylabel("Horizontal position [m]", fontsize=8);
fig = PyPlot.gcf();
title("Experimental setup")

######################################### Set up JUDI operators ###################################

d_lin = J*dm

dm1 = J'*d_lin

dm2 = J'*P*J*dm

figure();
subplot(2,1,1);
PyPlot.plot(wavelet);title("zoomed-in wavelet in time domain");
subplot(2,1,2);
PyPlot.plot(abs.(fftshift(fft(wavelet))));title("wavelet in frequency domain");

figure();
subplot(2,2,1)
imshow(dm1.data');title("J'*J*dm")
subplot(2,2,2)
imshow(dm2.data');title("J'*cumsum*J*dm")
subplot(2,2,3)
imshow(abs.(fftshift(fft(dm1.data')))/norm(abs.(fftshift(fft(dm1.data'))),Inf),cmap="jet",vmin=0,vmax=1)
subplot(2,2,4)
imshow(abs.(fftshift(fft(dm2.data')))/norm(abs.(fftshift(fft(dm2.data'))),Inf),cmap="jet",vmin=0,vmax=1)

maxit = 10

x1 = 0f0 .* dm
start1 = J'*P*d_lin
x1,his1 = cg!(x1,M'*J'*P*J*M,start1,maxiter=maxit,log=true,verbose=true)
x2 = 0f0 .* dm
start2 = J'*d_lin
x2,his2 = cg!(x2,M'*J'*J*M,start2,maxiter=maxit,log=true,verbose=true)
res1 = his1[:resnorm]
res2 = his2[:resnorm]

u1 = 20*log.(res1./norm(start1))
u2 = 20*log.(res2./norm(start2))

figure();
xlabel("iterations")
ylabel("ratio")
PyPlot.plot(res1/norm(start1),label="CG")
PyPlot.plot(res2/norm(start2),label="P-CG")
title("normalized least squares residual")
legend()

figure();
xlabel("iterations")
ylabel("dB")
PyPlot.plot(u1,label="CG")
PyPlot.plot(u2,label="P-CG")
title("normalized log-based least squares residual")
legend()

figure();
subplot(2,2,1)
imshow(reshape(M*x1,n)');title("CG")
subplot(2,2,2)
imshow(reshape(M*x2,n)');title("P-CG")
subplot(2,2,3)
imshow(abs.(fftshift(fft(reshape(M*x2,n)')))/norm(abs.(fftshift(fft(reshape(M*x2,n)'))),Inf),cmap="jet",vmin=0,vmax=1)
subplot(2,2,4)
imshow(abs.(fftshift(fft(reshape(M*x2,n)')))/norm(abs.(fftshift(fft(reshape(M*x2,n)'))),Inf),cmap="jet",vmin=0,vmax=1)


"""
Calculate the euler characteristic of solid, working on V, copEV, copFE matrices
"""
function test_eulero(V, copEV, copFE)
    if !isempty(V)
        vertex=size(V,1)
    else
        println("WARNING!!! V matrix is empty!!!")
    end
    if !isempty(copEV)
        edges=size(copEV,1)
    else
        println("WARNING!!! copEV matrix is empty!!!")
    end
    if !isempty(copFE)
        faces = size(copFE,1)
    else
        println("WARNING!!! copFE matrix is empty!!!")
    end
    bc=size(Lar.Arrangement.biconnected_components(copEV))[1]
    println("EULERIAN CHARACTERISTIC-->$(vertex-edges+faces) BICON COMPS==$(bc) VERTEX = $(vertex) EDGES = $(edges) FACES = $(faces)")
end
"""
check the euler characteristic of solid for a single face at time, working on V, copEV, copFE matrices
"""
function test_eulero_face4face(V, copEV, copFE)
    edges=Set{Int64}()
    vertices=Set{Int64}()
    counter=0
    dict=Dict()
    array=[]
    array2=[]
    ciEV=[]
    ciFE=[]
    if !isempty(V)
        vertex=size(V)[1]
    end
    if !isempty(copEV)
        ciEV=findall(!iszero,copEV);
    end
    if !isempty(copFE)
        ciFE=findall(!iszero,copFE);
    end
    for j in 1:size(copFE)[1]
        for i in ciFE
            if i[1]==j
                push!(array,[i[1],i[2]])
            end
        end
    end
    for j in 1:size(copEV)[1]
        for i in ciEV
            if i[1]==j
                push!(array2,[i[1],i[2]])
            end
        end
    end
    prev=0
    for x in array
        if prev!=x[1]
            if length(vertices)!==length(edges)
                println("WARNING !!! TEST FAILED")
            end
            vertices=Set{Int64}()
            edges=Set{Int64}()
        end
        for y in array2
            if x[2]==y[1]
                push!(edges,y[1])
                push!(vertices,y[2])
            end

        end
        prev=x[1]
    end
end


@doc raw"""
    p = characteristic_polynomial(g_ol)

Determine the characteristic polynomial associated with open loop transfer function `g_ol`.

The characteristic polynomial is $1+g_{ol}(s)$. The roots of the characteristic polynomial determine the character of the response of the closed loop system to bounded inputs.

# Arguments
* `g_ol::TransferFunction`: open loop transfer function

# Returns
a polynomial of type `Polynomial`

# Example
```jldoctest
g_ol = 4 / (s + 3) / (s + 2) / (s + 1)
characteristic_polynomial(g_ol)
# output
Polynomials.Polynomial(10.0 + 11.0*s + 6.0*s^2 + 1.0*s^3)
```
"""
function characteristic_polynomial(g_ol::TransferFunction)
    if g_ol.time_delay != 0.0
        error("not implemented for time delay")
    end
    # TODO wut if time delay
    return g_ol.denominator + g_ol.numerator
end


using Plots
push!(LOAD_PATH, "./src")
using Basis2DTri
using Basis1D
theme(:wong)
gr(aspect_ratio=1,legend=:bottomright,axis=nothing,border=:none,ticks=nothing,legendfontsize=14,xlims=[1.5,4],ylims=[0.8,3],
   markerstrokewidth=0,markersize=1)
r1D, w1D = gauss_quad(0,0,2)
Nfp = length(r1D) # number of points per face
e = ones(Nfp) # vector of all ones
z = zeros(Nfp) # vector of all zeros
rf = [r1D; -r1D; -e];
sf = [-e; r1D; -r1D];

r_tri,s_tri,_ = Basis2DTri.quad_nodes_tri(4)
plot([-1;-1].+3,[-1;1].+2,seriescolor=:black,lw=3,label = nothing)
plot!([-1;1].+3,[1;-1].+2,seriescolor=:black,lw=3,label = nothing)
plot!([-1;1].+3,[-1;-1].+2,seriescolor=:black,lw=3,label = nothing)
#scatter!(r_tri.+3,s_tri.+2,markersize=8,markercolor=:deepskyblue,markerstrokecolor=:black,markerstrokewidth=4,label = nothing)
scatter!(rf.+3,sf.+2,markersize=8,markershape=:rect,markercolor=:darkorange1,markerstrokecolor=:black,markerstrokewidth=4,label = nothing)

savefig("triangle_facequad_nodes.png")


#========================================================================================#
"""
	NBodies

A system of N interacting bodies moving in dim dimensions over time T in res timesteps.

Author: Niall Palfreyman, 24/3/2022.
"""
module NBodies

# Externally callable methods of NBodies
export NBody, add!, simulate

#-----------------------------------------------------------------------------------------
# Module types:

"""
	NBody

An NBody system capable of containing multiple (N) bodies that gravitationally interact
with each other.
"""
mutable struct NBody
	N							# Number of bodies
	T							# Duration of simulation
	res							# Timestep resolution
	xi							# Initial positions of bodies
	pi							# Initial momenta of bodies
	m							# Masses of bodies

	"Construct a new NBody"
	function NBody( T=40, resolution=20000)
		# Initialise all fields of the decoding apparatus:
		new(
			0,					# Initially no bodies in the system
			T,					# Duration
			resolution,			# Timestep resolution
			[],					# Initial positions
			[],					# Initial momenta
			[],					# Masses of bodies
		)
	end
end

#-----------------------------------------------------------------------------------------
# Module methods:

"""
	addbody!( nbody::NBody, xi, pi, m=1)

Add to the system a new body with initial position and momentum xi, pi, and with mass m.
"""
function addbody!( nbody::NBody, xi, pi, m=1)
	push!( nbody.xi, xi)
	push!( nbody.pi, pi)
	push!( nbody.m, m)
	nbody.N  += 1
end

"""
	relpos( locations)

A useful helper function that calculates relative positions of a vector x of locations.
"""
function relpos( locations)
	locPerBody = repeat(locations,1,length(locations))
	locPerBody - permutedims(locPerBody)
end

"""
	simulate( nbody::NBody)

Run a simulation of the given NBody system over duration divided into res timesteps.
"""
function simulate( nbody)
	dt = nbody.T / nbody.res
	t = 0:dt:nbody.T
	x = Vector{typeof(nbody.xi)}(undef,nbody.res+1)
	p = Vector{typeof(nbody.pi)}(undef,nbody.res+1)

	# Pre-calculate products of masses:
	G = 1
	Gmm = G * nbody.m * nbody.m'

	# Initialisation:
	x[1] = nbody.xi
	p[1] = nbody.pi

	# Simulation:
	for n = 1:nbody.res
		# Simulate timestep n:
		x[n+1] = x[n] + dt*p[n]./nbody.m

		diff = relpos(x[n])
		invSq = abs.(diff'.*diff) .^ (3/2)
		for i in 1:length(x[n]) invSq[i,i] += 1 end
		forceFROMEachSource = -Gmm .* diff ./ invSq
		forceONEachSource = sum( forceFROMEachSource, dims=2)
		p[n+1] = p[n] + dt * forceONEachSource[:]
	end

	(t,x,p)
end

#-----------------------------------------------------------------------------------------
using GLMakie

"""
	demo()

Demonstrate simulation of a simple 2-body problem.
"""
function demo()
	# Initialise the system:
	nb = NBody( 2.25, 10000)
	addbody!( nb, -1, 0)
	addbody!( nb,  1, 0)

	# Run the simulation:
	(t,x,p) = simulate(nb)

	# Display some fancy graphics:
	fig = Figure(resolution=(600, 600))
	ax = Axis(fig[1, 1], xlabel = "Time (t)", title = "N-body Trajectories")

	lines!(ax, t, map(el->el[1],x), color=:blue)		# Body 1
	lines!(ax, t, map(el->el[2],x), color=:red)			# Body 2
	display(fig)
end

end		# of NBodies


mutable struct SimulationSamples
  input::InputParameters
  #output::SimResult
  sampleddataDF::DataFrame
end

function run1simulationsample(IP::InputParameters, minclonesize, maxclonesize)

  sresult = tumourgrow_birthdeathsample(IP.b, IP.d, IP.Nmax, IP.μ; numclones = IP.numclones, s = IP.selection, tevent = IP.tevent, clonalmutations = IP.clonalmutations);

  return sresult, IP
end

function getsamples(;nsample = 1, nclones = 1, ploidy = 2, read_depth = 100.0, detectionlimit = 5 ./read_depth, clonalmutations = 100.0, μ = 10.0, d = 0.0, b = log(2), ρ = 0.0, Nmax = 10^3, samplesize = round(Int64, 0.1 * Nmax), s = repeat([1.0], inner = nclones), tevent = collect(1.0:0.5:100.0)[1:nclones], cellularity = 1.0, fixedmu = false)

  nclones == length(s) || error("Number of clones is $(nclones), size of selection coefficient array is $(length(s)), these must be the same size ")

  nclones == length(tevent) || error("Number of clones is $(nclones), size of selection coefficient array is $(length(tevent)), these must be the same size ")

  IP = InputParameters(nclones,
  Nmax,
  detectionlimit,
  ploidy,
  read_depth,
  clonalmutations,
  s,
  μ,
  b,
  d,
  tevent,
  ρ,
  cellularity,
  fixedmu)

  simresult, IP = run1simulationsample(IP, 0.0, 1.0)

  samples = sample(simresult.cells, samplesize * nsample, replace = false)

  samplesvec = []
  for i in 0:(nsample - 1)
      push!(samplesvec, samples[(i * samplesize + 1):(samplesize * (i+ 1))])
  end

  mutations = unique(CancerSeqSim.cellsconvert(samples)[1]);
  DF = DataFrame(mutations = 1:maximum(mutations))

  for i in 1:nsample
      DF[Symbol("sample$i")] = counts(CancerSeqSim.cellsconvert(samplesvec[i])[1], 1:maximum(mutations))
  end

  DFs = sampledhistmultiple(Array(DF[:sample1]), Array(DF[:mutations]), samplesize, 1)
  for i in 2:nsample
      DFtemp = sampledhistmultiple(Array(DF[Symbol("sample$i")]), Array(DF[:mutations]), samplesize, i, detectionlimit = IP.detectionlimit, ploidy = IP.ploidy, read_depth = IP.read_depth, cellularity = IP.cellularity)
      DFs = join(DFs, DFtemp, on = :mutations)
  end

  return SimulationSamples(IP, DFs), simresult, DF

end

function sampledhistmultiple(AF, mutations, samplesize, nsample; detectionlimit = 0.1, ploidy = 2.0, read_depth = 100.0, cellularity = 1.0)

    AF = AF./ploidy

    AF = AF .* cellularity

    ind = AF.>(detectionlimit * samplesize)
    AF = AF[ind]
    mutations = mutations[ind]

    samp_percent = read_depth/samplesize

    depth = rand(Binomial(samplesize,samp_percent), length(AF))

    samp_alleles = map((n, p) -> rand(Binomial(n, p)), depth, AF/samplesize)

    VAF = samp_alleles./depth

    DF = DataFrame(mutations = mutations)
    DF[Symbol("VAF$(nsample)")] = VAF

    return DF
end


using Distributions: Distributions
using UnitfulAtomic: austrip
using Random: Random

struct BoltzmannVelocityDistribution{T,S} <: VelocityDistribution
    dist::Distributions.MvNormal{T}
    size::NTuple{S,Int}
end

"""
    BoltzmannVelocityDistribution(temperature, masses)

Obtain velocities sampled from the Boltzmann distribution for given temperature and masses.

This is effectively just a multivariate normal distribution but when combined with the
[`DynamicalDistribution`](@ref) it ensure samples are generated with the correct size.
If using for ring polymer dynamics, ensure the temperature is given as nbeads*temperature.
"""
function BoltzmannVelocityDistribution(temperature, masses, size)
    dist = Distributions.MvNormal(sqrt.(austrip(temperature) ./ masses))
    BoltzmannVelocityDistribution(dist, size)
end

function Base.rand(rng::AbstractRNG, s::BoltzmannVelocityDistribution{T,S}) where {T,S}
    out = Array{T,S}(undef, s.size...)
    if S == 2
        @views for i=1:size(out, 1)
            Random.rand!(rng, s.dist, out[i,:])
        end
    elseif S == 3
        @views for i=1:size(out, 3)
            for j=1:size(out, 1)
                Random.rand!(rng, s.dist, out[j,:,i])
            end
        end
    end
    return out
end

select_item(x::BoltzmannVelocityDistribution, ::Integer, ::NTuple) = rand(x)


"""
    struct ComplexSpace <: EuclideanSpace{ℂ}

A standard complex vector space ℂ^d with Euclidean inner product and no additional
structure. It is completely characterised by its dimension and whether its the normal space
or its dual (which is canonically isomorphic to the conjugate space).
"""
struct ComplexSpace <: EuclideanSpace{ℂ}
  d::Int
  dual::Bool
end
ComplexSpace(d::Integer = 0; dual = false) = ComplexSpace(Int(d), dual)
function ComplexSpace(dim::Pair; dual = false)
    if dim.first === Trivial()
        return ComplexSpace(dim.second; dual = dual)
    else
        msg = "$(dim) is not a valid dimension for ComplexSpace"
        throw(SectorMismatch(msg))
    end
end
function ComplexSpace(dims::AbstractDict; kwargs...)
    if length(dims) == 0
        return ComplexSpace(0; kwargs...)
    elseif length(dims) == 1
        return ComplexSpace(first(dims); kwargs...)
    else
        msg = "$(dims) is not a valid dimension dictionary for ComplexSpace"
        throw(SectorMismatch(msg))
    end
end

# convenience constructor
Base.getindex(::ComplexNumbers) = ComplexSpace
Base.:^(::ComplexNumbers, d::Int) = ComplexSpace(d)

# Corresponding methods:
#------------------------
dim(V::ComplexSpace) = V.d
isdual(V::ComplexSpace) = V.dual
Base.axes(V::ComplexSpace) = Base.OneTo(dim(V))

Base.conj(V::ComplexSpace) = ComplexSpace(dim(V), !isdual(V))

Base.oneunit(::Type{ComplexSpace}) = ComplexSpace(1)
⊕(V1::ComplexSpace, V2::ComplexSpace) = isdual(V1) == isdual(V2) ?
    ComplexSpace(dim(V1)+dim(V2), isdual(V1)) :
    throw(SpaceMismatch("Direct sum of a vector space and its dual does not exist"))
fuse(V1::ComplexSpace, V2::ComplexSpace) = ComplexSpace(V1.d*V2.d)
flip(V::ComplexSpace) = dual(V)

infimum(V1::ComplexSpace, V2::ComplexSpace) = isdual(V1) == isdual(V2) ?
    ComplexSpace(min(dim(V1), dim(V2)), isdual(V1)) :
    throw(SpaceMismatch("Infimum of space and dual space does not exist"))
supremum(V1::ComplexSpace, V2::ComplexSpace) = isdual(V1) == isdual(V2) ?
    ComplexSpace(max(dim(V1), dim(V2)), isdual(V1)) :
    throw(SpaceMismatch("Supremum of space and dual space does not exist"))

Base.show(io::IO, V::ComplexSpace) = print(io, isdual(V) ? "(ℂ^$(V.d))'" : "ℂ^$(V.d)")


using DelimitedFiles
using WriteVTK
using MeshCore: AbsShapeDesc, P1, L2, L3
using MeshCore: VecAttrib
using MeshCore: shapedesc, nshapes,  IncRel
using MeshCore: ShapeColl
using LinearAlgebra: norm
using StaticArrays

"""
    L2blockx(xs::Vector{T}; intbytes = 8) where {T}

Generate a graded mesh on an interval.

Keyword argument `intbytes` controls the size of the integer indexes.

# Return
By convention the function returns an incidence relation (`connectivity`)
between the elements (`connectivity.left`) and the vertices
(`connectivity.right`). The geometry is stored as the attribute "geom" of the
vertices.
"""
function L2blockx(xs::Vector{T}; intbytes = 8) where {T}
	inttype = Int64
    if intbytes == 4
        inttype = Int32
    end
    nL = length(xs)-1;
    nnodes = (nL+1);
    ncells = (nL);
    xys = zeros(T, nnodes, 1);
    conns = zeros(inttype, ncells, 2);

    f=1;
    for i in 1:(nL+1)
        xys[f, 1] = xs[i]
        f=f+1;
    end

    gc=1;
    for i in 1:nL
       conns[gc, 1] = i
       conns[gc, 2] = (i+1)
       gc=gc+1;
    end

    C = inttype.(conns[1:gc-1, :]);
    N, TC = size(xys, 2), eltype(xys)
    locs =  VecAttrib([SVector{N, TC}(xys[i, :]) for i in 1:size(xys, 1)])
    vertices = ShapeColl(P1, length(locs), "vertices")
    vertices.attributes["geom"] = locs
    elements = ShapeColl(L2, size(C, 1), "elements")
    return IncRel(elements, vertices, C)
end

"""
    L2block(Length, nL; kwargs...)

Generate a quadrilateral mesh  of the 2D block.

See also: L2blockx
"""
function L2block(Length, nL; kwargs...)
    return L2blockx(collect(linearspace(0.0, Length, nL+1)); kwargs...);
end


"""
Create a random partition for a matrix `A`. Each node is assigned to a module at
random. If the smallest dimension of the matrix is *m*, there will be between 1
and *m* modules.

This method tends to work well for small to medium networks (and is faster than
`partition_lp`).
"""
function partition_random(A::Array{Int64, 2})
  # Smaller margin of A determines the number of modules
  min_part = minimum(size(A))
  # Generates a matrix with int(0) for the module partition
  S = zeros(Int64, sum(size(A)), min_part)
  # Pick modules at random and put them in the partition
  init_part = rand(1:1:min_part,sum(size(A)))
  for i = 1:length(init_part)
    S[i,init_part[i]] = 1
  end
  # Only non-empty modules are allowed
  # And it returns a Modular object
  return no_empty_modules!(Modular(A, S))
end


using .RationalFunctionFields: modular_reduction, rational_reconstruction, CRT,
                                CRT_reconstruction

@testset "reduction + reconstruction tests" begin
    modulo = 2^31 - 1

    FF = Sing.N_ZpField(modulo)

    xs  = [ Sing.QQ(rand(1:10000), rand(1:10000)) for _ in 1:10 ]
    xrs = modular_reduction(xs, FF)
    ys  = rational_reconstruction(xrs, BigInt(modulo))

    @test xs == ys

end

@testset "complex reduction tests" begin
    R, (x, y) = AA.PolynomialRing(Sing.QQ, ["x", "y"])

    modulo = 2^31 - 1
    FF = AA.GF(2^31 - 1)

    fs = [
        R(0),
        R(1),
        7x^2*y + 5x + 3y,
        (x + y)^5
    ]

    frs = modular_reduction(fs, FF)

    fr = frs[3]
    @test AA.base_ring(fr) == FF

    frs = rational_reconstruction(frs, BigInt(modulo))

    @test fs == frs
end

@testset "CRT tests" begin

    moduli = [11]
    xs = [3, 4, 5, 10]

    for x in xs
        rems = [mod(x, mi) for mi in moduli]
        @test CRT(moduli, rems) == x
    end

    moduli = Primes.primes(5, 20)
    m = prod(moduli)

    xs = [ rand(1:(m-1)) for _ in 1:100 ]

    for x in xs
        rems = [mod(x, mi) for mi in moduli]

        @test CRT(moduli, rems) == x
    end

    moduli = Primes.primes(2^5, 2^6)
    m = prod(moduli)

    xs = [ rand(1:(m-1)) for _ in 1:100 ]

    for x in xs
        rems = [mod(x, mi) for mi in moduli]

        @test CRT(moduli, rems) == x
    end

end

@testset "CRT basis reconstruction" begin
    moduli = [7, 5]

    FF1 = Sing.N_ZpField(moduli[1])
    FF2 = Sing.N_ZpField(moduli[2])

    R1, (x, y) = AA.PolynomialRing(FF1, ["x1", "y1"])
    R11, (a, b) = AA.PolynomialRing(AA.FractionField(R1), ["a1", "b1"])

    fs1 = [
        4*a^2*b^2 + (x + 6y)//(x + y)*a + (3x + 2x*y + 1*y)//(y)*b + 5y^2,
        1a + 3b + 2a*b + 2x//(1y)
    ]

    R2, (x, y) = AA.PolynomialRing(FF2, ["x2", "y2"])
    R22, (a, b) = AA.PolynomialRing(AA.FractionField(R2), ["a2", "b2"])

    fs2 = [
        2*a^2*b^2 + (1x + 3y)//(x + y)*a + (3x + 1x*y + 4*y)//(y)*b + 5y^2,
        1a + 3b + 1a*b + 4x//(y)
    ]

    basis = CRT_reconstruction([fs1, fs2], moduli)

    yoverx = AA.parent(basis[1])
    basepolyrings = AA.base_ring(yoverx)
    ground = AA.ZZ

    x, y = AA.gens(AA.base_ring(basepolyrings))
    a, b = AA.gens(yoverx)

    # to check when number is in the denominator
    ans = [
           32*a^2*b^2 + (x + 13*y)//(x + y)*a + (16*x*y + 3*x + 29*y)//(y)*b + 5*y^2,
           1a + 3b + 16a*b + 9x//(y)
    ]

    @test basis == ans

end


using DiffEqBayesStan, OrdinaryDiffEq, ParameterizedFunctions,
      RecursiveArrayTools, Distributions, MCMCChains, Test

println("\n\nFour parameter case\n")

f1 = @ode_def begin
    dx = a*x - b*x*y
    dy = -c*y + d*x*y
end a b c d

u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())

t = collect(range(1,stop=10,length=50))
randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(1.5,0.01),0.5,2),truncated(Normal(1.0,0.01),0.5,1.5),
          truncated(Normal(3.0,0.01),0.5,4),truncated(Normal(1.0,0.01),0.5,2)]

bayesian_result = stan_inference(prob1,t,data,priors;
    num_chains=1,num_samples=1000,num_warmups=1000,
    vars =(DiffEqBayesStan.StanODEData(),InverseGamma(4,1)))

@test mean(get(bayesian_result.chains,:theta_1)[1]) ≈ 1.5 atol=1e-1
@test mean(get(bayesian_result.chains,:theta_2)[1]) ≈ 1.0 atol=1e-1
@test mean(get(bayesian_result.chains,:theta_3)[1]) ≈ 3.0 atol=1e-1
@test mean(get(bayesian_result.chains,:theta_4)[1]) ≈ 1.0 atol=1e-1


#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2019
#
# This is a Julia package for discrete event simulation
#
using DiscreteEvents, Distributions, Random, .Threads, Test

num_customers = 10   # total number of customers generated
num_servers = 2      # number of servers
μ = 1.0 / 2          # service rate
λ = 0.9              # arrival rate
arrival_dist = Exponential(1/λ)  # interarrival time distriubtion
service_dist = Exponential(1/μ); # service time distribution

# describe the server process
function server(clk::Clock, id::Int, input::Channel, output::Channel, service_dist::Distribution)
    job = take!(input)
    # _bench[end] || now!(clk, ()->@printf("%5.3f: server %d serving customer %d\n", tau(clk), id, job))
    delay!(clk, rand(service_dist))
    # _bench[end] || now!(clk, ()->@printf("%5.3f: server %d finished serving %d\n", tau(clk), id, job))
    put!(output, job)
end

# model arrivals
function arrivals(clk::Clock, queue::Channel, num_customers::Int, arrival_dist::Distribution)
    for i = 1:num_customers # initialize customers
        delay!(clk, rand(arrival_dist))
        put!(queue, i)
        # _bench[end] || now!(clk, ()->@printf("%5.3f: customer %d arrived\n", tau(clk), i))
    end
end

function run_model(arrival_dist, service_dist, num_customers, num_servers, t)
    onthread(2) do
        Random.seed!(8710)   # set random number seed for reproducibility
        clock = Clock()
        for i in 1:num_servers
            process!(clock, Prc(i, server, i, input, output, service_dist))
        end
        process!(clock, Prc(0, arrivals, input, num_customers, arrival_dist), 1)
        run!(clock, t)
    end
end

# create a matrix of random numbers with n row and
# a column of random numbers generated on each thread
function prand(n::Int)
    rd = zeros(Int64, n, nthreads())
    @threads for i in 1:nthreads()
        rd[:,i] = rand(Int64, n)
    end
    return rd
end

input = Channel{Int}(Inf)
output = Channel{Int}(Inf)

if (VERSION ≥ v"1.3") && (nthreads() > 1)
    println("... testing utilities ...")
    run_model(arrival_dist, service_dist, 10, 2, 20)
    @test length(output) == 10

    pseed!(123)
    x = prand(1)
    @test x != prand(1)
    pseed!(123)
    @test x == prand(1)
end


using DelimitedFiles

const data = readdlm("day11.txt", String)

# floor will be 0, empty will be 1 and occupied will be 2

function create_matrix(data::Array{<:AbstractString,2})
    # We will pad the data with 0s around
    height = length(data) + 2
    width = length(data[1]) + 2
    mat = zeros(Int, (height, width))
    for (i, line) in enumerate(data)
        for (j, c) in enumerate(data[i])
            if c == 'L'
                # +1 because padding
                mat[i + 1 , j + 1] = 1
            end
        end
    end
    return mat
end

function process(mat::Array{Int,2}, i::Int, j::Int)
    state = mat[i, j]
    # empty
    if state == 1
        # if any is occupied, stay empty, else occupied
        new_state = any(mat[xi, xj] == 2 for xi ∈ i - 1:i + 1, xj ∈ j - 1:j + 1) ? 1 : 2
    # occupied
    elseif state == 2
        # If 4 adjacent (i.e. 5 including itself) are occupied, becomes empty
        new_state = sum(mat[xi, xj] == 2 for xi ∈ i - 1:i + 1, xj ∈ j - 1:j + 1) >= 5 ? 1 : 2
    else
        new_state = state
    end
    return new_state
end

function one_pass(mat::Array{Int,2})
    h, w = size(mat)
    new_mat = zeros(Int, (h, w))
    for i ∈ 2:h - 1, j ∈ 2:w - 1
        new_mat[i, j] = process(mat, i, j)
    end
    return new_mat
end

function wait_for_stable(mat::Array{Int,2})
    new_mat = one_pass(mat)
    while new_mat != mat
        mat = new_mat
        new_mat = one_pass(mat)
    end
    return count(==(2), mat)
end

const mat = create_matrix(data)

println("Part 1: $(wait_for_stable(mat))")

# Now we want the first non empty seat in each of the 8 directions

function look_in_direction(mat::Array{Int,2}, i::Int, j::Int, xi::Int, xj::Int)
    maxi, maxj = size(mat)
    while true
        new_i, new_j = i + xi, j + xj
        if new_i ∈ [1, maxi] || new_j ∈ [1, maxj]
            # we reached the end
            return 0
        end
        if mat[new_i, new_j] != 0
            return mat[new_i, new_j]
        end
        i, j = new_i, new_j
    end
end

function process_p2(mat::Array{Int,2}, i::Int, j::Int)
    state = mat[i, j]
    # empty
    if state == 1
        # if any is occupied, stay empty, else occupied
        new_state = any(look_in_direction(mat, i, j, xi, xj) == 2 for xi ∈ - 1:1, xj ∈ - 1:1) ? 1 : 2
    # occupied
    elseif state == 2
        # If 5 adjacent (i.e. 6 including itself) are occupied, becomes empty
        new_state = sum(look_in_direction(mat, i, j, xi, xj) == 2 for xi ∈ - 1:1, xj ∈ - 1:1) >= 6 ? 1 : 2
    else
        new_state = state
    end
    return new_state
end

function one_pass_p2(mat::Array{Int,2})
    h, w = size(mat)
    new_mat = zeros(Int, (h, w))
    for i ∈ 2:h - 1, j ∈ 2:w - 1
        new_mat[i, j] = process_p2(mat, i, j)
    end
    return new_mat
end

function wait_for_stable_p2(mat::Array{Int,2})
    new_mat = one_pass(mat)
    while new_mat != mat
        mat = new_mat
        new_mat = one_pass_p2(mat)
    end
    return count(==(2), mat)
end

println("Part 2: $(wait_for_stable_p2(mat))")


"""
    BayesianLinearRegressor{Tmw, TΛw}

A Bayesian Linear Regressor is a distribution over linear functions given by
```julia
w ~ Normal(mw, Λw)
f(x) = dot(x, w)
```
where `mw` and `Λw` are the mean and precision of `w`, respectively.
"""
struct BayesianLinearRegressor{Tmw<:AbstractVector,TΛw<:AbstractMatrix} <: AbstractGP
    mw::Tmw
    Λw::TΛw
end

const FiniteBLR = FiniteGP{<:BayesianLinearRegressor}

# All code below implements the primary + secondary AbstractGPs.jl APIs.

x_as_colvecs(fx::FiniteBLR) = x_as_colvecs(fx.x)

x_as_colvecs(x::ColVecs) = x

x_as_colvecs(x::RowVecs) = ColVecs(x.X')

function x_as_colvecs(x::T) where {T<:AbstractVector}
    return error(
        "$T is not a subtype of AbstractVector that is known. Please provide either a",
        "ColVecs or RowVecs.",
    )
end

AbstractGPs.mean(fx::FiniteBLR) = x_as_colvecs(fx).X' * fx.f.mw

function AbstractGPs.cov(fx::FiniteBLR)
    α = _cholesky(fx.f.Λw).U' \ x_as_colvecs(fx).X
    return Symmetric(α' * α + fx.Σy)
end

function AbstractGPs.var(fx::FiniteBLR)
    α = _cholesky(fx.f.Λw).U' \ x_as_colvecs(fx).X
    return vec(sum(abs2, α; dims=1)) .+ diag(fx.Σy)
end

AbstractGPs.mean_and_cov(fx::FiniteBLR) = (mean(fx), cov(fx))

AbstractGPs.mean_and_var(fx::FiniteBLR) = (mean(fx), var(fx))

function AbstractGPs.rand(rng::AbstractRNG, fx::FiniteBLR, samples::Int)
    X = x_as_colvecs(fx).X
    w = fx.f.mw .+ _cholesky(fx.f.Λw).U \ randn(rng, size(X, 1), samples)
    return X' * w .+ _cholesky(fx.Σy).U' * randn(rng, size(X, 2), samples)
end

function AbstractGPs.logpdf(fx::FiniteBLR, y::AbstractVector{<:Real})
    _, Bt, δy, logpdf_δy, Λεy = __compute_inference_quantities(fx, y)
    return -(logdet(Λεy) - sum(abs2, Λεy.U' \ (Bt'δy))) / 2 + logpdf_δy
end

function AbstractGPs.posterior(fx::FiniteBLR, y::AbstractVector{<:Real})
    Uw, Bt, δy, _, Λεy = __compute_inference_quantities(fx, y)

    # Compute posterior over decorrelated weights.
    mεy = Λεy \ (Bt'δy)

    # Compute posterior over weights.
    T = Λεy.U * Uw
    return BayesianLinearRegressor(fx.f.mw + Uw \ mεy, Symmetric(T'T))
end

# Computation utilised in both `logpdf` and `posterior`.
function __compute_inference_quantities(fx::FiniteBLR, y::AbstractVector{<:Real})
    X = x_as_colvecs(fx).X
    length(y) == size(X, 2) || throw(error("length(y) != size(fx.x.X, 2)"))
    blr = fx.f
    N = length(y)

    Uw = _cholesky(blr.Λw).U
    Σy = _cholesky(fx.Σy)

    Bt = Σy.U' \ (Uw' \ X)'
    δy = Σy.U' \ (y - mean(fx))

    logpdf_δy = -(N * log(2π) + logdet(Σy) + sum(abs2, δy)) / 2

    Λεy = _cholesky(Symmetric(Bt'Bt + I))

    return Uw, Bt, δy, logpdf_δy, Λεy
end

# Random function sample generation
# Following the Random API: https://docs.julialang.org/en/v1/stdlib/Random/#Hooking-into-the-Random-API
struct BLRFunctionSample{Tw<:AbstractVector}
    w::Tw
end

(s::BLRFunctionSample)(X::AbstractMatrix{<:Real}) = X's.w
(s::BLRFunctionSample)(X::ColVecs) = X.X's.w
(s::BLRFunctionSample)(X::RowVecs) = X.X * s.w

function Random.Sampler(
    ::Type{<:AbstractRNG}, blr::BayesianLinearRegressor, ::Random.Repetition
)
    return blr
end

function Random.rand(rng::AbstractRNG, blr::BayesianLinearRegressor)
    w = blr.mw .+ _cholesky(blr.Λw).U \ randn(rng, size(blr.mw))
    return BLRFunctionSample(w)
end

function Random.rand(rng::AbstractRNG, blr::BayesianLinearRegressor, dims::Dims)
    ws = blr.mw .+ _cholesky(blr.Λw).U \ randn(rng, (only(size(blr.mw)), prod(dims)))
    bs = [BLRFunctionSample(w) for w in eachcol(ws)]
    return reshape(bs, dims)
end

function Random.rand!(
    rng::AbstractRNG, A::AbstractArray{<:BLRFunctionSample}, blr::BayesianLinearRegressor
)
    ws = blr.mw .+ _cholesky(blr.Λw).U \ randn(rng, (only(size(blr.mw)), prod(size(A))))
    for i in LinearIndices(A)
        @inbounds A[i] = BLRFunctionSample(ws[:, i])
    end
    return A
end


module TestStat
  include("../src/stat.jl")
  using DataFrames
  using StatsBase
  using Base.Test

  nona = @data [1, 2, 3, 4, 5]
  onena = @data [1, NA, 3, 4, 5]
  twona = @data [1, NA, 3, NA, 5]

  # Using 0 NA
  # Track StatsBase zscore results. They use sample stddev as default, so we
  # should too for consistency. Throw an error if this changes.
  statsbase_zscore = @data zscore(nona)
  statsbase_zscore_pop = @data zscore(nona, 3., √2)
  @test isequal(nazscore(nona), statsbase_zscore)
  @test isequal(nazscore_pop(nona), statsbase_zscore_pop)
  # Give the objectively correct results as a sanity check in case StatsBase
  # gets whacky.
  expected_nona = @data [
    -1.2649110640673518,
    -0.6324555320336759,
    0.0,
    0.6324555320336759,
    1.2649110640673518
  ]
  expected_nona_pop = @data [
    -1.414213562373095,
    -0.7071067811865475,
    0.0,
    0.7071067811865475,
    1.414213562373095
  ]
  @test isequal(nazscore(nona), expected_nona)
  @test isequal(nazscore_pop(nona), expected_nona_pop)

  # Using 1 NA
  expected_onena = @data [
    -1.3174650984805198,
    NA,
    -0.14638501094227999,
    0.43915503282683993,
    1.02469507659596
  ]
  expected_onena_pop = @data [
    -1.52127765851133,
    NA,
    -0.1690308509457033,
    0.50709255283711,
    1.1832159566199232
  ]

  @test isequal(nazscore(onena), expected_onena)
  @test isequal(nazscore_pop(onena), expected_onena_pop)

  # Using 2 NA
  expected_twona = @data [-1., NA, 0., NA, 1.]
  expected_twona_pop = @data [
    -1.224744871391589,
    NA,
    0.0,
    NA,
    1.224744871391589
  ]

  @test isequal(nazscore(twona), expected_twona)
  @test isequal(nazscore_pop(twona), expected_twona_pop)
end


include("../src/SINDY.jl")
using LinearAlgebra
using DataDrivenDiffEq

# Sample and test(integration tests) some of the functionalities for Types
ds = LorenzSystem()
ds1 = ODE1()
ds2 = ODE3()
opt = STLSQ(0.01)
lib = TrigBasis()
lib1 = PolyTrigBasis()

# Integration tests - lorenz attractor
begin
        data = datagen(LorenzSystem())
        savefig("./figures/Traj_Lorenz.svg")
        v = differentiate(data,TotalVariationalDerivativative())
        savefig("./figures/Vel_Lorenz.svg")
        b_type = PolyTrigBasis()
        θ = basis(munge(data.u),b_type)
        LinearAlgebra.cond(θ)
        ξ = _optimize(θ,v,STLSQ(0.1))
        pprint(ξ,b_type)

        # This function needs to be defined manually from the output of pretty print.
        function lorenz_remake!(du,u,p,t)
                du[1] = -9.9756*u[1] + 9.97594*u[2]
                du[2] = 27.7373*u[1] - 0.9388*u[2] - 0.99*u[1]*u[3]
                du[3] = -2.65769*u[3] +  0.99615*u[1]*u[2]
        end

        u₀ = [1,0,0]
        p = [10,28,8/3]
        tspan = (0,100)
        t = 0.01:0.01:100

        prob = ODEProblem(lorenz_remake!,u₀,tspan,p)
        sol = solve(prob,Tsit5(),saveat=t)
        fig = plot(sol,vars=(1,2,3),title="Lorenz Attractor remade")
        display(fig)
        savefig("./figures/Remade_Lorenz.svg")
end

#TODO: The basis function is the only manual part of the implementation. This has to change.
#TODO: Include the forcing function expressions in the defn of the Dynamical System.

# Integration tests - Lotka Volterra

begin
        data = datagen(LotkaVolterra())
        savefig("./figures/Traj_Lotka.svg")
        v = differentiate(data,TotalVariationalDerivativative())
        savefig("./figures/Vel_Lotka.svg")
        b_type = PolynomialBasis()
        θ = basis(munge(data.u),b_type)
        LinearAlgebra.cond(θ)
        ξ = _optimize(θ,v,STLSQ(0.01))
        pprint(ξ,b_type)

        function lk!(du,u,p,t)
                du[1] = 0.7003*u[1] - 0.30014*u[1]*u[2]
                du[2] = -0.299953*u[2] + 0.399979*u[1]*u[2]
        end

        u₀ = [1,1]
        p = []
        tspan = (0,20)
        t = 0.01:0.01:20
        prob = ODEProblem(lk!,u₀,tspan,p)
        sol = solve(prob,Tsit5(),saveat=t)
        fig = plot(sol,vars=(1,2),title="Lotka Volterra remade")
        display(fig)
        savefig("./figures/Remade_Lotka.svg")
end

# Integration - Helicopter data

begin
        data = permutedims(datagen(HelicopterData("./src/data/heli.csv")))
        t = Vector(0.01:0.01:131)
        savefig("./figures/traj_heli.svg")
        plot(t,data[1,:],xlabel="t",ylabel="θ",title="θ")
        savefig("./figures/heli_θ.png")
        plot(t,data[2,:],xlabel="t",ylabel="ψ",title="ψ")
        savefig("./figures/heli_ψ.png")

        # v = differentiate(data,t,TotalVariationalDerivativative())
        # savefig("./figures/velo_heli.svg")
        # v,s = collocate_data(data,t)
        v = zeros(size(data))
        for i = 1:size(data,1)
                int = CubicSpline(data[i,:],t)
                v[i,:] = DataInterpolations.derivative.((int,),t)
        end

        b_type = TrigBasis()
        θ = basis(data,b_type)
        LinearAlgebra.cond(θ)
        ξ = _optimize(θ,v,STLSQ(0.001))
        pprint(ξ,b_type)

        # function helicopter!(du,u,p,t)
        #         du[1] = 0.0546156*sin(u[2]) -0.0431238*cos(u[1]) + 0.022303*cos(u[2]) + 0.0176159*tan(u[1])
        #         du[2] = 0.819696*sin(u[1]) -0.041357*sin(u[2]) + 0.17279*cos(u[1]) +  0.0282611*cos(u[2]) -0.469918*tan(u[1])
        # end

        # function helicopter!(du,u,p,t)
        #         du[1] = 0.0544514*sin(u[2]) -0.0432586*cos(u[1]) + 0.022315*cos(u[2]) + 0.0170897*tan(u[1])
        #         du[2] = 0.78517*sin(u[1]) -0.0393751*sin(u[2]) + 0.170829*cos(u[1]) +  0.0281283*cos(u[2]) -0.439536*tan(u[1])
        # end

        function helicopter!(du,u,p,t)
                du[1] = 0.00168582*sin(u[1]) + 0.0546215*sin(u[2]) -0.0431765 *cos(u[1]) + 0.022313*cos(u[2]) + 0.0160404*tan(u[1])
                du[2] = 0.797168*sin(u[1]) -0.0402192*sin(u[2]) + 0.171233*cos(u[1]) +  0.0281401*cos(u[2]) -0.450491*tan(u[1])
        end

        u₀ = [-0.67,0.67]
        p = []
        tspan = (0,131)
        t = 0.01:0.01:131
        prob = ODEProblem(helicopter!,u₀,tspan,p)
        sol = solve(prob,Tsit5(),saveat=t)

        # Validate the measurement

        temp = munge(sol.u)
        plot(sol.t,temp[1,:],xlabel="t",ylabel="θ",title="θ")
        savefig("./figures/cθ_learned.png")
        plot(sol.t,temp[2,:],xlabel="t",ylabel="ψ",title="ψ")
        savefig("./figures/cψ_learned.png")

        fig = plot(sol,vars=(1,2),title="Helicopter model remade",xlabel="ψ",ylabel="θ")
        display(fig)
        savefig("./figures/cRemade_heli.svg")
end


using Base.Cartesian

##############################################
# Accelerated Convolutions for Efficient Multi-Scale Time to Contact Computation in Julia
# Alexander Amini, Alan Edelman, Berthold Horn
##############################################
@generated function fastconv(E::Array{T,N}, k::Array{T,N}) where {T,N}
    quote
        retsize = [size(E)...] + [size(k)...] .- 1
        retsize = tuple(retsize...)
        ret = zeros(T, retsize)

        convn!(ret,E,k)
        return ret

    end
end

# in place helper operation to speedup memory allocations
@generated function convn!(out::Array{T}, E::Array{T,N}, k::Array{T,N}) where {T,N}
    quote
        @inbounds begin
            @nloops $N x E begin
                @nloops $N i k begin
                    (@nref $N out d->(x_d + i_d - 1)) += (@nref $N E x) * (@nref $N k i)
                end
            end
        end
        return out
    end
end


using GeometricIntegrators.BasisFunctions
using GeometricIntegrators.CommonFunctions
using GeometricIntegrators.Config
using GeometricIntegrators.Integrators
using GeometricIntegrators.Quadratures
using GeometricIntegrators.Solvers
using GeometricIntegrators.TestProblems.Oscillator
using GeometricIntegrators.Utils
using Test

using GeometricIntegrators.TestProblems.Oscillator: Δt, nt, refx, refq, refp

iode = oscillator_iode()

QGau4 = GaussLegendreQuadrature(4)
BGau4 = LagrangeBasis(nodes(QGau4))

### CGVI Integrators ###

cgint = IntegratorCGVI(iode, BGau4, QGau4, Δt)
cgsol = integrate(cgint, nt)

@test rel_err(cgsol.q, refx) < 1E-7

### DGVI Integrators ###

dgint = IntegratorDGVI(iode, BGau4, QGau4, Δt)
dgsol = integrate(dgint, nt)

@test rel_err(dgsol.q, refx) < 1E-7


using Distributions

interaction_dict = Dict()

interaction_dict["binary"] = [-1.0,1.0]
interaction_dict["normal"] = Normal()
interaction_dict["uniform"] = Uniform(-1.0,1.0)


using CompScienceMeshes, BEAST
using LinearAlgebra

Γ = meshcuboid(1.0, 1.0, 1.0, 0.15)
X = raviartthomas(Γ)

κ,  η  = 1.0, 1.0
κ′, η′ = 1.4κ, η/1.4

T  = Maxwell3D.singlelayer(wavenumber=κ)
T′ = Maxwell3D.singlelayer(wavenumber=κ′)
K  = Maxwell3D.doublelayer(wavenumber=κ)
K′ = Maxwell3D.doublelayer(wavenumber=κ′)

# d = normalize(x̂+ŷ+ẑ)
d = ẑ
# p = normalize(d × ẑ)
p = x̂
E = Maxwell3D.planewave(direction=d, polarization=p, wavenumber=κ)
H = -1/(im*κ*η)*curl(E)

e, h = (n × E) × n, (n × H) × n

@hilbertspace j m
@hilbertspace k l

α, α′ = 1/η, 1/η′
pmchwt = @discretise(
    (η*T+η′*T′)[k,j] +      (K+K′)[k,m] -
         (K+K′)[l,j] + (α*T+α′*T′)[l,m] == h[k] + e[l],
    j∈X, m∈X, k∈X, l∈X)

u = solve(pmchwt)

nX = numfunctions(X)
uj = u[1:nX]
um = u[nX+1:end]

Θ, Φ = range(0.0,stop=2π,length=100), 0.0
ffpoints = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for θ in Θ for ϕ in Φ]

# Don't forgt the far field comprises two contributions
ffm = potential(MWFarField3D(κ*im), ffpoints, um, X)
ffj = potential(MWFarField3D(κ*im), ffpoints, uj, X)
ff = η*im*κ*ffj + im*κ*cross.(ffpoints, ffm)

using Plots
plot(xlabel="theta")
plot!(Θ,norm.(ff),label="far field")

import PlotlyJS
using LinearAlgebra
fcrj, _ = facecurrents(uj,X)
PlotlyJS.plot(patch(Γ, norm.(fcrj)))


import RDatasets
import DataFrames: DataFrame
using LowRankModels

# pick a data set
df = RDatasets.dataset("psych", "msq")

# make a GLRM on the whole dataframe using type imputation
auto_glrm, labels = GLRM(df, 3)
# fit!(auto_glrm)  - this doesn't work yet, because the data contains some trivial columns (ordinals with <2 levels)

# now we'll try it without type imputation
# we'll just fit four of the columns, to try out all four data types
dd = DataFrame([df[s] for s in [:TOD, :Scale, :Vigorous, :Wakeful]])
dd[!,end] = (dd[:,end].==1)
datatypes = [:real, :cat, :ord, :bool]

# fit it!
glrm = GLRM(dd, 2, datatypes)

println("initializing")
init_svd!(glrm)

println("fitting")
X, Y, ch = fit!(glrm)

# print results
println(ch.objective)

println("imputing")
impute(glrm)

println("crossvalidating")
cross_validate(glrm, do_obs_check=false, init=init_svd!)


using StructuralEquationModels, Distributions, Random, Optim, LineSearches
import StructuralEquationModels as SEM
include("helper.jl")

x = Symbol.("x".*string.(1:13))

S = [:x1 0 0 0 0 0 0 0
     0 :x2 0 0 0 0 0 0
     0 0 :x3 0 0 0 0 0
     0 0 0 :x4 0 0 0 0
     0 0 0 0 :x5 0 0 0
     0 0 0 0 0 :x6 0 0
     0 0 0 0 0 0 :x7 :x9
     0 0 0 0 0 0 :x9 :x8]

F = [1.0 0 0 0 0 0 0 0
     0 1.0 0 0 0 0 0 0
     0 0 1.0 0 0 0 0 0
     0 0 0 1.0 0 0 0 0
     0 0 0 0 1.0 0 0 0
     0 0 0 0 0 1.0 0 0]

A = [0 0 0 0 0 0 1.0   0
     0 0 0 0 0 0 :x10 0
     0 0 0 0 0 0 :x11 0
     0 0 0 0 0 0 0    1.0
     0 0 0 0 0 0 0    :x12
     0 0 0 0 0 0 0    :x13
     0 0 0 0 0 0 0    0
     0 0 0 0 0 0 0    0]

ram_matrices = RAMMatrices(;A = A, S = S, F = F, parameters = x, colnames = nothing)

true_val = [repeat([1], 8)
            0.4
            repeat([0.8], 4)]

start_val = [repeat([1], 9)
             repeat([0.5], 4)]

imply_ml = RAMSymbolic(;specification = ram_matrices, start_val = start_val)

imply_ml.Σ_function(imply_ml.Σ, true_val)

true_dist = MultivariateNormal(imply_ml.Σ)

Random.seed!(1234)
x = transpose(rand(true_dist, 100000))
semobserved = SemObsCommon(data = x)

loss_ml = SemLoss(SEM.SemML(;observed = semobserved, n_par = length(start_val)))

diff =
    SemDiffOptim(
        BFGS(;linesearch = BackTracking(order=3), alphaguess = InitialHagerZhang()),# m = 100),
        Optim.Options(
            ;f_tol = 1e-10,
            x_tol = 1.5e-8))

model_ml = Sem(semobserved, imply_ml, loss_ml, diff)
model_ml(true_val, true, false, false)
solution_ml = sem_fit(model_ml)

@test isapprox(true_val, solution(solution_ml); atol = .05)


using JuMP
using SCS
using LinearAlgebra

q = 4 # dimension of estimate space
p = 8 # number of experimental vectors
nmax = 3 # upper bound on lambda
n = 12

V = randn(q, p)

eye = Matrix{Float64}(I, q, q);

aOpt = Model(with_optimizer(SCS.Optimizer, verbose = 0))
@variable(aOpt, np[1:p], lower_bound = 0, upper_bound = nmax)
@variable(aOpt, u[1:q], lower_bound = 0)

@constraint(aOpt, sum(np) <= n)
for i = 1:q
    @SDconstraint(aOpt, [V * diagm(0 => np ./ n) * V' eye[:, i]; eye[i, :]' u[i]] >= 0)
end

@objective(aOpt, Min, sum(u))

optimize!(aOpt)

@show objective_value(aOpt);
@show value.(np);

eOpt = Model(with_optimizer(SCS.Optimizer, verbose = 0))
@variable(eOpt, np[1:p], lower_bound = 0, upper_bound = nmax)
@variable(eOpt, t)

@SDconstraint(eOpt, V * diagm(0 => np ./ n) * V' - (t .* eye) >= 0)
@constraint(eOpt, sum(np) <= n)

@objective(eOpt, Max, t)

optimize!(eOpt)

@show objective_value(eOpt);
@show value.(np);

dOpt = Model(with_optimizer(SCS.Optimizer, verbose = 0))
@variable(dOpt, np[1:p], lower_bound = 0, upper_bound = nmax)
@variable(dOpt, t)
@objective(dOpt, Max, t)
@constraint(dOpt, sum(np) <= n)
E = V * diagm(0 => np ./ n) * V'
@constraint(dOpt, [t, 1, (E[i, j] for i in 1:q for j in 1:i)...] in MOI.LogDetConeTriangle(q))

optimize!(dOpt)

@show objective_value(dOpt);
@show value.(np);


using ContactStructures

plotlyjs()
plot_along_param_curve(t -> (2*cos(2π*t),2*sin(2π*t),0), standard_structure, 0.0625;  midpoint_marker_color=:red, ε=0.35)
plot!(xlims=(-3, 3), ylims=(-3,3), zlims=(-3,3),legend=false)


# asymmetric test kasa method
for _ in 1:10
    local x0 = 10*rand()
    local y0 = 10*rand()
    local r = rand()+1
    local x = r.*[-1.0,0,0,1,sqrt(2)/2,sqrt(2)/2] .+ x0
    local y = r.*[0.0,1,-1,0,sqrt(2)/2,-sqrt(2)/2] .+ y0
    coefs = CircleFit.kasa(x,y)
    @test x0 ≈ coefs[1]
    @test y0 ≈ coefs[2]
    @test r ≈ coefs[3]
end

# symmetric test kasa method
for _ in 1:10
    local x0 = 10*rand()
    local y0 = 10*rand()
    local r = rand()+1
    local x = r.*[-1.0,0,0,1] .+ x0
    local y = r.*[0.0,1,-1,0] .+ y0
    coefs = CircleFit.kasa(x,y)
    @test x0 ≈ coefs[1]
    @test y0 ≈ coefs[2]
    @test r ≈ coefs[3]
end


"""
Calculates the optimal feedback gains K,d as well as the 2nd Order approximation of the
Cost-to-Go, using a backward Riccati-style recursion. (non-allocating)
"""
function backwardpass!(solver::iLQRSolver2{T,QUAD,L,O,n,n̄,m}) where {T,QUAD<:QuadratureRule,L,O,n,n̄,m}
	N = solver.N

    # Objective
    obj = solver.obj
    model = solver.model

    # Extract variables
    Z = solver.Z; K = solver.K; d = solver.d;
    G = solver.G
    S = solver.S
	Quu_reg = solver.Quu_reg
	Qux_reg = solver.Qux_reg

    # Terminal cost-to-go
	Q = solver.Q[N]
    S[N].xx .= Q.xx
    S[N].x .= Q.x

    # Initialize expecte change in cost-to-go
    ΔV = @SVector zeros(2)

    k = N-1
    while k > 0
        ix = Z[k]._x
        iu = Z[k]._u

		# Get error state expanions
		fdx,fdu = error_expansion(solver.D[k], model)
		Q = solver.Q[k]

		# Calculate action-value expansion
		_calc_Q!(Q, S[k+1], S[k], fdx, fdu)

		# Regularization
		Quu_reg .= Q.uu #+ solver.ρ[1]*I
		Quu_reg .+= solver.ρ[1]*Diagonal(@SVector ones(m))
		Qux_reg .= Q.ux

	    if solver.opts.bp_reg
	        vals = eigvals(Hermitian(Quu_reg))
	        if minimum(vals) <= 0
	            @warn "Backward pass regularized"
	            regularization_update!(solver, :increase)
	            k = N-1
	            ΔV = @SVector zeros(2)
	            continue
	        end
	    end

        # Compute gains
		_calc_gains!(K[k], d[k], Quu_reg, Qux_reg, Q.u)

        # Calculate cost-to-go (using unregularized Quu and Qux)
		ΔV += _calc_ctg!(S[k], Q, K[k], d[k])

        k -= 1
    end

    regularization_update!(solver, :decrease)

    return ΔV

end

function static_backwardpass!(solver::iLQRSolver2{T,QUAD,L,O,n,n̄,m}) where {T,QUAD<:QuadratureRule,L,O,n,n̄,m}
	N = solver.N

    # Objective
    obj = solver.obj
    model = solver.model

    # Extract variables
    Z = solver.Z; K = solver.K; d = solver.d;
    G = solver.G
    S = solver.S
	Quu_reg = SMatrix(solver.Quu_reg)
	Qux_reg = SMatrix(solver.Qux_reg)

    # Terminal cost-to-go
	Q = error_expansion(solver.Q[N], model)
	Sxx = SMatrix(Q.xx)
	Sx = SVector(Q.x)

    # Initialize expected change in cost-to-go
    ΔV = @SVector zeros(2)

    k = N-1
    while k > 0
        ix = Z[k]._x
        iu = Z[k]._u

		# Get error state expanions
		fdx,fdu = error_expansion(solver.D[k], model)
		fdx,fdu = SMatrix(fdx), SMatrix(fdu)
		Q = error_expansion(solver.Q[k], model)

		# Calculate action-value expansion
		Q = _calc_Q!(Q, Sxx, Sx, fdx, fdu)

		# Regularization
		Quu_reg, Qux_reg = _bp_reg!(Q, fdx, fdu, solver.ρ[1], solver.opts.bp_reg_type)

	    if solver.opts.bp_reg
	        vals = eigvals(Hermitian(Quu_reg))
	        if minimum(vals) <= 0
	            @warn "Backward pass regularized"
	            regularization_update!(solver, :increase)
	            k = N-1
	            ΔV = @SVector zeros(2)
	            continue
	        end
	    end

        # Compute gains
		K_, d_ = _calc_gains!(K[k], d[k], Quu_reg, Qux_reg, Q.u)

        # Calculate cost-to-go (using unregularized Quu and Qux)
		Sxx, Sx, ΔV_ = _calc_ctg!(Q, K_, d_)
		ΔV += ΔV_
        k -= 1
    end

    regularization_update!(solver, :decrease)

    return ΔV

end

function _bp_reg!(Quu_reg::SizedMatrix{m,m}, Qux_reg, Q, fdx, fdu, ρ, ver=:control) where {m}
    if ver == :state
        Quu_reg .= Q.uu #+ solver.ρ[1]*fdu'fdu
		mul!(Quu_reg, Transpose(fdu), fdu, ρ, 1.0)
        Qux_reg .= Q.ux #+ solver.ρ[1]*fdu'fdx
		mul!(Qux_reg, fdu', fdx, ρ, 1.0)
    elseif ver == :control
        Quu_reg .= Q.uu #+ solver.ρ[1]*I
		Quu_reg .+= ρ*Diagonal(@SVector ones(m))
        Qux_reg .= Q.ux
    end
end

function _bp_reg!(Q, fdx, fdu, ρ, ver=:control)
    if ver == :state
		Quu_reg = Q.uu + ρ * fdu'fdu
		Qux_reg = Q.ux + ρ * fdu'fdx
    elseif ver == :control
		Quu_reg = Q.uu + ρ * I
        Qux_reg = Q.ux
    end

	Quu_reg, Qux_reg
end

function _calc_Q!(Q, S1, S, fdx, fdu)
	# Compute the cost-to-go, stashing temporary variables in S[k]
    # Qx =  Q.x[k] + fdx'S.x[k+1]
	mul!(Q.x, Transpose(fdx), S1.x, 1.0, 1.0)

    # Qu =  Q.u[k] + fdu'S.x[k+1]
	mul!(Q.u, Transpose(fdu), S1.x, 1.0, 1.0)

    # Qxx = Q.xx[k] + fdx'S.xx[k+1]*fdx
	mul!(S.xx, Transpose(fdx), S1.xx)
	mul!(Q.xx, S.xx, fdx, 1.0, 1.0)

    # Quu = Q.uu[k] + fdu'S.xx[k+1]*fdu
	mul!(S.ux, Transpose(fdu), S1.xx)
	mul!(Q.uu, S.ux, fdu, 1.0, 1.0)

    # Qux = Q.ux[k] + fdu'S.xx[k+1]*fdx
	mul!(S.ux, Transpose(fdu), S1.xx)
	mul!(Q.ux, S.ux, fdx, 1.0, 1.0)

	return nothing
end

function _calc_Q!(Q::StaticExpansion, Sxx, Sx, fdx::SMatrix, fdu::SMatrix)
	Qx = Q.x + fdx'Sx
	Qu = Q.u + fdu'Sx
	Qxx = Q.xx + fdx'Sxx*fdx
	Quu = Q.uu + fdu'Sxx*fdu
	Qux = Q.ux + fdu'Sxx*fdx
	StaticExpansion(Qx,Qxx,Qu,Quu,Qux)
end

function _calc_gains!(K::SizedArray, d::SizedArray, Quu::SizedArray, Qux::SizedArray, Qu)
	LAPACK.potrf!('U',Quu.data)
	K .= Qux
	d .= Qu
	LAPACK.potrs!('U', Quu.data, K.data)
	LAPACK.potrs!('U', Quu.data, d.data)
	K .*= -1
	d .*= -1
	# return K,d
end

function _calc_gains!(K, d, Quu::SMatrix, Qux::SMatrix, Qu::SVector)
	K_ = -Quu\Qux
	d_ = -Quu\Qu
	K .= K_
	d .= d_
	return K_,d_
end

function _calc_ctg!(S, Q, K, d)
	# S.x[k]  =  Qx + K[k]'*Quu*d[k] + K[k]'* Qu + Qux'd[k]
	tmp1 = S.u
	S.x .= Q.x
	mul!(tmp1, Q.uu, d)
	mul!(S.x, Transpose(K), tmp1, 1.0, 1.0)
	mul!(S.x, Transpose(K), Q.u, 1.0, 1.0)
	mul!(S.x, Transpose(Q.ux), d, 1.0, 1.0)

	# S.xx[k] = Qxx + K[k]'*Quu*K[k] + K[k]'*Qux + Qux'K[k]
	tmp2 = S.ux
	S.xx .= Q.xx
	mul!(tmp2, Q.uu, K)
	mul!(S.xx, Transpose(K), tmp2, 1.0, 1.0)
	mul!(S.xx, Transpose(K), Q.ux, 1.0, 1.0)
	mul!(S.xx, Transpose(Q.ux), K, 1.0, 1.0)
	transpose!(Q.xx, S.xx)
	S.xx .+= Q.xx
	S.xx .*= 0.5

    # calculated change is cost-to-go over entire trajectory
	t1 = d'Q.u
	mul!(Q.u, Q.uu, d)
	t2 = 0.5*d'Q.u
    return @SVector [t1, t2]
end

function _calc_ctg!(Q::StaticExpansion, K::SMatrix, d::SVector)
	Sx = Q.x + K'Q.uu*d + K'Q.u + Q.ux'd
	Sxx = Q.xx + K'Q.uu*K + K'Q.ux + Q.ux'K
	Sxx = 0.5*(Sxx + Sxx')
	# S.x .= Sx
	# S.xx .= Sxx
	t1 = d'Q.u
	t2 = 0.5*d'Q.uu*d
	return Sxx, Sx, @SVector [t1, t2]
end


export istotally_real, istotally_complex, conjugates, conjugates_real,
       conjugates_complex, conjugates_log, complex_conjugation

################################################################################
#
#  Totally real/complex predicates
#
################################################################################

@doc Markdown.doc"""
    istotally_real(K::NumberField) -> Bool

Returns true if and only if $K$ is totally real, that is, if all roots of the
defining polynomial are real.
"""
function istotally_real(K::NumField)
  return iszero(signature(K)[2])
end

@doc Markdown.doc"""
    istotally_complex(K::AnticNumberField) -> Bool

Returns true if and only if $K$ is totally complex, that is, if all roots of the
defining polynomial are not real.
"""
function istotally_complex(K::NumField)
  return iszero(signature(K)[1])
end

################################################################################
#
#  Conjugates and real embeddings
#
################################################################################

@doc Markdown.doc"""
    conjugates(x::nf_elem, abs_tol::Int) -> Vector{acb}

Compute the conjugates of `x` as elements of type `acb`.
Recall that we order the complex conjugates
$\sigma_{r+1}(x),...,\sigma_{r+2s}(x)$ such that
$\sigma_{i}(x) = \overline{sigma_{i + s}(x)}$ for $r + 1 \leq i \leq r + s$.

Every entry `y` of the vector returned satisfies
`radius(real(y)) < 2^-abs_tol` and `radius(imag(y)) < 2^-abs_tol` respectively.
"""
function conjugates(x::NumFieldElem, abs_tol::Int = 32, T = arb)
  if T === arb
    return conjugates_arb(x, abs_tol)
  else
    error("Cannot return conjugates as type $T")
  end
end

# This is for quick and dirty computations
function __conjugates_arb(x::nf_elem, prec::Int = 32)
  K = parent(x)
  d = degree(K)
  r1, r2 = signature(K)
  conjugates = Array{acb}(undef, r1 + 2*r2)

  c = conjugate_data_arb_roots(K, -1)

  CC = AcbField(prec, cached = false)
  RR = ArbField(prec, cached = false)

  xpoly = arb_poly(parent(K.pol)(x), prec)

  for i in 1:r1
    o = RR()
    ccall((:arb_poly_evaluate, libarb), Nothing,
          (Ref{arb}, Ref{arb_poly}, Ref{arb}, Int),
           o, xpoly, c.real_roots[i], prec)

    if !isfinite(o)
      error("oops")
    end
    conjugates[i] = CC(o)
  end

  for i in 1:r2
    tacb = CC()
    ccall((:arb_poly_evaluate_acb, libarb), Nothing,
          (Ref{acb}, Ref{arb_poly}, Ref{acb}, Int),
           tacb, xpoly, c.complex_roots[i], prec)

    if !isfinite(tacb)
      error("oops")
    end

    conjugates[r1 + i] = tacb
    conjugates[r1 + i + r2] = conj(conjugates[r1 + i])
  end

  return conjugates
end

function conjugates_arb(x::nf_elem, abs_tol::Int = 32)
  K = parent(x)
  d = degree(K)
  r1, r2 = signature(K)
  conjugates = Array{acb}(undef, r1 + 2*r2)
  target_tol = abs_tol
  abs_tol = Int(floor(abs_tol * 1.1))

  while true
    prec_too_low = false
    c = conjugate_data_arb_roots(K, abs_tol)

    if abs_tol < 0
      error("Precision overflow in conjugates_arb_log")
    end

    CC = AcbField(abs_tol, cached = false)
    RR = ArbField(abs_tol, cached = false)

    xpoly = arb_poly(parent(K.pol)(x), abs_tol)

    for i in 1:r1
      o = RR()
      ccall((:arb_poly_evaluate, libarb), Nothing,
            (Ref{arb}, Ref{arb_poly}, Ref{arb}, Int),
             o, xpoly, c.real_roots[i], abs_tol)

      if !isfinite(o) || !radiuslttwopower(o, -target_tol)
        abs_tol = 2*abs_tol
        prec_too_low = true
        break
      end
      conjugates[i] = CC(o)
    end

    if prec_too_low
      continue
    end

    for i in 1:r2
      tacb = CC()
      ccall((:arb_poly_evaluate_acb, libarb), Nothing,
            (Ref{acb}, Ref{arb_poly}, Ref{acb}, Int),
             tacb, xpoly, c.complex_roots[i], abs_tol)

      if !isfinite(tacb) || !radiuslttwopower(tacb, -target_tol)
        abs_tol = 2*abs_tol
        prec_too_low = true
        break
      end

      conjugates[r1 + i] = tacb
      conjugates[r1 + i + r2] = conj(conjugates[r1 + i])
    end

    if prec_too_low
      continue
    end

    for i in 1:d
      expand!(conjugates[i], -target_tol)
    end
    return conjugates
  end
end

@doc Markdown.doc"""
    conjugates_arb_real(x::nf_elem, abs_tol::Int) -> Vector{arb}

Compute the real conjugates of `x` as elements of type `arb`.

Every entry `y` of the array returned satisfies
`radius(y) < 2^-abs_tol`.
"""
function conjugates_real(x::nf_elem, abs_tol::Int = 32, T = arb)
  if T === arb
    return conjugates_arb_real(x, abs_tol)
  else
    error("Cannot return real conjugates as type $T")
  end
end

function conjugates_arb_real(x::nf_elem, abs_tol::Int = 32)
  r1, r2 = signature(parent(x))
  c = conjugates_arb(x, abs_tol)
  z = Array{arb}(undef, r1)

  for i in 1:r1
    z[i] = real(c[i])
  end

  return z
end

@doc Markdown.doc"""
    conjugates_complex(x::nf_elem, abs_tol::Int) -> Vector{acb}

Compute the complex conjugates of `x` as elements of type `acb`.
Recall that we order the complex conjugates
$\sigma_{r+1}(x),...,\sigma_{r+2s}(x)$ such that
$\sigma_{i}(x) = \overline{sigma_{i + s}(x)}$ for $r + 1 \leq i \leq r + s$.

Every entry `y` of the array returned satisfies
`radius(real(y)) < 2^-abs_tol` and `radius(imag(y)) < 2^-abs_tol`.
"""
function conjugates_complex(x::nf_elem, abs_tol::Int = 32, T = arb)
  if T === arb
    return conjugates_arb_complex(x, abs_tol)
  else
    error("Cannot return real conjugates as type $T")
  end
end

function conjugates_arb_complex(x::nf_elem, abs_tol::Int)
  r1, r2 = signature(parent(x))
  c = conjugates_arb(x, abs_tol)
  z = Vector{acb}(undef, r2)

  for i in (r1 + 1):(r1 + r2)
    z[i - r1] = c[i]
  end

  return z
end

################################################################################
#
#  Logarithms of conjugates
#
################################################################################

@doc Markdown.doc"""
    conjugates_arb_log(x::nf_elem, abs_tol::Int) -> Array{arb, 1}

Returns the elements
$(\log(\lvert \sigma_1(x) \rvert),\dotsc,\log(\lvert\sigma_r(x) \rvert),
\dotsc,2\log(\lvert \sigma_{r+1}(x) \rvert),\dotsc,
2\log(\lvert \sigma_{r+s}(x)\rvert))$ as elements of type `arb` with radius
less then `2^-abs_tol`.
"""
function conjugates_log(x::nf_elem, abs_tol::Int = 32, T = arb)
  if T === arb
    return conjugates_arb_log(x, abs_tol)
  else
    throw(error("Cannot return real conjugates as type ", T))
  end
end

function conjugates_arb_log(x::nf_elem, abs_tol::Int)
  K = parent(x)
  c = conjugate_data_arb_roots(K, abs_tol)
  r1 = length(c.real_roots)
  r2 = length(c.complex_roots)
  d = degree(K)
  target_tol = abs_tol

  # TODO: Replace this using multipoint evaluation of libarb
  z = Array{arb}(undef, r1 + r2)
  while true
    prec_too_low = false
    c = conjugate_data_arb_roots(K, abs_tol)
    if abs_tol > 2^20
      error("Something wrong in conjugates_arb_log")
    end
    xpoly = arb_poly(parent(K.pol)(x), abs_tol)
    RR = ArbField(abs_tol, cached = false)
    for i in 1:r1
      o = RR()
      ccall((:arb_poly_evaluate, libarb), Nothing,
            (Ref{arb}, Ref{arb_poly}, Ref{arb}, Int),
            o, xpoly, c.real_roots[i], abs_tol)
      abs!(o, o)
      log!(o, o)
      z[i] = o

      if !isfinite(z[i]) || !radiuslttwopower(z[i], -target_tol)
        abs_tol = 2*abs_tol
        prec_too_low = true
        break
      end
    end

    if prec_too_low
      continue
    end

    CC = AcbField(abs_tol, cached = false)

    tacb = CC()
    for i in 1:r2
      oo = RR()
      ccall((:arb_poly_evaluate_acb, libarb), Nothing,
            (Ref{acb}, Ref{arb_poly}, Ref{acb}, Int),
            tacb, xpoly, c.complex_roots[i], abs_tol)
      abs!(oo, tacb)
      log!(oo, oo)
      mul2exp!(oo, oo, 1)
      z[r1 + i] = oo

      if !isfinite(z[r1 + i]) || !radiuslttwopower(z[r1 + i], -target_tol)
        abs_tol = 2*abs_tol
        prec_too_low = true
        break
      end
    end

    if prec_too_low
      continue
    end

    for i in 1:r1 + r2
      zz = deepcopy(z[i])
      expand!(z[i], -target_tol)
    end
    return z
  end
end

function conjugates_arb_log(x::nf_elem, R::ArbField)
  z = conjugates_arb_log(x, R.prec)
  return map(R, z)
end

################################################################################
#
#  Minkowski map
#
################################################################################

@doc Markdown.doc"""
    minkowski_map(a::nf_elem, abs_tol::Int) -> Array{arb, 1}

Returns the image of $a$ under the Minkowski embedding.
Every entry of the array returned is of type `arb` with radius less then
`2^(-abs_tol)`.
"""
function minkowski_map(a::T, abs_tol::Int = 32) where T <: NumFieldElem
  z = _minkowski_map_and_apply(a, abs_tol, identity)
  return z
end

# The following function computes the minkowski_map, applies G to the output.
# G mus be a function (::Vector{arb}, abs_tol::Int) -> Bool, *
# where the first return value indicates if the result is good enough
function _minkowski_map_and_apply(a, abs_tol, G, work_tol::Int = abs_tol)
  K = parent(a)
  A = Array{arb}(undef, absolute_degree(K))
  c = conjugates_arb(a, work_tol)::Vector{acb}
  r, s = signature(K)

  for i = 1:r
    @assert isreal(c[i])
    A[i] = real(c[i])
    if !radiuslttwopower(A[i], -abs_tol)
      error("this should not happen")
    end
  end

  if work_tol > 2^18 || abs_tol > 2^18
    throw(error("asdsd"))
  end

  #R = ArbField(precision(parent(c[1])), false)
  R = ArbField(2 * work_tol, cached = false)
  sqrt2 = sqrt(R(2))

  for i in 1:s
    t = c[r + i]
    A[r + 2*i - 1] = sqrt2 * real(t)
    A[r + 2*i] = sqrt2 * imag(t)
    if !radiuslttwopower(A[r + 2*i], -abs_tol)
      return _minkowski_map_and_apply(a, abs_tol, G, Int(floor(work_tol * 2)))
    end
  end

  if typeof(G) === typeof(identity)
    return A
  else
    b, B = G(A, abs_tol)
    if b
      return B
    else
      return _minkowski_map_and_apply(a, abs_tol, G, 2*work_tol)
    end
  end
end

################################################################################
#
#  T_2
#
################################################################################

function t2(x::S, abs_tol::Int = 32, T = arb) where S <: NumFieldElem
  if T === arb
    g = function(w, abs_tol)
      z = mapreduce(y -> y^2, +, w)
      return radiuslttwopower(z, -abs_tol), z
    end
    c = _minkowski_map_and_apply(x, abs_tol, g)
    return c
  else
    throw(NotImplemented())
  end
end

############################################################################
#
#  Signs of real embeddings
#
############################################################################

#@doc Markdown.doc"""
##    _signs(a::nf_elem) -> Array{Int, 1}
#> For a non-zero element $a$ return the signs of all real embeddings.
#"""
function _signs(a::nf_elem)
  if iszero(a)
    error("element must not be zero")
  end

  p = 16
  r1, r2 = signature(parent(a))

  if r1 == 0
    return Int[]
  end

  s = Array{Int}(undef, r1)
  while true
    c = conjugates(a, p)
    done = true
    for i=1:r1
      if contains(reim(c[i])[1], 0)
        p *= 2
        done = false
        break
      end
      s[i] = ispositive(real(c[i])) ? 1 : -1
    end
    if done
      return s
    end
  end
end

#@doc Markdown.doc"""
##    signs(a::FacElem{nf_elem, AnticNumberField}) -> Array{Int, 1}
#> For a non-zero element $a$ in factored form,
#> return the signs of all real embeddings.
#"""
function _signs(a::FacElem{nf_elem, AnticNumberField})
  r1, r2 = signature(base_ring(a))
  if r1 == 0
    return Int[]
  end

  s = ones(Int, r1)

  for (k, e) = a.fac
    if iseven(e)
      continue
    end
    s .*= _signs(k)
  end
  return s
end

@doc Markdown.doc"""
    complex_conjugation(K::AnticNumberField)

Given a totally complex normal number field, this function returns an
automorphism which is the restriction of complex conjugation at one embedding.
"""
function complex_conjugation(K::AnticNumberField; auts::Vector{NfToNfMor} = NfToNfMor[])
  if !isempty(auts)
    A = auts
  else
    A = automorphisms(K)
    if length(A) < degree(K)
      error("Number field must be normal")
    end
  end
  a = gen(K)
  d = length(A)
  !istotally_complex(K) && error("Number field must be totally complex")
  #First, quick and dirty. If only one automorphism works, then we return it
  p = 32
  while true
    c = conjugates(a, p)
    for i = 1:d
      if !isinvolution(A[i])
        continue
      end
      cc = conj.(conjugates(image_primitive_element(A[i]), p))
      for k = 1:d
        if overlaps(c[k], cc[k])
          found = true
          for j = 1:d
            if j == k
              continue
            end
            if overlaps(c[j], cc[k])
              found = false
              break
            end
          end
          if found
            return A[i]
          end
        end
      end
    end
    p = 2 * p
    if p > 2^18
      error("Precision too high in complex_conjugation")
    end
  end
  error("something went wrong!")
end

function _find_complex_conjugation(K::AnticNumberField, A::Vector{NfToNfMor})
  a = gen(K)
  #First, quick and dirty. If only one automorphism works, then we return it
  p = 32
  while true
    c = conjugates(a, p)
    overlap = false
    for i = 1:length(A)
      if !isinvolution(A[i])
        continue
      end
      cc = conj.(conjugates(image_primitive_element(A[i]), p))
      for k = 1:degree(K)
        if overlaps(c[k], cc[k])
          overlap = true
          found = true
          for j = 1:degree(K)
            if j == k
              continue
            end
            if overlaps(c[j], cc[k])
              found = false
              break
            end
          end
          if found
            return true, A[i]
          end
        end
      end
    end
    if !overlap
      break
    end
    if p > 2^18
      error("Precision too high in complex_conjugation")
    end
  end
  return false, A[1]
end

function iscomplex_conjugation(f::NfToNfMor)
  K = domain(f)
  @assert K == codomain(f)
  !istotally_complex(K) && error("Number field must be totally complex")
  p = 32
  d = degree(K)
  a = gen(K)
  img_a = image_primitive_element(f)
  while true
    c = conjugates(a, p)
    cc = conj.(conjugates(img_a, p))
    for i = 1:degree(K)
      if !overlaps(c[i], cc[i])
        return false
      end
    end
    #Now I need to assure that the precision is enough.
    found = true
    for j in 1:d-1
      if overlaps(c[d], cc[j])
        found = false
        break
      end
    end
    if found
      return true
    end
    p = 2 * p
    if p > 2^18
      error("Precision too high in complex_conjugation")
    end
  end

end