text:
#
# Generic FE space related methods

abstract type DistributedFESpace <: FESpace end

function FESpaces.get_vector_type(fs::DistributedFESpace)
  @abstractmethod
end

function FESpaces.get_free_dof_ids(fs::DistributedFESpace)
  @abstractmethod
end

function FESpaces.FEFunction(f::DistributedFESpace,::AbstractVector)
  @abstractmethod
end

function FESpaces.EvaluationFunction(f::DistributedFESpace,::AbstractVector)
  @abstractmethod
end

function FESpaces.get_fe_basis(f::DistributedFESpace)
  @abstractmethod
end

function FESpaces.get_trial_fe_basis(f::DistributedFESpace)
  @abstractmethod
end

function FESpaces.zero_free_values(f::DistributedFESpace)
  V = get_vector_type(f)
  vec = allocate_vector(V,get_free_dof_ids(f))
  fill!(vec,zero(eltype(vec)))
end

FESpaces.num_free_dofs(f::DistributedFESpace) = length(get_free_dof_ids(f))

function Base.zero(f::DistributedFESpace)
  free_values = zero_free_values(f)
  isconsistent = true
  FEFunction(f,free_values,isconsistent)
end

function dof_wise_to_cell_wise!(cell_wise_vector,dof_wise_vector,cell_to_ldofs,cell_prange)
  map_parts(cell_wise_vector,
          dof_wise_vector,
          cell_to_ldofs,
          cell_prange.partition) do cwv,dwv,cell_to_ldofs,partition
    cache  = array_cache(cell_to_ldofs)
    ncells = length(cell_to_ldofs)
    ptrs = cwv.ptrs
    data = cwv.data
    gdof = 0
    for cell in partition.oid_to_lid
      ldofs = getindex!(cache,cell_to_ldofs,cell)
      p = ptrs[cell]-1
      for (i,ldof) in enumerate(ldofs)
        if ldof > 0
          data[i+p] = dwv[ldof]
        end
      end
    end
  end
end

function cell_wise_to_dof_wise!(dof_wise_vector,cell_wise_vector,cell_to_ldofs,cell_range)
  map_parts(dof_wise_vector,
            cell_wise_vector,
            cell_to_ldofs,
            cell_range.partition) do dwv,cwv,cell_to_ldofs,partition

    gdof = 0
    cache = array_cache(cell_to_ldofs)
    for cell in partition.hid_to_lid
      ldofs = getindex!(cache,cell_to_ldofs,cell)
      p = cwv.ptrs[cell]-1
      for (i,ldof) in enumerate(ldofs)
        if ldof > 0
          dwv[ldof] = cwv.data[i+p]
        end
      end
    end
  end
end

function dof_wise_to_cell_wise(dof_wise_vector,cell_to_ldofs,cell_prange)
    cwv=map_parts(dof_wise_vector,cell_to_ldofs,cell_prange.partition) do dwv,cell_to_ldofs,partition
      cache = array_cache(cell_to_ldofs)
      ncells = length(cell_to_ldofs)
      ptrs = Vector{Int32}(undef,ncells+1)
      for cell in 1:ncells
        ldofs = getindex!(cache,cell_to_ldofs,cell)
        ptrs[cell+1] = length(ldofs)
      end
      PArrays.length_to_ptrs!(ptrs)
      ndata = ptrs[end]-1
      data = Vector{Int}(undef,ndata)
      gdof = 0
      PArrays.Table(data,ptrs)
    end
    dof_wise_to_cell_wise!(cwv,dof_wise_vector,cell_to_ldofs,cell_prange)
    cwv
end

function generate_gids(
cell_range::PRange,
cell_to_ldofs::AbstractPData{<:AbstractArray},
nldofs::AbstractPData{<:Integer})

neighbors = cell_range.exchanger.parts_snd
ngcells = length(cell_range)

# Find and count number owned dofs
ldof_to_part, nodofs = map_parts(
  cell_range.partition,cell_to_ldofs,nldofs) do partition,cell_to_ldofs,nldofs

  ldof_to_part = fill(Int32(0),nldofs)
  cache = array_cache(cell_to_ldofs)
  for cell in 1:length(cell_to_ldofs)
    owner = partition.lid_to_part[cell]
    ldofs = getindex!(cache,cell_to_ldofs,cell)
    for ldof in ldofs
      if ldof>0
        #TODO this simple approach concentrates dofs
        # in the last part and creates inbalances
        ldof_to_part[ldof] = max(owner,ldof_to_part[ldof])
      end
    end
  end
  nodofs = count(p->p==partition.part,ldof_to_part)
  ldof_to_part, nodofs
end

cell_ldofs_to_part = dof_wise_to_cell_wise(ldof_to_part,
                                            cell_to_ldofs,
                                            cell_range)
# Exchange the dof owners
exchange!(cell_ldofs_to_part,cell_range.exchanger)

cell_wise_to_dof_wise!(ldof_to_part,
                        cell_ldofs_to_part,
                        cell_to_ldofs,
                        cell_range)

# Find the global range of owned dofs
first_gdof, ngdofsplus1 = xscan(+,reduce,nodofs,init=1)
ngdofs = ngdofsplus1 - 1

# Distribute gdofs to owned ones
parts = get_part_ids(nodofs)
ldof_to_gdof = map_parts(
  parts,first_gdof,ldof_to_part) do part,first_gdof,ldof_to_part

  offset = first_gdof-1
  ldof_to_gdof = Vector{Int}(undef,length(ldof_to_part))
  odof = 0
  gdof = 0
  for (ldof,owner) in enumerate(ldof_to_part)
    if owner == part
      odof += 1
      ldof_to_gdof[ldof] = odof
    else
      ldof_to_gdof[ldof] = gdof
    end
  end
  for (ldof,owner) in enumerate(ldof_to_part)
    if owner == part
      ldof_to_gdof[ldof] += offset
    end
  end
  ldof_to_gdof
end

# Create cell-wise global dofs
cell_to_gdofs = dof_wise_to_cell_wise(ldof_to_gdof,
                                      cell_to_ldofs,
                                      cell_range)

# Exchange the global dofs
exchange!(cell_to_gdofs,cell_range.exchanger)

# Distribute global dof ids also to ghost
map_parts(
  parts,
  cell_to_ldofs,cell_to_gdofs,ldof_to_gdof,ldof_to_part,cell_range.partition) do part,
  cell_to_ldofs,cell_to_gdofs,ldof_to_gdof,ldof_to_part,partition

  gdof = 0
  cache = array_cache(cell_to_ldofs)
  for cell in partition.hid_to_lid
    ldofs = getindex!(cache,cell_to_ldofs,cell)
    p = cell_to_gdofs.ptrs[cell]-1
    for (i,ldof) in enumerate(ldofs)
      if ldof > 0 && ldof_to_part[ldof] == partition.lid_to_part[cell]
        ldof_to_gdof[ldof] = cell_to_gdofs.data[i+p]
      end
    end
  end
end

dof_wise_to_cell_wise!(cell_to_gdofs,ldof_to_gdof,cell_to_ldofs,cell_range)

exchange!(cell_to_gdofs,cell_range.exchanger)

cell_wise_to_dof_wise!(ldof_to_gdof,
                        cell_to_gdofs,
                        cell_to_ldofs,
                        cell_range)

# Setup dof partition
dof_partition = map_parts(parts,ldof_to_gdof,ldof_to_part) do part,ldof_to_gdof,ldof_to_part
  IndexSet(part,ldof_to_gdof,ldof_to_part)
end

# map_parts(parts,dof_partition,cell_to_ldofs,cell_to_gdofs,cell_range.partition) do part, partition, cell_to_ldofs,cell_to_gdofs, cell_range
#   if (part==3)
#     println("XXXX $(part)")
#     println(partition)
#     println(cell_to_ldofs)
#     println(cell_to_gdofs)
#     println(cell_range)
#   end
# end

# Setup dof exchanger
dof_exchanger = Exchanger(dof_partition,neighbors)

# Setup dof range
dofs = PRange(ngdofs,dof_partition,dof_exchanger)

dofs
end

# FEFunction related

struct DistributedFEFunctionData{T<:AbstractVector} <:GridapType
  free_values::T
end

const DistributedSingleFieldFEFunction = DistributedCellField{A,<:DistributedFEFunctionData{T}} where {A,T}

function FESpaces.get_free_dof_values(uh::DistributedSingleFieldFEFunction)
  uh.metadata.free_values
end

# Single field related

struct DistributedSingleFieldFESpace{A,B,C} <: DistributedFESpace
  spaces::A
  gids::B
  vector_type::Type{C}
  function DistributedSingleFieldFESpace(
    spaces::AbstractPData{<:SingleFieldFESpace},
    gids::PRange,
    vector_type::Type{C}) where C
    A = typeof(spaces)
    B = typeof(gids)
    new{A,B,C}(spaces,gids,vector_type)
  end
end

local_views(a::DistributedSingleFieldFESpace) = a.spaces

function FESpaces.get_vector_type(fs::DistributedSingleFieldFESpace)
  fs.vector_type
end

function FESpaces.get_free_dof_ids(fs::DistributedSingleFieldFESpace)
  fs.gids
end

function FESpaces.FEFunction(
  f::DistributedSingleFieldFESpace,free_values::AbstractVector,isconsistent=false)
  local_vals = consistent_local_views(free_values,f.gids,isconsistent)
  fields = map_parts(FEFunction,f.spaces,local_vals)
  metadata = DistributedFEFunctionData(free_values)
  DistributedCellField(fields,metadata)
end

function FESpaces.EvaluationFunction(
  f::DistributedSingleFieldFESpace,free_values::AbstractVector,isconsistent=false)
  local_vals = consistent_local_views(free_values,f.gids,isconsistent)
  fields = map_parts(EvaluationFunction,f.spaces,local_vals)
  metadata = DistributedFEFunctionData(free_values)
  DistributedCellField(fields,metadata)
end

function FESpaces.get_fe_basis(f::DistributedSingleFieldFESpace)
  fields = map_parts(get_fe_basis,f.spaces)
  DistributedCellField(fields)
end

function FESpaces.get_trial_fe_basis(f::DistributedSingleFieldFESpace)
  fields = map_parts(get_trial_fe_basis,f.spaces)
  DistributedCellField(fields)
end

function FESpaces.TrialFESpace(f::DistributedSingleFieldFESpace)
  spaces = map_parts(TrialFESpace,f.spaces)
  DistributedSingleFieldFESpace(spaces,f.gids,f.vector_type)
end

function FESpaces.TrialFESpace(f::DistributedSingleFieldFESpace,fun)
  spaces = map_parts(f.spaces) do s
    TrialFESpace(s,fun)
  end
  DistributedSingleFieldFESpace(spaces,f.gids,f.vector_type)
end

function FESpaces.TrialFESpace(fun,f::DistributedSingleFieldFESpace)
  spaces = map_parts(f.spaces) do s
    TrialFESpace(fun,s)
  end
  DistributedSingleFieldFESpace(spaces,f.gids,f.vector_type)
end

function generate_gids(
  model::DistributedDiscreteModel,
  spaces::AbstractPData{<:SingleFieldFESpace})
  cell_to_ldofs = map_parts(get_cell_dof_ids,spaces)
  nldofs = map_parts(num_free_dofs,spaces)
  generate_gids(model.gids,cell_to_ldofs,nldofs)
end

function FESpaces.interpolate(u,f::DistributedSingleFieldFESpace)
  free_values = zero_free_values(f)
  interpolate!(u,free_values,f)
end

function FESpaces.interpolate!(
  u,free_values::AbstractVector,f::DistributedSingleFieldFESpace)
  map_parts(f.spaces,local_views(free_values)) do V,vec
    interpolate!(u,vec,V)
  end
  FEFunction(f,free_values)
end

# Factories

function FESpaces.FESpace(model::DistributedDiscreteModel,reffe;kwargs...)
  spaces = map_parts(model.models) do m
    FESpace(m,reffe;kwargs...)
  end
  gids =  generate_gids(model,spaces)
  vector_type = _find_vector_type(spaces,gids)
  DistributedSingleFieldFESpace(spaces,gids,vector_type)
end

function _find_vector_type(spaces,gids)
  #TODO Now the user can select the local vector type but not the global one
  # new kw-arg global_vector_type ?
  # we use PVector for the moment
  local_vector_type = get_vector_type(get_part(spaces))
  T = eltype(local_vector_type)
  A = typeof(map_parts(i->local_vector_type(undef,0),gids.partition))
  B = typeof(gids)
  vector_type = PVector{T,A,B}
end

# Assembly

function FESpaces.collect_cell_matrix(
  trial::DistributedFESpace,
  test::DistributedFESpace,
  a::DistributedDomainContribution)
  map_parts(
    collect_cell_matrix,
    local_views(trial),
    local_views(test),
    local_views(a))
end

function FESpaces.collect_cell_vector(
  test::DistributedFESpace, a::DistributedDomainContribution)
  map_parts(
    collect_cell_vector,local_views(test),local_views(a))
end

function FESpaces.collect_cell_matrix_and_vector(
  trial::DistributedFESpace,
  test::DistributedFESpace,
  biform::DistributedDomainContribution,
  liform::DistributedDomainContribution)
  map_parts(collect_cell_matrix_and_vector,
    local_views(trial),
    local_views(test),
    local_views(biform),
    local_views(liform))
end

function FESpaces.collect_cell_matrix_and_vector(
  trial::DistributedFESpace,
  test::DistributedFESpace,
  biform::DistributedDomainContribution,
  liform::DistributedDomainContribution,
  uhd)
  map_parts(collect_cell_matrix_and_vector,
    local_views(trial),
    local_views(test),
    local_views(biform),
    local_views(liform),
    local_views(uhd))
end

function FESpaces.collect_cell_vector(
  test::DistributedFESpace,l::Number)
  map_parts(local_views(test)) do s
    collect_cell_vector(s,l)
  end
end

function FESpaces.collect_cell_matrix_and_vector(
  trial::DistributedFESpace,
  test::DistributedFESpace,
  mat::DistributedDomainContribution,
  l::Number)
  map_parts(
    local_views(trial),local_views(test),local_views(mat)) do u,v,m
    collect_cell_matrix_and_vector(u,v,m,l)
  end
end

function FESpaces.collect_cell_matrix_and_vector(
  trial::DistributedFESpace,
  test::DistributedFESpace,
  mat::DistributedDomainContribution,
  l::Number,
  uhd)
  map_parts(
    local_views(trial),local_views(test),local_views(mat),local_views(uhd)) do u,v,m,f
    collect_cell_matrix_and_vector(u,v,m,l,f)
  end
end

struct DistributedSparseMatrixAssembler{A,B,C,D,E,F} <: SparseMatrixAssembler
  strategy::A
  assems::B
  matrix_builder::C
  vector_builder::D
  rows::E
  cols::F
end

local_views(a::DistributedSparseMatrixAssembler) = a.assems

FESpaces.get_rows(a::DistributedSparseMatrixAssembler) = a.rows
FESpaces.get_cols(a::DistributedSparseMatrixAssembler) = a.cols
FESpaces.get_matrix_builder(a::DistributedSparseMatrixAssembler) = a.matrix_builder
FESpaces.get_vector_builder(a::DistributedSparseMatrixAssembler) = a.vector_builder
FESpaces.get_assembly_strategy(a::DistributedSparseMatrixAssembler) = a.strategy

function FESpaces.symbolic_loop_matrix!(A,a::DistributedSparseMatrixAssembler,matdata)
  map_parts(symbolic_loop_matrix!,local_views(A,a.rows,a.cols),a.assems,matdata)
end

function FESpaces.numeric_loop_matrix!(A,a::DistributedSparseMatrixAssembler,matdata)
  map_parts(numeric_loop_matrix!,local_views(A,a.rows,a.cols),a.assems,matdata)
end

function FESpaces.symbolic_loop_vector!(b,a::DistributedSparseMatrixAssembler,vecdata)
  map_parts(symbolic_loop_vector!,local_views(b,a.rows),a.assems,vecdata)
end

function FESpaces.numeric_loop_vector!(b,a::DistributedSparseMatrixAssembler,vecdata)
  map_parts(numeric_loop_vector!,local_views(b,a.rows),a.assems,vecdata)
end

function FESpaces.symbolic_loop_matrix_and_vector!(A,b,a::DistributedSparseMatrixAssembler,data)
  map_parts(symbolic_loop_matrix_and_vector!,local_views(A,a.rows,a.cols),local_views(b,a.rows),a.assems,data)
end

function FESpaces.numeric_loop_matrix_and_vector!(A,b,a::DistributedSparseMatrixAssembler,data)
  map_parts(numeric_loop_matrix_and_vector!,local_views(A,a.rows,a.cols),local_views(b,a.rows),a.assems,data)
end

# Parallel Assembly strategies

function local_assembly_strategy(::SubAssembledRows,rows,cols)
  DefaultAssemblyStrategy()
end

# When using this one, make sure that you also loop over ghost cells.
# This is at your own risk.
function local_assembly_strategy(::FullyAssembledRows,rows,cols)
  rows_lid_to_ohid = rows.lid_to_ohid
  GenericAssemblyStrategy(
    identity,
    identity,
    row->rows_lid_to_ohid[row]>0,
    col->true)
end

# Assembler high level constructors
function FESpaces.SparseMatrixAssembler(
  local_mat_type,
  local_vec_type,
  trial::DistributedFESpace,
  test::DistributedFESpace,
  par_strategy=SubAssembledRows())

  Tv = local_vec_type
  T = eltype(Tv)
  Tm = local_mat_type
  cols = trial.gids.partition
  rows = test.gids.partition
  assems = map_parts(local_views(test),local_views(trial),rows,cols) do v,u,rows,cols
    local_strategy = local_assembly_strategy(par_strategy,rows,cols)
    SparseMatrixAssembler(Tm,Tv,u,v,local_strategy)
  end
  matrix_builder = PSparseMatrixBuilderCOO(Tm,par_strategy)
  vector_builder = PVectorBuilder(Tv,par_strategy)
  rows = get_free_dof_ids(test)
  cols = get_free_dof_ids(trial)
  DistributedSparseMatrixAssembler(par_strategy,assems,matrix_builder,vector_builder,rows,cols)
end

function FESpaces.SparseMatrixAssembler(
  trial::DistributedFESpace,
  test::DistributedFESpace,
  par_strategy=SubAssembledRows())
  Tv = get_vector_type(get_part(local_views(trial)))
  T = eltype(Tv)
  Tm = SparseMatrixCSC{T,Int}
  SparseMatrixAssembler(Tm,Tv,trial,test,par_strategy)
end


text:
module FMRIOpto

using Reexport
@reexport using DataFrames, CSV, CategoricalArrays, BrowseTables
@reexport using Statistics, StatsBase, Plots, StatsPlots, HypothesisTests
@reexport using MixedModels

include("utilities.jl")
include("functions.jl")

export nanmean, nansem, safemean, safesem, iscolumn, model_res
export average_presentation, define_bold, plot_bold, odour_response, stim_interaction, analysis_bold

end # module

ext: jl


text:
struct ℝ{n} <: Manifold end

dim(::Type{ℝ{n}}) where n = n

# Creates aliases R1, TR1, etc.
for n = 1:12
    @eval const $(Symbol("R$n")) = ℝ{$n}
    @eval const $(Symbol("TR$n")) = T{ℝ{$n}}
    @eval export $(Symbol("R$n"))
    @eval export $(Symbol("TR$n"))
end
TR{n} = T{ℝ{n}}

# Assume all charts are global for now
isglobal(::Chart{k,ℝ{n}}) where {k,n} = true

emb_to_chart(p, C::Chart{1,ℝ{n}}) where n = p
chart_to_emb(p, C::Chart{1,ℝ{n}}) where n = p

# Not a diffeomorphism for all of R^n, but useful for manual checks
emb_to_chart(p, C::Chart{2,ℝ{n}}) where n = p.^2
chart_to_emb(p, C::Chart{2,ℝ{n}}) where n = p.^(1/2)

# Not a diffeomorphism for all of R^n, but useful for manual checks
emb_to_chart(p, C::Chart{3,ℝ{n}}) where n = p.^3
chart_to_emb(p, C::Chart{3,ℝ{n}}) where n = p.^(1/3)

emb_to_chart(p, C::Chart{4,ℝ{n}}) where n = p.^3 + p
function chart_to_emb(p, C::Chart{4,ℝ{n}}) where n
    a = sqrt(3)*sqrt.(27*p.^2 .+ 4) .+ 9*p
    (a/18).^(1/3) .- ((2/3)./a).^(1/3)
end

global_frame_coeffs(p, C::Chart{1,ℝ{n}}) where n = SMatrix{n,n,eltype(p)}(I)
global_coframe_coeffs(p, C::Chart{1,ℝ{n}}) where n = SMatrix{n,n,eltype(p)}(I)

function global_frame_coeffs(p, C::Chart{I,ℝ{n}}) where {n,I}
    C1 = Chart{1,ℝ{n}}()
    p1 = chart_transition(p, C, C1)
    B1 = global_frame_coeffs(p, C1)
    global_frame_coeffs_chart_transition(p1, B1, C1, C)
end

function global_coframe_coeffs(p, C::Chart{I,ℝ{n}}) where {n,I}
    C1 = Chart{1,ℝ{n}}()
    p1 = chart_transition(p, C, C1)
    β1 = global_frame_coeffs(p, C1)
    global_coframe_coeffs_chart_transition(p1, β1, C1, C)
end

inv_global_frame_coeffs(p, C) = inv(global_coframe_coeffs(p, C))
inv_global_coframe_coeffs(p, C) = inv(global_frame_coeffs(p, C))


text:
using Manifolds, Manopt, Test, ManifoldsBase

@testset "gradients" begin
    @testset "Circle (Allocating)" begin
        M = Circle()
        N = PowerManifold(M, 4)
        x = [0.1, 0.2, 0.3, 0.5]
        tvTestξ = [-1.0, 0.0, 0.0, 1.0]
        @test grad_TV(N, x) == tvTestξ
        @test grad_TV(M, (x[1], x[1])) ==
              (zero_tangent_vector(M, x[1]), zero_tangent_vector(M, x[1]))
        @test norm(N, x, grad_TV(N, x, 2) - tvTestξ) ≈ 0
        tv2Testξ = [0.0, 0.5, -1.0, 0.5]
        @test grad_TV2(N, x) == tv2Testξ
        @test norm(N, x, forward_logs(N, x) - [0.1, 0.1, 0.2, 0.0]) ≈ 0 atol = 10^(-16)
        @test norm(
            N,
            x,
            grad_intrinsic_infimal_convolution_TV12(N, x, x, x, 1.0, 1.0)[1] -
            [-1.0, 0.0, 0.0, 1.0],
        ) ≈ 0
        @test norm(N, x, grad_intrinsic_infimal_convolution_TV12(N, x, x, x, 1.0, 1.0)[2]) ≈
              0
        x2 = [0.1, 0.2, 0.3]
        N2 = PowerManifold(M, size(x2)...)
        @test grad_TV2(N2, x2) == zeros(3)
        @test grad_TV2(N2, x2, 2) == zeros(3)
        @test grad_TV(M, (0.0, 0.0), 2) == (0.0, 0.0)
        # 2d forward logs
        N3 = PowerManifold(M, 2, 2)
        N3C = PowerManifold(M, 2, 2, 2)
        x3 = [0.1 0.2; 0.3 0.5]
        x3C = cat(x3, x3; dims=3)
        tC = cat([0.2 0.3; 0.0 0.0], [0.1 0.0; 0.2 0.0]; dims=3)
        @test norm(N3C, x3C, forward_logs(N3, x3) - tC) ≈ 0 atol = 10^(-16)

        M = Circle()
        p = 0
        q = π / 4
        @test grad_distance(M, p, q) == q - p
    end
    @testset "Sphere (Mutating)" begin
        M = Sphere(2)
        p = [0.0, 0.0, 1.0]
        q = [0.0, 1.0, 0.0]
        r = [1.0, 0.0, 0.0]
        @testset "Gradient of the distance function" begin
            X = zero_tangent_vector(M, p)
            grad_distance!(M, X, p, q)
            Y = grad_distance(M, p, q)
            Z = [0.0, 0.0, -π / 2] # known solution
            @test X == Y
            @test X == Z
        end
        @testset "Gradient of total variation" begin
            Y = grad_TV(M, (p, q))
            Z = [[0.0, -1.0, 0.0], [0.0, 0.0, -1.0]]
            X = similar.(Z)
            grad_TV!(M, X, (p, q))
            @test [y for y in Y] == X
            @test [y for y in Y] ≈ Z
            N = PowerManifold(M, NestedPowerRepresentation(), 3)
            s = [p, q, r]
            Y2 = grad_TV(N, s)
            Z2 = [[0.0, -1.0, 0.0], [-1.0, 0.0, -1.0], [0.0, -1.0, 0.0]]
            X2 = zero_tangent_vector(N, s)
            grad_TV!(N, X2, s)
            @test Y2 == Z2
            @test X2 == Z2
            Y2a = grad_TV(N, s, 2)
            X2a = zero_tangent_vector(N, s)
            grad_TV!(N, X2a, s, 2)
            @test Y2a == X2a
            N2 = PowerManifold(M, NestedPowerRepresentation(), 2)
            Y3 = grad_TV(M, (p, q), 2)
            X3 = zero_tangent_vector(N2, [p, q])
            grad_TV!(M, X3, (p, q), 2)
            @test [y for y in Y3] == X3
            Y4 = grad_TV(M, (p, p))
            X4 = zero_tangent_vector(N2, [p, q])
            grad_TV!(M, X4, (p, p))
            @test [y for y in Y4] == X4
        end
        @testset "Grad of second order total variation" begin
            N = PowerManifold(M, NestedPowerRepresentation(), 3)
            s = [p, q, r]
            X = zero_tangent_vector(N, s)
            grad_TV2!(M, X, s)
            Y = grad_TV2(M, s)
            Z = -1 / sqrt(2) .* [[0.0, 1.0, 0.0], [1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]
            @test Y == X
            @test Y ≈ Z
            Y2 = grad_TV2(M, s, 2)
            Z2 = -1.110720734539 .* [[0.0, 1.0, 0.0], [1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]
            @test Y2 ≈ Z2
            s2 = [p, shortest_geodesic(M, p, q, 0.5), q]
            @test grad_TV2(M, s2) == [zero_tangent_vector(M, se) for se in s2]
        end
    end
end


text:
#=
Here we solve the equation:
```math
 q + dot(∇, uq) = 0
 p - dot(∇, up) = 0
```
on a sphere to test the conservation of the numerics

The boundary conditions are `p = q` when `dot(n, u) > 0` and
`q = p` when `dot(n, u) < 0` (i.e., `p` flows into `q` and vice-sersa).
=#

using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGBalanceLawDiscretizations
using CLIMA.DGBalanceLawDiscretizations.NumericalFluxes
using CLIMA.MPIStateArrays
using CLIMA.ODESolvers
using CLIMA.GenericCallbacks
using LinearAlgebra
using StaticArrays
using Logging, Printf, Dates
using Random

const _nstate = 2
const _q, _p = 1:_nstate
const stateid = (qid = _q, pid = _p)
const statenames = ("q", "p")

const _nauxstate = 3
@inline function velocity_initialization!(vel, x, y, z)
    @inbounds begin
        r = x^2 + y^2 + z^2
        vel[1] = cos(10 * π * x) * sin(10 * π * y) + cos(20 * π * z)
        vel[2] = exp(sin(π * r))
        vel[3] = sin(π * (x + y + z))
    end
end

# physical flux function
@inline function flux!(F, Q, _, vel, _)
    @inbounds begin
        u, v, w = vel[1], vel[2], vel[3]
        q, p = Q[_q], Q[_p]

        F[1, _q], F[2, _q], F[3, _q] = u * q, v * q, w * q
        F[1, _p], F[2, _p], F[3, _p] = -u * p, -v * p, -w * p
    end
end
# physical flux function
@inline function numerical_flux!(F, nM, QM, _, velM, QP, _, velP, _)
    @inbounds begin
        uM, vM, wM = velM[1], velM[2], velM[3]
        unM = nM[1] * uM + nM[2] * vM + nM[3] * wM
        uP, vP, wP = velP[1], velP[2], velP[3]
        unP = nM[1] * uP + nM[2] * vP + nM[3] * wP
        un = (unP + unM) / 2

        if un > 0
            F[_q], F[_p] = un * QM[_q], -un * QP[_p]
        else
            F[_q], F[_p] = un * QP[_q], -un * QM[_p]
        end
    end
end
@inline function numerical_boundary_flux!(F, nM, QM, _, vel, QP, _, _, _, _)
    @inbounds begin
        u, v, w = vel[1], vel[2], vel[3]
        un = nM[1] * vel[1] + nM[2] * vel[2] + nM[3] * vel[3]

        if un > 0
            F[_q], F[_p] = un * QM[_q], -un * QM[_q]
        else
            F[_q], F[_p] = un * QM[_p], -un * QM[_p]
        end
    end
end

function run(mpicomm, ArrayType, N, Nhorz, Rrange, timeend, FT, dt)

    topl = StackedCubedSphereTopology(mpicomm, Nhorz, Rrange)

    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
        meshwarp = Topologies.cubedshellwarp,
    )

    # spacedisc = data needed for evaluating the right-hand side function
    spacedisc = DGBalanceLaw(
        grid = grid,
        length_state_vector = _nstate,
        flux! = flux!,
        numerical_flux! = numerical_flux!,
        numerical_boundary_flux! = numerical_boundary_flux!,
        auxiliary_state_length = _nauxstate,
        auxiliary_state_initialization! = velocity_initialization!,
    )

    # This is a actual state/function that lives on the grid
    initialcondition!(Q, x, y, z, _...) = begin
        @inbounds Q[_q], Q[_p] = rand(), rand()
    end
    Q = MPIStateArray(spacedisc, initialcondition!)

    lsrk = LSRK54CarpenterKennedy(spacedisc, Q; dt = dt, t0 = 0)

    eng0 = norm(Q)
    sum0 = weightedsum(Q)
    @info @sprintf """Starting
    norm(Q₀) = %.16e
    sum(Q₀) = %.16e""" eng0 sum0

    max_mass_loss = FT(0)
    max_mass_gain = FT(0)
    cbmass = GenericCallbacks.EveryXSimulationSteps(1) do
        cbsum = weightedsum(Q)
        max_mass_loss = max(max_mass_loss, sum0 - cbsum)
        max_mass_gain = max(max_mass_gain, cbsum - sum0)
        nothing
    end
    solve!(Q, lsrk; timeend = timeend, callbacks = (cbmass,))

    # Print some end of the simulation information
    engf = norm(Q)
    sumf = weightedsum(Q)
    @info @sprintf """Finished
    norm(Q)            = %.16e
    norm(Q) / norm(Q₀) = %.16e
    norm(Q) - norm(Q₀) = %.16e
    max mass loss      = %.16e
    max mass gain      = %.16e
    initial mass       = %.16e
    """ engf engf / eng0 engf - eng0 max_mass_loss max_mass_gain sum0
    max(max_mass_loss, max_mass_gain) / sum0
end

using Test
let
    CLIMA.init()
    ArrayTypes = (CLIMA.array_type(),)

    mpicomm = MPI.COMM_WORLD

    dt = 1e-4
    timeend = 100 * dt

    polynomialorder = 4

    Nhorz = 4
    Rrange = 1.0:0.25:2.0

    dim = 3
    @testset "$(@__FILE__)" for ArrayType in ArrayTypes
        for FT in (Float64,) #Float32)
            Random.seed!(0)
            @info (ArrayType, FT, dim)
            delta_mass = run(
                mpicomm,
                ArrayType,
                polynomialorder,
                Nhorz,
                Rrange,
                timeend,
                FT,
                dt,
            )
            @test abs(delta_mass) < 1e-15
        end
    end
end

nothing

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
default_type(ts::Type...) = default_type(promote_type(ts...))
default_type(::Type{T}) where {T} = T
default_type(::Type{<:Number}) = Float64

recode(enc::AbstractEncoding, enc_new::AbstractEncoding, x) =
    _recode(enc, enc_new, x)

Broadcast.broadcasted(::typeof(recode), enc, enc_new, x) =
    broadcast(_recode, Ref(enc), Ref(enc_new), x)

_recode(enc::TwoClassEncoding, enc_new::TwoClassEncoding, x) =
    ispositive(enc, x) ? positive_label(enc_new) : negative_label(enc_new)

classify(enc::TwoClassEncoding, score, t) =
    _classify(enc, score, t)

Broadcast.broadcasted(::typeof(classify), enc, score, t) =
    broadcast(_classify, Ref(enc), score, t)

_classify(enc::TwoClassEncoding, score, t) =
    score .>= t ? positive_label(enc) : negative_label(enc)


text:
using DataFrames, InteractiveFixedEffectModels, CSV, LinearAlgebra, Test

precision = 2e-1
df = DataFrame(CSV.File(joinpath(dirname(pathof(InteractiveFixedEffectModels)), "../dataset/Cigar.csv")))

for method in [:dogleg, :levenberg_marquardt, :gauss_seidel]
	println(method)
	model = @formula Sales ~ Price + ife(State, Year, 1)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ [328.1653237715761, -1.0415042260420706] .- 1)  < precision
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.228 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 851.514 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -9.7870 - 1) < precision
	@test result.r2_within > 0.0

	show(result)
	model = @formula Sales ~ Price + ife(State, Year, 2)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ [163.01350, -0.40610] .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.227 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 184.1897 - 1) < precision
	@test norm(result.augmentdf[1, :residuals]  / -1.774 - 1) < precision
	@test result.r2_within > 0.0

	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ -0.425389 .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1])  / 0.2474 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 123.460 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -5.222 - 1) < precision
	@test result.r2_within > 0.0

	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(Year)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ -0.3744296120563005 .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1])  / 0.1918 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 102.336 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -2.2153 - 1) < precision
	@test result.r2_within > 0.0

	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State) + fe(Year)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ -0.524157 .- 1) < precision
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.256 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 60.0481 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -5.614 - 1) < precision
	@test result.r2_within > 0.0

	# subset
	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State)
	result = regife(df, model, method = method, subset = df.State .<= 30, save = true)
	@test size(result.augmentdf, 1) == size(df, 1)
	@test norm(abs(result.augmentdf[1, :factors1]) / 0.25965 - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings1]) / 107.832 - 1) < precision
	@test norm(abs(result.augmentdf[1, :factors2])  / 0.2655  - 1) < precision
	@test norm(abs(result.augmentdf[1, :loadings2]) / 15.551 - 1) < precision
	@test norm(result.augmentdf[1, :residuals] / -3.8624  - 1) < precision
	@test norm(result.augmentdf[1, :fe_State] / 131.6162 - 1) < precision
end

# check high dimentional fixed effects are part of factor models
df.State2 = categorical(df.State)
@test_throws ErrorException regife(df, @formula(Sales ~ Price + ife(State, Year, 2) + fe(State2)))

# local minima
using DataFrames, InteractiveFixedEffectModels, Test
precision = 2e-1

for method in [:levenberg_marquardt, :dogleg]
	println(method)

	local df = DataFrame(CSV.File(joinpath(dirname(pathof(InteractiveFixedEffectModels)), "../dataset/Cigar.csv")))
	df.State = categorical(df.State)
	df.Year = categorical(df.Year)
	model = @formula Sales ~ Price + ife(State, Year, 1) + fe(State)
	result = regife(df, model, weights = :Pop, method = method, save = true)
	model = @formula Sales ~ Price + ife(State, Year, 2) + fe(State)
	result = regife(df, model, weights = :Pop, method = method, save = true)

	local df = DataFrame(CSV.File(joinpath(dirname(pathof(InteractiveFixedEffectModels)), "../dataset/EmplUK.csv")))
	df.id1 = df.Firm
	df.id2 = df.Year
	df.pid1 = categorical(df.id1)
	df.pid2 = categorical(df.id2)
	df.y = df.Wage
	df.x1 = df.Emp
	df.w = df.Output
	model = @formula y ~ x1 + ife(pid1, pid2, 2)
	result = regife(df, model, method = method, save = true)
	@test norm(result.coef ./ [4.53965, -0.0160858] .- 1) < precision
	model = @formula y ~ x1 + ife(pid1, pid2, 2)
	result = regife(df, model, weights = :w, method = method, save = true)
	@test norm(result.coef ./ [3.47551,-0.017366] .- 1) < precision
	model = @formula y ~ x1 + ife(pid1, pid2, 1)
	result = regife(df, model, weights = :w, method = method, save = true)
	@test norm(result.coef ./ [ -2.62105, -0.0470005] .- 1) < precision
end

# Check formula can be constructed programatically
using StatsModels
result1 = regife(df, @formula(Sales ~ NDI + fe(State) + ife(State, Year, 2)))
result2 = regife(df, Term(:Sales) ~ Term(:NDI) + fe(Term(:State)) + ife(Term(:State), Term(:Year), 2))
@test result1.coef ≈ result2.coef


text:
"""
    abstract type AbstractNuclearModel

Supertype for all (radially symmetric) nuclear models. The underlying charge distribution of
the nuclear model is assumed to be normalized to unity (i.e. corresponding to ``Z = 1``).

Each concrete nuclear model of type `T <: AbstractNuclearModel` should implement the
following method

* [`potential(::T, r::Real) -> Float64`](@ref potential):
  return the value of the potential at the radius `r`.

Additional, each nuclear model may also implement

* [`density(::T, r::Real) -> Float64`](@ref density):
  return the value of the (normalized) charge density at the radius `r`.

* [`rms(::T)`](@ref rms):
  return the root-mean-square radius ``\\sqrt{\\langle r^2 \\rangle}`` of the underlying
  charge distribution.

* [`from_rms(::Type{T}, rms)`](@ref from_rms):
  construct an instance of the nuclear model with the specified root-mean-square radius.

!!! note "Units"

    All quantities are assumed to be in atomic units.
"""
abstract type AbstractNuclearModel end
Base.Broadcast.broadcastable(x::AbstractNuclearModel) = Ref(x)

"""
    potential(::AbstractNuclearModel, r::Real) -> Float64

Return the value of the (normalized) charge density at the radius `r`.
"""
function potential end

"""
    density(::AbstractNuclearModel, r::Real) -> Float64

Return the value of the (normalized) charge density at the radius `r`.
"""
function density end

"""
    rms(::AbstractNuclearModel)

Return the root-mean-square radius ``\\sqrt{\\langle r^2 \\rangle}`` of the underlying
charge distribution.
"""
function rms end

"""
    from_rms(::Type{<:AbstractNuclearModel}, rms)

Construct an instance of the nuclear model with the specified root-mean-square radius.
"""
function from_rms end

"""
    struct PointNucleus <: AbstractNuclearModel

Represent a point nucleus with the potential

```math
V(r) = - Z / r,
\\quad
\\rho(r) = \\frac{Z}{4\\pi r^2} \\delta(r)
```

This nuclear model has no parameters. Also, this model does not define a method for
`density` due to the delta-function nature of the charge distribution.

# Constructors

* `PointNucleus()`: construct an instance of `PointNucleus`

# Examples

```jldoctest
julia> using AtomicMiscellany: PointNucleus, rms, potential

julia> rms(PointNucleus())
0.0

julia> potential(PointNucleus(), 2)
-0.5
```
"""
struct PointNucleus <: AbstractNuclearModel end
function potential(::PointNucleus, r::Real)
    r >= 0 || throw(DomainError(r, "r must be positive"))
    return -1/r
end
rms(::PointNucleus) = 0.0

"""
    struct UniformShellNucleus <: AbstractNuclearModel

Represents a nuclear model where all the charge is unformly distributed over an infinitely
thin shell at radius `R`. The nuclear potential and charge distribution are given by

```math
V(r) = \\begin{cases}
    -Z/R, & 0 \\leq r \\leq R \\\\
    -Z/r, & r > R
\\end{cases},
\\quad
\\rho(r) = \\frac{Z}{4\\pi r^2} \\delta(r - R)
```

Note that this model does not define a method for `density` due to the delta-function
nature of the charge distribution.

## Constructors

* `UniformShellNucleus(R::Real)`: construct a shell nucleus with radius `R`.

* `from_rms(UniformShellNucleus, rms)`: constructs a shell nucleus with the radius
  determined from the root-mean-square radius ``\\sqrt{\\langle r^2 \\rangle}`` of the
  charge distribution

  ```math
  R = \\sqrt{\\langle r^2 \\rangle}
  ```
"""
struct UniformShellNucleus <: AbstractNuclearModel
    R :: Float64
    function UniformShellNucleus(R::Real)
        R > 0 || throw(DomainError(R, "R must have a positive value"))
        new(R)
    end
end
from_rms(::Type{UniformShellNucleus}, rms) = UniformShellNucleus(rms)
function potential(m::UniformShellNucleus, r::Real)
    r >= 0 || throw(DomainError(r, "r must be positive"))
    r < m.R ? -1/m.R : -1/r
end
rms(m::UniformShellNucleus) = m.R

"""
    struct UniformSphericalNucleus <: AbstractNuclearModel

Represents a nuclear charge distribution where the charge is homogeneously distributed in
a sphere of radius ``R``. The potential and charge distributions are given by

```math
$(raw"""
V(r) = \begin{cases}
    -\frac{Z}{2R}\left[3 - \left(\frac{r}{R}\right)^2\right], & 0 \leq r \leq R \\
    -\frac{Z}{r}, & r > R
\end{cases},
\quad
\rho(r) = \begin{cases}
    \frac{3Z}{4\pi R^3}, & 0 \leq r \leq R \\
    0, & r > R
\end{cases}
""")
```

## Constructors

* `UniformSphericalNucleus(R::Real)`: construct a homogeneous spherical nucleus with radius
  `R`.

* `from_rms(UniformSphericalNucleus, rms)`: constructs a homogeneous spherical nucleus with
  the radius determined from the root-mean-square radius ``\\sqrt{\\langle r^2 \\rangle}``
  of the charge distribution

  ```math
  R = \\sqrt{\\frac{5}{3}} \\sqrt{\\langle r^2 \\rangle}
  ```

"""
struct UniformSphericalNucleus <: AbstractNuclearModel
    R :: Float64
    function UniformSphericalNucleus(R::Real)
        R > 0 || throw(DomainError(R, "R must have a positive value"))
        new(R)
    end
end
from_rms(::Type{UniformSphericalNucleus}, rms) = UniformSphericalNucleus(rms*sqrt(5/3))
function potential(m::UniformSphericalNucleus, r::Real)
    r >= 0 || throw(DomainError(r, "r must be positive"))
    R = m.R
    r < R ? ((r / R)^2 - 3) / (2 * R) : -1/r
end
function density(m::UniformSphericalNucleus, r::Real)
    r >= 0 || throw(DomainError(r, "r must be positive"))
    R = m.R
    r <= R ? 3 / (4π * R^3) : 0.0
end
rms(m::UniformSphericalNucleus) = m.R * sqrt(3/5)

"""
    struct GaussianNucleus <: AbstractNuclearModel

Represents a Gaussian-shaped nuclear charge distribution, size of which is determined by
the parameter ``R```. The potential and charge distributions are given by

```math
$(raw"""
V(r) = - \frac{Z}{r} \mathrm{erf}\left(\frac{r}{R}\right),
\quad
\rho(r) = \frac{Z}{\pi^{3/2} R^3}\exp\left[-\left(\frac{r}{R}\right)^2\right]
""")
```

## Constructors

* `GaussianNucleus(R::Real)`: constructs a Gaussian nuclear model of size `R`.

* `from_rms(UniformSphericalNucleus, rms)`: constructs a Gaussian nucleus with the ``R``
  parameter determined from the root-mean-square radius ``\\sqrt{\\langle r^2 \\rangle}``
  of the charge distribution

  ```math
  R = \\sqrt{\\frac{2}{3}} \\sqrt{\\langle r^2 \\rangle}
  ```

"""
struct GaussianNucleus <: AbstractNuclearModel
    R :: Float64

    function GaussianNucleus(R::Real)
        R > 0 || throw(DomainError(R, "R must have a positive value"))
        new(R)
    end
end
from_rms(::Type{GaussianNucleus}, rms) = GaussianNucleus(rms * sqrt(2/3))
function potential(m::GaussianNucleus, r::Real)
    r >= 0 || throw(DomainError(r, "r must be positive"))
    R = m.R
    # we need to special case r=0 due to the 1/r behaviour, even though the limit is fine.
    # We can get the correct value by considering the Taylor series of erf()
    # erf(x) = 2/√π (x - x^3 / 3 + ...)
    r > 0 ? - erf(r/R) / r : - 2/(R * sqrt(π))
end
function density(m::GaussianNucleus, r::Real)
    r >= 0 || throw(DomainError(r, "r must be positive"))
    R = m.R
    exp(-(r/R)^2) / (π^(3/2) * R^3)
end
rms(m::GaussianNucleus) = m.R * sqrt(3/2)

# """
#     struct FermiNucleus <: AbstractNuclearModel

# ...
# """
# struct FermiNucleus <: AbstractNuclearModel
#     c :: Float64
#     b :: Float64

#     function FermiNucleus() error("Not implemented.") end
# end
# from_rms(::Type{FermiNucleus}, rms) = PointNucleus() #FermiNucleus(0.0, 0.0)
# function potential(m::FermiNucleus, r::Real)
#     r >= 0 || throw(DomainError(r, "r must be positive"))
#     0.0
# end
# function density(m::FermiNucleus, r::Real)
#     r >= 0 || throw(DomainError(r, "r must be positive"))
#     0.0
# end
# rms(m::FermiNucleus) = m.R * sqrt(3/2)

"""
    struct Andrae2000

Type to dispatch on data from the following paper:

* Andrae, D. _"Finite Nuclear Charge Density Distributions in Electronic Structure
  Calculations for Atoms and Molecules."_ Physics Reports 336, no. 6 (October 2000):
  413–525. <https://doi.org/10.1016/S0370-1573(00)00007-7>.
"""
struct Andrae2000 end


text:
################################################################################
#
#   pohligHellman.jl : functions for Pohlig-Hellman algorithm
#
################################################################################

export pohligHellman

"""
    pohligHellmanPrime{T <: PolyElem}(prime::Int, n::Int, k::Int,
                                      mid::T, elem::T, defPol::T)

Compute the discrete logarithm of `elem` mod `prime^n` in basis `gen`.

Where `prime` si a prime number dividing `card` and `n` is the largest integer
such that `prime^n` divides `card`. See reference **[2]** for more information about
this algorithm, the references can be found in the documentation of the module.
"""
function pohligHellmanPrime{T <: PolyElem}(prime::Int, n::Int, k::Int,
                                           mid::T, elem::T, defPol::T)

    f = div(k, prime^n)
    g = powmod(mid, f, defPol)
    inverse = gcdinv(g, defPol)[2]
    tmp = powmod(elem, f, defPol)
    g = powmod(g, prime^(n-1), defPol)

    # We compute a table of the `prime`-th roots of unit
    arr = Array(T, prime)
    for i in 0:(prime-1)
        arr[i + 1] = powmod(g, i, defPol)
    end

    res = 0

    # We compute the coefficients `b_i` such that x = Σ b_i × prime^i, meaning
    # that we compute `x` in basis `prime`
    d = prime^n
    for i in 0:(n-1)
        d = div(d, prime)
        b = findfirst(arr, powmod(tmp, d, defPol)) - 1
        res += b*prime^i
        tmp = mulmod(tmp, powmod(inverse, b*prime^i, defPol), defPol)
    end

    return (fmpz(res), fmpz(prime)^n)
end

"""
    pohligHellman{T}(card::Integer, gen::T, elem::T, defPol::T)

Compute the discrete logarithm of `elem` in the basis `gen`, modulo the small
prime factors of card - 1.
"""
function pohligHellman{T}(card::Integer, gen::T, elem::T, defPol::T)

    # We find the small (meaning, less that log(card)) prime factors of card
    q::BigInt = characteristic(base_ring(parent(gen)))
    l::Int = ceil(Integer, (q^2-1)/2)
    l = max(l, ceil(Integer, log2(card)))
    A = Array{Int, 1}()
    for i in primes(l)
        if card%i == 1
            push!(A, i)
        end
    end

    B = Array(Int, (2, length(A)))

    for i in 1:length(A)
        prime = A[i]
        B[1, i] = A[i]
        ex = 1
        while divisible(fmpz(card-1), prime^(ex+1))
            ex += 1
        end
        B[2, i] = ex
    end

    k = prod(Int[B[1, i]^B[2, i] for i in 1:length(A)])
    d = div(card-1, k)
    mid = powmod(gen, d, defPol)
    elemid = powmod(elem, d, defPol)

    # We set some variables to fmpz to use `crt`
    res::Nemo.fmpz = 0
    n::Nemo.fmpz = 0

    # And we compute the discrete logarithm of `elem` in the basis `gen`, for
    # each prime factor, using pohligHellmanPrime, then we compute our final
    # result using Chinese remindering theorem
    if length(A) > 1
        a = pohligHellmanPrime(B[1, 1], B[2, 1], k, mid, elemid, defPol)
        b = pohligHellmanPrime(B[1, 2], B[2, 2], k, mid, elemid, defPol)
        res, n = crt(a[1], a[2], b[1], b[2]), a[2]*b[2]
        for i in 1:(length(A)-2)
            a = pohligHellmanPrime(B[1, i+2], B[2, i+2], k, mid, elemid, defPol)
            res, n = crt(res, n, a[1], a[2]), n*a[2]
        end
        return res, n
    else
        return pohligHellmanPrime(B[1, 1], B[2, 1], k, mid, elemid, defPol)
    end
end

pohligHellman{T}(card::Nemo.fmpz, gen::T, elem::T, defPol::T) =
pohligHellman(BigInt(card), gen, elem, defPol)

# Functions precomputing parameters and roots of unity for the Pohlig Hellman
# algorithm

"""
    pohligHellmanParam{T}(card::Integer, gen::T, defPol::T)

Compute some necessary parameters for the Pohlig Hellman algorithm.
"""
function pohligHellmanParam{T}(card::Integer, gen::T, defPol::T)

    # We find the small (meaning, less that log(card)) prime factors of card-1
    q::BigInt = characteristic(base_ring(parent(gen)))
    l::Int = ceil(Integer, (q^2-1)/2)
    l = max(l, ceil(Integer, log2(card)))
    A = Array{Int, 1}()
    for i in primes(l)
        if card%i == 1
            push!(A, i)
        end
    end

    # We compute the exponent of these factors
    B = Array(Int, (2, length(A)))

    for i in 1:length(A)
        prime = A[i]
        B[1, i] = A[i]
        ex = 1
        while divisible(fmpz(card-1), prime^(ex+1))
            ex += 1
        end
        B[2, i] = ex
    end

    # We compute the product of the small factors and compute a big powering, in
    # order to do smaller powering later
    k = prod(Int[B[1, i]^B[2, i] for i in 1:length(A)])
    d = div(card-1, k)
    mid = powmod(gen, d, defPol)

    return k, d, mid, B
end

pohligHellmanParam{T}(card::Nemo.fmpz, gen::T, defPol::T) =
pohligHellmanParam(BigInt(card), gen, defPol)

"""
    pohligHellmanRoots{T <: PolyElem}(B::Array{Int, 2}, k::Int,
                                      mid::T, defPol::T)

Compute the roots of unity and inverses needed in the Pohlig-Hellman algorithm.
"""
function pohligHellmanRoots{T <: PolyElem}(B::Array{Int, 2}, k::Int,
                                           mid::T, defPol::T)

    # We setup a one dimensional array and a two dimensional array
    nbprimes = size(B)[2]
    inverses = Array(T, nbprimes)
    biggestPrime = B[1, nbprimes]
    arr = Array(T, (biggestPrime, nbprimes)) # We create an array that is too big !

    # We fill the one dimensional array with inverses of the generator and the
    # two dimensional array with roots of unity. Both are need for the
    # Pohlig-Hellman algorithm.
    for j in 1:nbprimes
        prime, n = B[1, j], B[2, j]
        f = div(k, prime^n)
        g = powmod(mid, f, defPol)
        inverses[j] = gcdinv(g, defPol)[2]
        g = powmod(g, prime^(n-1), defPol)

        # We compute a table of the `prime`-th roots of unit
        for i in 0:(prime-1)
            arr[i + 1, j] = powmod(g, i, defPol)
        end
    end

    return arr, inverses
end

"""
pohligHellmanPrimePrec{T <: PolyElem}(prime::Int, n::Int, k::Int,
                                      mid::T, elem::T, defPol::T,
                                      arr::Array{T, 1}, inverse::T)

Compute the discrete logarithm of `elem` mod `prime^n` in basis `gen` with
precomputed informations.

Where `prime` si a prime number dividing `card` and `n` is the largest integer
such that `prime^n` divides `card`. See reference **[2]** for more information about
this algorithm, the references can be found in the documentation of the module.

"""
function pohligHellmanPrimePrec{T <: PolyElem}(prime::Int, n::Int, k::Int,
                                           mid::T, elem::T, defPol::T,
                                           arr::Array{T, 1}, inverse::T)

    f = div(k, prime^n)
    tmp = powmod(elem, f, defPol)

    res = 0

    # We compute the coefficients `b_i` such that x = Σ b_i × prime^i, meaning
    # that we compute `x` in basis `prime`
    d = prime^n
    for i in 0:(n-1)
        d = div(d, prime)
        b = findfirst(arr, powmod(tmp, d, defPol)) - 1
        res += b*prime^i
        tmp = mulmod(tmp, powmod(inverse, b*prime^i, defPol), defPol)
    end

    return (fmpz(res), fmpz(prime)^n)
end

"""
    pohligHellmanPrec{T}(card::Integer, mid::T, elem::T, defPol::T, k::Int,
                         d::Integer, B::Array{Int, 2}, arr::Array{T, 2},
                         inverses::Array{T, 1})

Compute the discrete logarithm of `elem` in the basis `gen`, modulo the small
prime factors of card - 1, with some precomputed informations.
"""
function pohligHellmanPrec{T}(card::Integer, mid::T, elem::T, defPol::T, k::Int,
                              d::Integer, B::Array{Int, 2}, arr::Array{T, 2},
                              inverses::Array{T, 1})

    elemid = powmod(elem, d, defPol)

    # We set some variables to fmpz to use `crt`
    res::Nemo.fmpz = 0
    n::Nemo.fmpz = 0

    # And we compute the discrete logarithm of `elem` in the basis `gen`, for
    # each prime factor, using pohligHellmanPrime, then we compute our final
    # result using Chinese remindering theorem
    if size(B)[2] > 1
        a = pohligHellmanPrimePrec(B[1, 1], B[2, 1], k, mid, elemid, defPol,
                                   arr[1:B[1, 1], 1], inverses[1])
        b = pohligHellmanPrimePrec(B[1, 2], B[2, 2], k, mid, elemid, defPol,
                                   arr[1:B[1, 2], 2], inverses[2])
        res, n = crt(a[1], a[2], b[1], b[2]), a[2]*b[2]
        for i in 1:(size(B)[2]-2)
            a = pohligHellmanPrimePrec(B[1, i+2], B[2, i+2], k, mid, elemid,
                                       defPol, arr[1:B[1, i+2], i+2], inverses[i+2])
            res, n = crt(res, n, a[1], a[2]), n*a[2]
        end
        return res, n
    else
        return pohligHellmanPrimePrec(B[1, 1], B[2, 1], k, mid, elemid, defPol,
                                      arr[1:B[1, 1], 1], inverses[1])
    end
end

pohligHellmanPrec{T}(card::Nemo.fmpz, gen::T, elem::T, defPol::T, k::Int,
                     d::Integer, B::Array{Int, 2}, arr::Array{T, 2},
                     inverses::Array{T, 1}) = pohligHellmanPrec(BigInt(card), gen, elem,
                                                  defPol, k, d, B, arr, inverses)


text:
"""
This file initializes the ClimateERA module by setting and determining the
ECMWF reanalysis parameters to be analyzed and the regions upon which the data
are to be extracted from.  Functionalities include:
    - Setting up of reanalysis module type
    - Setting up of reanalysis parameters to be analyzed
    - Setting up of time steps upon which data are to be downloaded
    - Setting up of region of analysis based on ClimateEasy

"""

function gnssancread(fnc::AbstractArray,epar::Dict,nhr::Integer,ii::Integer,yr::Integer)

    nhr = daysinyear(yr-1) * nhr;

    dyr1 = erancread(fnc[ii],epar); dim = size(dyr1)

    if !(ii==1); d2 = erancread(fnc[ii-1],epar);
    else;        d2 = ones(dim[1],dim[2],nhr) * NaN;
    end

    return cat(d2,d1,dims=3);

end

function gnssancsave(dysm::AbstractArray,ddhr::AbstractArray,dvar::AbstractArray,
                    fnc::AbstractString,nlon::Integer,nlat::Integer,nhr::Integer)

    fnc = replace(fnc,"era"=>"gnssa");
    if isfile(fnc)
       @info "$(Dates.now()) - Unfinished netCDF file $(fnc) detected.  Deleting."
       rm(fnc);
    end

    nccreate(fnc,"longitude")

end


text:
using NewtonsMethod
using Test

@testset "NewtonsMethod.jl" begin
    # Write your own tests here.

# return back nothing if non-convergence: wrote in function

# test of non-convergence

f(x)=2+x^2
f_prime(x)=2*x
t1=newtonroot(f, f_prime; x_0=0, tolerance=1E-7, maxiter=1000)[1]
t2=newtonroot1(f; x_0=0, tolerance=1E-7, maxiter=1000)[1]
@test t1 == nothing
@test t2 == nothing

# root of a known function
f(x)=x^2-2x+1
f_prime(x)=2*x-2
t3=newtonroot(f, f_prime; x_0=0, tolerance=1E-7, maxiter=1000)[1]
@test t3 ≈ 1.0

# automatic differentiation
t4=newtonroot2(f; x_0=0, tolerance=1E-7, maxiter=1000)[1]
@test t4 ≈ 1.0

# finding those roots with a BigFloat
t5=newtonroot(f, f_prime; x_0=BigFloat(x0), tolerance=1E-7, maxiter=1000)[1]
@test t5 ≈ 1.0

# several of above tests

f(x)=x^2-2x
f_prime(x)=2*x-2

t3=newtonroot(f, f_prime; x_0=1.5, tolerance=1E-7, maxiter=1000)[1]
@test t3 ≈ 2.0
t4=newtonroot2(f; x_0=1.5, tolerance=1E-7, maxiter=1000)[1]
@test t4 ≈ 2.0
t5=newtonroot(f, f_prime; x_0=BigFloat(1.5), tolerance=1E-7, maxiter=1000)[1]
@test t5 ≈ 2.0

f(x)=x^2-3x+2
f_prime(x)=2*x-3

t3=newtonroot(f, f_prime; x_0=0, tolerance=1E-7, maxiter=1000)[1]
@test t3 ≈ 1.0
t4=newtonroot2(f; x_0=0, tolerance=1E-7, maxiter=1000)[1]
@test t4 ≈ 1.0
t5=newtonroot(f, f_prime; x_0=BigFloat(0), tolerance=1E-7, maxiter=1000)[1]
@test t5 ≈ 1.0

# ensure that the maxiter is working
f(x)=x^2-2x+1
f_prime(x)=2*x-2

t6=newtonroot(f, f_prime; x_0=0, tolerance=1E-7, maxiter=5)[1]
@test t6 == nothing

# ensure that tol is working
t7=newtonroot(f, f_prime; x_0=0, tolerance=1E-7, maxiter=1000)[2]
@test t7 <= 1E-7

end


text:
#! /usr/bin/env julia

include("/home/steve/Software/julia/base/test.jl")
using Test

f() = 1
g() = 2

println(macroexpand(:(@test f() <= g() <=6 <=7)))
println(macroexpand(:(@test f() >= g() >= 6)))

@test f() <= g() <=6 <=7
@test f() >= g() >= 6

max_stars_repo_name: sjkelly/JuliaPlayground

max_issues_repo_name: sjkelly/JuliaPlayground

max_forks_repo_name: sjkelly/JuliaPlayground


text:
# Automatic lane change system

# ## Model

module automatic_lane_change_system #jl
using ReachabilityAnalysis, ModelingToolkit

vars = @variables vx, sx, ax, w, vy, sy

@taylorize function mode0!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.01 * sx - 0.103 + 0.84 - 0.3 * vx - 0.5 * ax
    du[4] = -2 * w
    du[5] = -2 * vy
    du[6] = 0.1 * vy
    return du
end

@taylorize function mode1!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.5 *1.4 - 0.5 * ax
    du[4] = 3 - 0.15 * w + 0.2 - 0.01 * sy
    du[5] = 7.5 - 0.45 * w + 0.5 - 0.025 * sy - 0.05 * vy
    du[6] = 0.1 * vy
    return du
end

@taylorize function mode2!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.5 * vx + 1.4 - 0.5 * ax
    du[4] = 0.15 * w + 0.2 - 0.01 * sy
    du[5] = 0.25 * w + 0.5 - 0.025 * sy - 0.05 * vy
    du[6] = 0.1 * vy
    return du
end

@taylorize function mode3!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.5 * vx + 1.4 - 0.5 * ax
    du[4] = -0.1 * w + 0.2 - 0.01 * sy
    du[5] = -0.25 * w + 0.5 - 0.025 * sy - 0.05 * vy
    du[6] = 0.1 * vy
    return du
end

@taylorize function mode4!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.5 * vx + 1.4 - 0.5 * ax
    du[4] = -3 - 0.15 * w + 0.2 - 0.01 * sy
    du[5] = -7.5 - 0.45 * w + 0.5 - 0.025 * sy - 0.05 * vy
    du[6] = 0.1 * vy
    return du
end

@taylorize function mode5!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.01 * sx - 0.103 + 0.84 - 0.3 * vx - 0.5 * ax
    du[4] = -2 * w
    du[5] = -2 * vy
    du[6] = 0.1 * vy
    return du
end

@taylorize function mode6!(du, u, params, t)
    vx, sx, ax, w, vy, sy = u
    du[1] = 0.1 * ax
    du[2] = vx - 2.5
    du[3] = -0.5 * vx + 1.4 - 0.5 * ax
    du[4] = -2 * w
    du[5] = -2 * vy
    du[6] = 0.1 * vy
    return du
end

function acc_model()
    HA = LightAutomaton(7)

    ## mode 0
    X = HalfSpace(sx + 10 < 0, vars)
    m0 = @system(x' = mode0!(x), dim: 6, x ∈ X)

    ## mode 1
    X = HalfSpace(sy < 12, vars)
    m1 = @system(x' = mode1!(x), dim: 6, x ∈ X)

    ## mode 2
    X = HalfSpace(vy > 0.05, vars)
    m2 = @system(x' = mode2!(x), dim: 6, x ∈ X)

    ## mode 3
    X = HalfSpace(vy + 0.05 < 0, vars)
    m3 = @system(x' = mode3!(x), dim: 6, x ∈ X)

    ## mode 4
    X = HalfSpace(sy > 3.5, vars)
    m4 = @system(x' = mode4!(x), dim: 6, x ∈ X)

    ## mode 5
    X = Universe(6)
    m5 = @system(x' = mode5!(x), dim: 6, x ∈ X)

    ## mode 6
    X = HalfSpace(vy < 0.05, vars)
    m6 = @system(x' = mode6!(x), dim: 6, x ∈ X)

    ## modo 0 → 1
    add_transition!(HA, 1, 2, 1)
    G = HalfSpace(sx + 10 >= 0, vars)
    R0 = ConstrainedIdentityMap(6, G)

    ## modo 1 → 2
    add_transition!(HA, 2, 3, 2)
    G = HalfSpace(sy >= 12, vars)
    R1 = ConstrainedIdentityMap(6, G)

    ## modo 4 → 3
    add_transition!(HA, 5, 4, 5)
    G = HalfSpace(sy <= 3.5, vars)
    R2 = ConstrainedIdentityMap(6, G)

    ## modo 5 → 4
    add_transition!(HA, 6, 5, 6)
    G = HalfSpace(sx >= 10, vars)
    R3 = ConstrainedIdentityMap(6, G)

    ## modo 2 → 5
    add_transition!(HA, 3, 6, 3)
    G = HalfSpace(sx >= 10, vars)
    R4 = ConstrainedIdentityMap(6, G)

    ## modo 3 → 6
    add_transition!(HA, 4, 7, 4)
    G = HalfSpace(vy + 0.07 >= 0, vars)
    R5 = ConstrainedIdentityMap(6, G)

    modes = [m1, m2, m3, m4, m5, m6]

    resets = [R0, R1, R2, R3, R4, R5]

    ## hybrid system
    S = HybridSystem(HA, modes, resets, fill(AutonomousSwitching(), 6))

    return S
end

function model(X0)
    H = acc_model()
    return IVP(H, X0)
end
end #jl

# ## References

max_stars_repo_name: mforets/ReachabilityLibrary

max_issues_repo_name: mforets/ReachabilityLibrary

max_forks_repo_name: mforets/ReachabilityLibrary


text:
using Random
using StableRNGs
using HDF5

import Base: iterate, length, IteratorEltype

export dataset
export SARTS
export SART
export DataSet

const SARTS = (:state, :action, :reward, :terminal, :next_state)
const SART = (:state, :action, :reward, :terminal)

"""
Represents a iterable dataset with the following fields:

`dataset`: Dict{Symbol, Any}, representation of the dataset as a Dictionary with style as `style`
`repo`: String, the repository from which the dataset is taken
`size`: Integer, the size of the dataset
`batch_size`: Integer, the size of the batches returned by `iterate`.
`style`: Tuple, the type of the NamedTuple, for now SARTS and SART is supported.
`rng`<: AbstractRNG.
`meta`: Dict, the metadata provided along with the dataset
`is_shuffle`: Bool, determines if the batches returned by `iterate` are shuffled.
"""
struct DataSet{T<:AbstractRNG}
    dataset::Dict{Symbol, Any}
    repo::String
    size::Integer
    batch_size::Integer
    style::Tuple
    rng::T
    meta::Dict
    is_shuffle::Bool
end

# TO-DO: include other functionality like is_sequential, will be implemented soon
# TO-DO: enable the users providing their own paths to datasets if they already have it
# TO-DO: add additional env arg to do complete verify function
"""
    dataset(dataset::String; style::Tuple, rng<:AbstractRNG, is_shuffle::Bool, max_iters::Int64, batch_size::Int64)

Creates a dataset of enclosed in a DataSet type and other related metadata for the `dataset` that is passed.
The `DataSet` type is an iterable that fetches batches when used in a for loop for convenience during offline training.

`dataset`: Dict{Symbol, Any}, Name of the datset.
`repo`: Name of the repository of the dataset.
`style`: the style of the iterator and the Dict inside DataSet that is returned.
`rng`: StableRNG
`max_iters`: maximum number of iterations for the iterator.
`is_shuffle`: whether the dataset is shuffled or not. `true` by default.
`batch_size`: batch_size that is yielded by the iterator. Defaults to 256.

The returned type is an infinite iterator which can be called using `iterate` and will return batches as specified in the dataset.
"""
function dataset(dataset::String;
    style=SARTS,
    repo = "d4rl",
    rng = StableRNG(123),
    is_shuffle = true,
    batch_size=256
)

    try
        @datadep_str repo*"-"*dataset
    catch
        throw("The provided dataset is not available")
    end

    path = @datadep_str repo*"-"*dataset

    @assert length(readdir(path)) == 1
    file_name = readdir(path)[1]

    data = h5open(path*"/"*file_name, "r") do file
        read(file)
    end

    # sanity checks on data
    verify(data)

    dataset = Dict{Symbol, Any}()
    meta = Dict{String, Any}()

    N_samples = size(data["observations"])[2]

    for (key, d_key) in zip(["observations", "actions", "rewards", "terminals"], Symbol.(["state", "action", "reward", "terminal"]))
            dataset[d_key] = data[key]
    end

    for key in keys(data)
        if !(key in ["observations", "actions", "rewards", "terminals"])
            meta[key] = data[key]
        end
    end

    return DataSet(dataset, repo, N_samples, batch_size, style, rng, meta, is_shuffle)

end

function iterate(ds::DataSet, state = 0)
    rng = ds.rng
    batch_size = ds.batch_size
    size = ds.size
    is_shuffle = ds.is_shuffle
    style = ds.style

    if is_shuffle
        inds = rand(rng, 1:size, batch_size)
        map((x)-> if x <= size x else 1 end, inds)
    else
        if (state+1) * batch_size <= size
            inds = state*batch_size+1:(state+1)*batch_size
        else
            return nothing
        end
        state += 1
    end

    batch = (state = copy(ds.dataset[:state][:, inds]),
    action = copy(ds.dataset[:action][:, inds]),
    reward = copy(ds.dataset[:reward][inds]),
    terminal = copy(ds.dataset[:terminal][inds]))

    if style == SARTS
        batch = merge(batch, (next_state = copy(ds.dataset[:state][:, (1).+(inds)]),))
    end

    return batch, state
end

take(ds::DataSet, n::Integer) = take(ds.dataset, n)
length(ds::DataSet) = ds.size
IteratorEltype(::Type{DataSet}) = EltypeUnknown() # see if eltype can be known (not sure about carla and adroit)

function verify(data::Dict{String, Any})
    for key in ["observations", "actions", "rewards", "terminals"]
        @assert (key in keys(data)) "Expected keys not present in data"
    end
    N_samples = size(data["observations"])[2]
    @assert size(data["rewards"]) == (N_samples,) || size(data["rewards"]) == (1, N_samples)
    @assert size(data["terminals"]) == (N_samples,) || size(data["terminals"]) == (1, N_samples)
end


text:
global PROC

push!(LOAD_PATH,pwd())
#push!(LOAD_PATH, joinpath(pwd(), "src"))

#push all subdirectories from src
map(d -> push!(LOAD_PATH, joinpath(pwd(), "./", d)),
    filter(d -> isdir(joinpath("./", d)), readdir("./")))

#push!(LOAD_PATH, joinpath(pwd(), "Demos"))

#include("usegpu.jl")
include("useproc.jl")

max_stars_repo_name: davidbarber/AutoDiff

max_issues_repo_name: davidbarber/AutoDiff

max_forks_repo_name: davidbarber/AutoDiff


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

module Meshes

using Tables
using StaticArrays
using SparseArrays
using CircularArrays
using SimpleTraits
using RecipesBase
using LinearAlgebra
using Random

using IterTools: ivec
using StatsBase: Weights
using SpecialFunctions: gamma
using Distances: PreMetric, Euclidean, Mahalanobis, evaluate
using ReferenceFrameRotations: angle_to_dcm
using NearestNeighbors: KDTree, BallTree, knn, inrange

# import categorical arrays as a temporary solution for plot recipes
using CategoricalArrays: CategoricalValue, levelcode

import Tables
import Random
import Base: values, ==, +, -, *
import StatsBase: sample
import Distances: evaluate
import NearestNeighbors: MinkowskiMetric

# Queryverse compatibility
import TableTraits
import IteratorInterfaceExtensions
const IIE = IteratorInterfaceExtensions

# IO utils
include("ioutils.jl")

# numerical tolerances
include("tolerances.jl")

# basic types
include("vectors.jl")
include("points.jl")
include("angles.jl")

# type traits
include("traits.jl")

# geometries
include("geometries.jl")

# collections
include("collections.jl")

# meshes
include("connectivities.jl")
include("topologies.jl")
include("toporelations.jl")
include("mesh.jl")

# mesh data
include("meshdata.jl")

# miscellaneous
include("conventions.jl")
include("paths.jl")
include("distances.jl")
include("neighborhoods.jl")
include("neighborsearch.jl")
include("supportfun.jl")
include("laplacian.jl")

# views and partitions
include("views.jl")
include("partitions.jl")

# algorithms
include("viewing.jl")
include("sampling.jl")
include("partitioning.jl")
include("intersections.jl")
include("discretization.jl")
include("simplification.jl")
include("refinement.jl")
include("smoothing.jl")
include("boundingboxes.jl")
include("hulls.jl")

# utilities
include("utils.jl")

# plot recipes
include("plotrecipes/domain.jl")
include("plotrecipes/data.jl")
include("plotrecipes/points.jl")
include("plotrecipes/geometries.jl")
include("plotrecipes/collections.jl")
include("plotrecipes/cartesiangrids.jl")
include("plotrecipes/partitions.jl")

export
  # points
  Point, Point1, Point2, Point3,
  Point1f, Point2f, Point3f,
  embeddim, paramdim,
  coordtype, coordinates,
  center, centroid,
  ⪯, ≺, ⪰, ≻,

  # vectors
  Vec, Vec1, Vec2, Vec3, Vec1f, Vec2f, Vec3f,

  # angles
  ∠,

  # domain traits
  Domain,
  embeddim, paramdim, coordtype,
  element, nelements,

  # data traits
  Data, Variable,
  domain, constructor, asarray,
  variables, name, mactype,

  # optional traits
  IsGrid, isgrid,

  # domain/data alias
  DomainOrData,

  # geometries
  Geometry,
  embeddim, paramdim, coordtype,
  measure, area, volume, boundary,
  center, centroid, isconvex, issimplex,

  # primitives
  Primitive,
  Line, Ray, Plane, BezierCurve,
  Box, Ball, Sphere, Cylinder,
  ncontrols, degree, Horner, DeCasteljau,
  radius, height, sides,
  measure, diagonal,

  # polytopes
  Polytope, Polygon, Polyhedron,
  Segment, Ngon, Triangle, Quadrangle,
  Pentagon, Hexagon, Heptagon,
  Octagon, Nonagon, Decagon,
  Chain, PolyArea,
  Tetrahedron, Pyramid, Hexahedron,
  vertices, nvertices,
  windingnumber, chains, segments,
  isclosed, issimple, hasholes,
  angles, innerangles, close!, open!,
  orientation, bridge, normal,

  # orientation algorithms
  WindingOrientation,
  TriangleOrientation,

  # point or geometry alias
  PointOrGeometry,

  # multi-types
  Multi,

  # collections
  Collection,
  PointSet,
  GeometrySet,

  # rotation conventions
  RotationConvention,
  GSLIB, Leapfrog, Datamine,
  TaitBryanExtr, TaitBryanIntr,
  EulerExtr, EulerIntr,
  axesseq, orientation, angleunits,
  mainaxis, isextrinsic, rotmat,

  # paths
  Path,
  LinearPath, RandomPath,
  ShiftedPath, SourcePath,
  traverse,

  # neighborhoods
  Neighborhood, MetricBall,
  NormBall, Ellipsoid,
  metric, radius,

  # neighbordhood search
  NeighborSearchMethod,
  BoundedNeighborSearchMethod,
  BallSearch,
  KNearestSearch,
  KBallSearch,
  BoundedSearch,
  search!, search,

  # miscellaneous
  supportfun,
  laplacematrix,

  # connectivities
  Connectivity,
  paramdim, indices,
  connect, materialize,
  issimplex,

  # topologies
  Topology,
  vertices, nvertices,
  faces, elements, facets,
  element, nelements,
  facet, nfacets,
  FullTopology,
  GridTopology,
  HalfEdgeTopology, HalfEdge,
  half4elem, half4vert,
  half4edge, half4pair,
  edge4pair,

  # topological relations
  TopologicalRelation,
  Boundary, Coboundary, Adjacency,

  # meshes
  Mesh,
  CartesianGrid, SimpleMesh,
  vertices, nvertices, topology,
  faces, elements, facets,
  element, nelements,
  facet, nfacets,
  topoconvert,
  spacing,

  # mesh data
  MeshData,
  meshdata,

  # views
  DomainView, DataView,

  # partitions
  Partition,
  indices, metadata,

  # viewing
  indices, slice,

  # sampling
  SamplingMethod,
  DiscreteSamplingMethod,
  ContinuousSamplingMethod,
  UniformSampling,
  WeightedSampling,
  BallSampling,
  RegularSampling,
  HomogeneousSampling,
  MinDistanceSampling,
  sample,

  # partitioning
  PartitionMethod,
  PredicatePartitionMethod,
  SPredicatePartitionMethod,
  RandomPartition,
  FractionPartition,
  BlockPartition,
  BisectPointPartition,
  BisectFractionPartition,
  BallPartition,
  PlanePartition,
  DirectionPartition,
  PredicatePartition,
  SpatialPredicatePartition,
  ProductPartition,
  HierarchicalPartition,
  partition, split, →,

  # intersections
  Intersection,
  NoIntersection,
  CrossingLines,
  OverlappingLines,
  CrossingSegments,
  MidTouchingSegments,
  CornerTouchingSegments,
  OverlappingSegments,
  OverlappingBoxes,
  FaceTouchingBoxes,
  CornerTouchingBoxes,
  intersecttype,
  hasintersect,

  # discretization
  DiscretizationMethod,
  FIST, Dehn1899,
  discretize,
  triangulate,

  # simplification
  SimplificationMethod,
  DouglasPeucker,
  simplify,
  decimate,

  # refinement
  RefinementMethod,
  QuadRefinement,
  CatmullClark,
  refine,

  # smoothing
  SmoothingMethod,
  TaubinSmoothing,
  smooth,

  # bounding boxes
  boundingbox,

  # hulls
  HullMethod,
  GrahamScan,
  hull,

  # utililities
  signarea,
  sideof

end # module


text:
# ---
# title: Time series
# cover: assets/time_series.png
# description: Visualizing time series data.
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
using Dates
set_aog_theme!() #src

x = today() - Year(1) : Day(1) : today()
y = cumsum(randn(length(x)))
z = cumsum(randn(length(x)))
df = (; x, y, z)
labels = ["series 1", "series 2", "series 3", "series 4", "series 5"]
plt = data(df) *
    mapping(:x, [:y, :z] .=> "value", color=dims(1) => renamer(labels) => "series ") *
    visual(Lines)
draw(plt)

#

x = now() - Hour(6) : Minute(1) : now()
y = cumsum(randn(length(x)))
z = cumsum(randn(length(x)))
df = (; x, y, z)
plt = data(df) *
    mapping(:x, [:y, :z] .=> "value", color=dims(1) => renamer(labels) =>"series ") *
    visual(Lines)
fg = draw(plt)

#

dates = Date(2022, 1, 1):Day(1):Date(2022, 1, 31)
trend = cumsum(randn(31))
df = map(1:2000) do _
    idx = rand(1:31)
    date = dates[idx]
    observation = trend[idx] + 2 * rand()
    return (; date, observation)
end
plt = data(df) * mapping(:date, :observation) * visual(BoxPlot)
draw(plt)

# save cover image #src
mkpath("assets") #src
save("assets/time_series.png", fg) #src


text:
using SpringerQuarantineBooks
# Empty because we don't want to download the books for C.I.


text:
using Base.Test
using StatsBase.sample
using DataFrames

@time @testset "causal search" begin
    X = readdlm(joinpath(@__DIR__, "X1.dat"))
    env = Vector{Int}(X[:,1])
    X = X[:,2:end]
    S = 1:size(X,2)
    result = causalSearch(X, X[:, 1], env, setdiff(S,1) , α=0.01)
    @test result.model_reject == true
    result = causalSearch(X, X[:, 2], env, setdiff(S,2) , α=0.01, p_max=4)
    @test result.S == [5]
end

@time @testset "causal search logistic" begin
    df = DataFrame(x1 = CategoricalArray([1  0  0  0  0  1  0  1  1  0  0  0  0  1  0  1][:]),
                x2 = [0.0  0.0  1.0  0.0  0.0  0.0  1.0  1.0  1.0  0.0  1.0  0.0  0.0  0.0  1.0  1.0][:],
                x3 = [0.0  0.0  1.0  1.0  0.0  0.0  1.0  0.0  0.0  0.0  1.0  1.0  0.0  0.0  1.0  0.0][:],
                y =  [1.0  1.0  1.0  0.0  0.0  1.0  1.0  1.0  0.0  1.0  1.0  0.0  0.0  1.0  1.0  0.0][:])
    env = repeat([1,2], inner=8)
    r1 = causalSearch(df, :y, env, method="logistic-LR", iterate_all=true)
    @test length(r1.S) == 0
    r2 = causalSearch(df, :y, env, method="logistic-SF", iterate_all=true)
    @test length(r2.S) == 0
end


text:
import Diffusion  # Requires to do Pkg.activate(".") in the project's top level folder.

import MPI
using CUDA

MPI.Init()
comm = MPI.COMM_WORLD

# Numerics
nx     = 128;     # Number of grid points in x dimension
ny     = 128;     # Number of grid points in y dimension
nt     = 1000;    # Number of time steps

# Physics
lx     = 10.0;    # Length of domain in x dimension
ly     = 10.0;    # Length of domain in y dimension
lam    = 1.0;     # Thermal conductivity
cp_min = 1.0;     # Minimal heat capacity
amp_T  = 100.0;   # Amplitude of Gaussian anomly in temperature field

Cp = zeros(Float64, nx, ny) # Heat capacity field

for ix = 1:nx
    for iy = 1:ny
        Cp[ix,iy] = cp_min + iy/ny;
    end
end

# Copy field Cp to GPU
Cp_d = CuArray(Cp)

# Simulate heat diffusion calling directly the diffusion routine.
Diffusion.diffusion2D(nx, ny, nt, lx, ly, lam, cp_min, amp_T, Cp_d, comm)

# Simulate heat diffusion calling the C interface of the diffusion routine.
Diffusion.julia_diffusion(Csize_t(nx), Csize_t(ny), Csize_t(nt), Cdouble(lx), Cdouble(ly), Cdouble(lam), Cdouble(cp_min), Cdouble(amp_T), pointer(Cp_d), comm.val)

MPI.Finalize()

max_stars_repo_name: omlins/libdiffusion

max_issues_repo_name: omlins/libdiffusion

max_forks_repo_name: omlins/libdiffusion


text:
import Base.LinAlg.LU, Base.getindex
import AxisAlgorithms: A_ldiv_B_md!, _A_ldiv_B_md!

### Tridiagonal inversion along a particular dimension, first offsetting the values by b

function A_ldiv_B_md!(dest, F, src, dim::Integer, b::AbstractVector)
    1 <= dim <= max(ndims(dest),ndims(src)) || throw(DimensionMismatch("The chosen dimension $dim is larger than $(ndims(src)) and $(ndims(dest))"))
    n = size(F, 1)
    n == size(src, dim) && n == size(dest, dim) || throw(DimensionMismatch("Sizes $n, $(size(src,dim)), and $(size(dest,dim)) do not match"))
    size(dest) == size(src) || throw(DimensionMismatch("Sizes $(size(dest)), $(size(src)) do not match"))
    check_matrix(F)
    length(b) == size(src,dim) || throw(DimensionMismatch("length(b) = $(length(b)), which does not match $(size(src,dim))"))
    R1 = CartesianRange(size(dest)[1:dim-1])  # these are not type-stable, so let's use a function barrier
    R2 = CartesianRange(size(dest)[dim+1:end])
    _A_ldiv_B_md!(dest, F, src, R1, R2, b)
end

# Filtering along the first dimension
function _A_ldiv_B_md!{T,CI<:CartesianIndex{0}}(dest, F::LU{T,Tridiagonal{T}}, src,  R1::CartesianRange{CI}, R2, b)
    n = size(F, 1)
    if n == 0
        return nothing
    end
    dl = F.factors.dl
    d  = F.factors.d
    du = F.factors.du
    dinv = 1./F.factors.d  # might not want to do this for small R2
    @inbounds for I2 in R2
        # Forward substitution
        dest[1, I2] = src[1, I2] + b[1]
        for i = 2:n  # Can't use @simd here
            dest[i, I2] = src[i, I2] + b[i] - dl[i-1]*dest[i-1, I2]
        end
        # Backward substitution
        dest[n, I2] *= dinv[n]
        for i = n-1:-1:1   # Can't use @simd here
            dest[i, I2] = (dest[i, I2] - du[i]*dest[i+1, I2])*dinv[i]
        end
    end
    nothing
end

# Filtering along any other dimension
function _A_ldiv_B_md!{T}(dest, F::LU{T,Tridiagonal{T}}, src, R1, R2, b)
    n = size(F, 1)
    dl = F.factors.dl
    d  = F.factors.d
    du = F.factors.du
    dinv = 1./d  # might not want to do this for small R1 and R2
    # Forward substitution
    @inbounds for I2 in R2
        @simd for I1 in R1
            dest[I1, 1, I2] = src[I1, 1, I2] + b[1]
        end
        for i = 2:n
            @simd for I1 in R1
                dest[I1, i, I2] = src[I1, i, I2] + b[i] - dl[i-1]*dest[I1, i-1, I2]
            end
        end
        # Backward substitution
        @simd for I1 in R1
            dest[I1, n, I2] *= dinv[n]
        end
        for i = n-1:-1:1
            @simd for I1 in R1
                dest[I1, i, I2] = (dest[I1, i, I2] - du[i]*dest[I1, i+1, I2])*dinv[i]
            end
        end
    end
    dest
end

### Woodbury solver with offset

function _A_ldiv_B_md!(dest, W::Woodbury, src,  R1, R2, b)
    _A_ldiv_B_md!(dest, W.A, src, R1, R2, b)
    tmp1 = AxisAlgorithms._A_mul_B_md(W.V, dest, R1, R2)
    tmp2 = AxisAlgorithms._A_mul_B_md(W.Cp, tmp1, R1, R2)
    tmp3 = AxisAlgorithms._A_mul_B_md(W.U, tmp2, R1, R2)
    # TODO?: would be nice to fuse the next two steps
    AxisAlgorithms._A_ldiv_B_md!(tmp3, W.A, tmp3, R1, R2)
    AxisAlgorithms.sub!(dest, tmp3)
end

function check_matrix{T}(F::LU{T,Tridiagonal{T}})
    for i = 1:size(F,1)
        F.ipiv[i] == i || error("For efficiency, pivoting is not supported")
    end
    du2 = F.factors.du2
    for i = 1:length(du2)
        du2[i] == 0 || error("For efficiency, du2 must be all zeros")
    end
end

check_matrix(F::Woodbury) = check_matrix(F.A)

check_matrix(M) = error("unsupported matrix type $(typeof(M))")


text:
using AxisArrays

@testset "With AxisArray" begin
    a = AxisArray(reshape(1:6, 3, 2), Axis{:foo}(2:4), Axis{:bar}(3:4))
    # 2-dimensional AxisArray{Int64,2,...} with axes:
    #     :foo, 2:4
    #     :bar, 3:4
    # And data, a 3×2 Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}:
    #  1  4
    #  2  5
    #  3  6

    v = Axis{:virt}(5:8)

    w = @inferred(EM.WrappedArray(a, (v, AxisArrays.axes(a)...)))

    for i = (5, 10, :a)
        @test @inferred(w[i, :, :]) == a
    end

    w2 = EM.WrappedArray(a, (AxisArrays.axes(a, 2), v, AxisArrays.axes(a, 1)))
    @test isa(@inferred(w2[1, 4, 3]), Number)
    @test w2[1, 4, 3] == a[3, 1]
end

@testset "With constant" begin
    using AxisArrays

    b = 2.0
    v = Axis{:virt}(5:8)

    w3 = EM.WrappedArray(b, v)
    @test @inferred(w3[5]) == b
    @test @inferred(w3[8]) == b
end

max_stars_repo_name: coroa/EnergyModels

max_issues_repo_name: coroa/EnergyModels

max_forks_repo_name: coroa/EnergyModels


text:
type Color
    r::Uint8
    g::Uint8
    b::Uint8
end

type Image
    pic::Array{Color,2}
end
Image(w::Int, h::Int) = Image(Array(Color, w, h))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
using Surrogates
using LinearAlgebra
using Test
using QuadGK
using Cubature

#1D
obj = x -> 3*x + log(x)
a = 1.0
b = 4.0
x = sample(2000,a,b,SobolSample())
y = obj.(x)
alpha = 2.0
n = 6
my_loba = LobacheskySurrogate(x,y,alpha,n,a,b)
val = my_loba(3.83)

#1D integral
int_1D = lobachesky_integral(my_loba,a,b)
int = quadgk(obj,a,b)
int_val_true = int[1]-int[2]
@test abs(int_1D - int_val_true) < 10^-5
add_point!(my_loba,3.7,12.1)
add_point!(my_loba,[1.23,3.45],[5.20,109.67])

#ND

obj = x -> x[1] + log(x[2])
lb = [0.0,0.0]
ub = [8.0,8.0]
alpha = [2.4,2.4]
n = 8
x = sample(3200,lb,ub,SobolSample())
y = obj.(x)
my_loba_ND = LobacheskySurrogate(x,y,alpha,n,lb,ub)

#ND

int_ND = lobachesky_integral(my_loba_ND,lb,ub)
int = hcubature(obj,lb,ub)
int_val_true = int[1]-int[2]
@test abs(int_ND - int_val_true) < 10^-1
add_point!(my_loba_ND,(10.0,11.0),4.0)
add_point!(my_loba_ND,[(12.0, 15.0),(13.0,14.0)],[4.0,5.0])
lobachesky_integrate_dimension(my_loba_ND,lb,ub,2)

obj = x -> x[1] + log(x[2]) + exp(x[3])
lb = [0.0,0.0,0.0]
ub = [8.0,8.0,8.0]
alpha = [2.4,2.4,2.4]
x = sample(50,lb,ub,SobolSample())
y = obj.(x)
n = 4
my_loba_ND = LobacheskySurrogate(x,y,alpha,n,lb,ub)
lobachesky_integrate_dimension(my_loba_ND,lb,ub,2)


text:
# For now dummy runtests.jl
# (Turing currently not available in test environment)


text:
# basic Pose2 convolution test

using RoME
using Test

##

@testset "test basic Pose2 convolution based on MvNormal" begin
##

fg = initfg()

addVariable!(fg, :x0, Pose2)
u0 = ProductRepr([0;0.0], [1 0; 0 1.])
M = getManifold(Pose2) # TODO add better dispatch to simplify
x0 = [AMP.makePointFromCoords(M,0.01*randn(3),u0) for _ in 1:100];
X0 = manikde!(M, x0)
initManual!(fg, :x0, X0)

addVariable!(fg, :x1, Pose2)
addFactor!(fg, [:x0;:x1], Pose2Pose2(MvNormal([10;0;pi], 0.1*diagm([1;1;1]))), graphinit=false)

## now do a convolution

X1_ = approxConvBelief(fg, :x0x1f1, :x1)
x1_ = getPoints(X1_) .|> x->AMP.makeCoordsFromPoint(M,x)

μ_x1_ = mean(X1_)
Σ_x1_ =  cov(M, getPoints(X1_))

@test isapprox( μ_x1_.parts[1], [10; 0], atol=0.2)
@test isapprox( μ_x1_.parts[2], [-1 0; 0 -1], atol=0.2)

@test isapprox( Σ_x1_, [0.15 0 0; 0 0.15 0; 0 0 0.15], atol=0.4)

# check X1_ rotation's coordinates is spread between +pi and -pi

ψ_ = (x->x[3]).(x1_)

# bi-modality should occur since ±pi is the same value on-manifold
@test 20 < sum(2 .< ψ_)
@test 20 < sum(ψ_ .< -2)
# numerical check
@test sum(-2 .< ψ_ .< 2) < 1
# on-manifold check
@test sum(3.15 .< abs.(ψ_)) < 1

## now test the deconvolution

pts, meas = approxDeconv(fg, :x0x1f1)
X12_ = manikde!(M, pts)

# check that deconv is good

@test mmd(M, pts, meas) < 0.001

##
end


text:
# Example showing the use of waveform triggering
using AcqSynth, PlotlyJS
const acq = AcqSynth

# As in the first example, we setup the board
boardnum = 0
acq.init_board(boardnum)

# Let's change to single channel mode
chan = acq.AIN0
acq.set_channels(boardnum,chan)
# We can verify that this worked properly
acq.get_channels(boardnum) == chan

# Check the frequency, in single channel mode, this will be twice the normal
acq.get_frequency(boardnum)

# Configure the board for waveform triggering
hysteresis = 32
threshold = 2^11-220 # middle value for a 12-bit board
ttype = acq.WAVEFORM_TRIGGER
# This function configures the waveform parameters
acq.set_waveform_trigger_params(boardnum,threshold,hysteresis)
# This function sets the trigger type
acq.set_trigger(boardnum,ttype,acq.RISING_EDGE,chan)

# Setup acquisition, the card will wait until a trigger is detected
numblocks = 2 # let's grab 2 blocks
acq.setup_acquire(boardnum,numblocks)

# Grab the data from the card
data = acq.get_samples_12(boardnum,numblocks) # transfer data as 12-bit samples

# Plot the stuff
t = (1:length(data))./4E9 # time vector, we are sampling at 4GSPS
plot(scatter(x=t[1:2000],y=data[1:2000],name="AIN0"))


text:
export ARC_step_computation

function ARC_step_computation(h, g, Δ; kwargs...)

    h2     = ((h) .^ 2)
    gdelta = (4 .* (g ./ Δ))
    discr1 = h2 .- gdelta
    discr2 = h2 .+ gdelta

    if true in (h .> .0)
      if true in (g .> .0)
        d = (-2 .* g) ./ (h .+ sqrt.(discr2))
      else
        d = (-2 .* g) ./ (h .+ sqrt.(discr1))
      end
    else
      if true in (g .< .0)
        d = (-h .+ sqrt.(discr1)) ./ (2 ./ Δ)
      else
        d = (-h .+ sqrt.(discr2)) ./ (-2 ./ Δ)
      end
    end

    return vec(d)
  end


text:
#####
##### Boundary conditions on pressure are derived from boundary conditions
##### on the north-south horizontal velocity, v.
#####

# Pressure boundary conditions are either zero flux (Neumann) or Periodic.
PressureBC(::BC)  = BoundaryCondition(Flux, nothing)
PressureBC(::PBC) = PeriodicBC()

function PressureBoundaryConditions(vbcs)
    x = CoordinateBoundaryConditions(PressureBC(vbcs.x.left), PressureBC(vbcs.x.right))
    y = CoordinateBoundaryConditions(PressureBC(vbcs.y.left), PressureBC(vbcs.y.right))
    z = CoordinateBoundaryConditions(PressureBC(vbcs.z.left), PressureBC(vbcs.z.right))
    return (x=x, y=y, z=z)
end


text:
function ∑tendencies!(q_tendencies, q, ∑tendencies_params, t)
    @unpack params, grid, aux_O2, aux, case, tri_diag = ∑tendencies_params

    assign!(q_tendencies, (:u, :v, :q_tot, :θ_liq), grid, 0.0)

    gm, en, ud, sd, al = allcombinations(q)

    for k in over_elems(grid)
        for (tend, force) in (
            (:u, :u_forcing),
            (:v, :v_forcing),
            (:θ_liq, :θ_liq_forcing),
            (:q_tot, :q_tot_forcing),
        )
            q_tendencies[tend, k] = aux[force, k]
        end
    end

    compute_tendencies_en_O2!(
        grid,
        q_tendencies,
        q,
        aux_O2,
        :tke,
        aux,
        tri_diag,
        params,
    )
    compute_tendencies_gm_scalars!(grid, q_tendencies, q, aux, params, tri_diag)
    compute_tendencies_ud!(grid, q_tendencies, q, aux, params)
end


text:
"""
    InstantiateEvaluateCovariance(cℓ::AbstractCℓ, ConvDens::AbstractConvolvedDensity,
    cosmogrid::CosmologicalGrid, ProbeA::String, ProbeB::String)

This function evaluates and returns the [`aₗₘCovariance`](@ref), according to the following
formula:
```math
\\Sigma_{i j}^{\\mathrm{AB}}(\\ell)=\\sqrt{\\frac{2}{(2 \\ell+1) \\Delta \\ell
f_{\\mathrm{sky}}}}\\left(C_{i j}^{\\mathrm{AB}}(\\ell)+N_{i j}^{\\mathrm{AB}}(\\ell)
\\right),
```
where ``\\mathrm{A}`` and ``\\mathrm{B}`` are the probes, ``\\mathrm{i}`` and
``\\mathrm{j}`` are the tomographic bins, ``\\ell`` and ``\\Delta\\ell`` are respectively the
central value and width of the considered multipole bin, ``f_{\\mathrm{sky}}`` is the sky
fraction covered by the field considered, ``N_{i j}^{\\mathrm{AB}}`` is the noise matrix.
"""
function InstantiateEvaluateCovariance(cℓ::AbstractCℓ, ConvDens::AbstractConvolvedDensity,
    cosmogrid::CosmologicalGrid, ProbeA::String, ProbeB::String)
    Cov = aₗₘCovariance()
    Cov.Cℓ = cℓ
    Cov.Noise      = zeros(size(cℓ.CℓArray))
    Cov.Covariance = zeros(size(cℓ.CℓArray))
    EvaluateNoise!(Cov, ConvDens, ProbeA::String, ProbeB::String)
    EvaluateCovariance!(Cov, cosmogrid)
    InvertCovariance!(Cov)
    return Cov
end

function EvaluateNoise!(Cov::aₗₘCovariance, ConvDens::AbstractConvolvedDensity,
    ProbeA::String, ProbeB::String)
    if ProbeA == ProbeB
        if ProbeA =="Lensing"
            probe_factor = 0.3
        else
            probe_factor = 1.0
        end
        for lidx in 1:length(Cov.Cℓ.CℓArray[:,1,1])
            for iidx in 1:length(Cov.Cℓ.CℓArray[1,:,1])
                Cov.Noise[lidx, iidx, iidx] = probe_factor^2 ./
                (ConvDens.SurfaceDensityArray[iidx]*3437.746771^2)
                #TODO: the multiplication here is to convert from square degrees to
                #steradians. This should be done better...
            end
        end
    end
end

function EvaluateCovariance!(Cov::aₗₘCovariance, cosmogrid::CosmologicalGrid)
    Cov.Covariance = (Cov.Cℓ.CℓArray + Cov.Noise)
    f_sky = 0.363610260832152
    #TODO: this is hardcoded...
    for (ℓidx, ℓvalue) in enumerate(cosmogrid.ℓBinCenters)
        Cov.Covariance[ℓidx,:,:] *= sqrt(2/((2* ℓvalue+1)*f_sky*cosmogrid.ℓBinWidths[ℓidx]))
    end
end

function InvertCovariance!(Cov::AbstractCovariance)
    for ℓidx in 1:length(Cov.Covariance[:,1,1])
        Cov.Covariance⁻¹[ℓidx,:,:] = inv(Cov.Covariance[ℓidx,:,:])
    end
end

"""
    InstantiateEvaluateCovariance(Covaₗₘ::aₗₘCovariance)

This function evaluates and returns the [`CℓCovariance`](@ref), using the following formula:

```math
\\boldsymbol{\\Xi}(\\ell)= \\left(\\boldsymbol{D}_{n}^{T}\\left( \\left(\\boldsymbol{\\Sigma}(\\ell)\\right)^{-1} \\otimes
\\left(\\boldsymbol{\\Sigma}(\\ell)\\right)^{-1}\\right) \\boldsymbol{D}_{n}\\right)^{-1}
```
where ``\\boldsymbol{D}_{n}`` is the [`DuplicationMatrix`](@ref).
"""
function InstantiateEvaluateCovariance(Covaₗₘ::aₗₘCovariance)
    Cov = CℓCovariance()
    ℓnumber = length(Covaₗₘ.Cℓ.CℓArray[:,1,1])
    inumber = length(Covaₗₘ.Cℓ.CℓArray[1,1,:])
    D = DuplicationMatrix(inumber)
    Dᵀ = Transpose(D)
    kronCovaₗₘ = zeros(inumber^2,inumber^2)
    DᵀkronCovaₗₘ = zeros(floor(Int,inumber*0.5*(inumber+1)),inumber^2)
    Cov.Covariance = zeros(ℓnumber, floor(Int,inumber*0.5*(inumber+1)),
    floor(Int,inumber*0.5*(inumber+1)))
    Cov.Covariance = zeros(size(Cov.Covariance))
    Cov.Covariance⁻¹ = zeros(size(Cov.Covariance))
    TempCov = zeros(floor(Int,inumber*0.5*(inumber+1)),
    floor(Int,inumber*0.5*(inumber+1)))
    for ℓ in 1:ℓnumber
        kronCovaₗₘ = kron(Covaₗₘ.Covariance⁻¹[ℓ,:,:], Covaₗₘ.Covariance⁻¹[ℓ,:,:])
        LinearAlgebra.mul!(DᵀkronCovaₗₘ, Dᵀ, kronCovaₗₘ)
        LinearAlgebra.mul!(TempCov, DᵀkronCovaₗₘ, D)
        Cov.Covariance⁻¹[ℓ,:,:] = TempCov
        Cov.Covariance[ℓ,:,:] = inv(Cov.Covariance⁻¹[ℓ,:,:])
        Matrix = (Cov.Covariance[ℓ,:,:] .+ transpose(Cov.Covariance[ℓ,:,:])) ./2
        Cov.Covariance[ℓ,:,:] = Matrix
        Cov.Covariance⁻¹[ℓ,:,:] = inv(Cov.Covariance[ℓ,:,:])
    end
    return Cov
end


text:
@enum OutboxKind send=1 acknowledge

abstract type EventDispatcher end

mutable struct RedisEventDispatcher <: EventDispatcher
    conn::Redis.RedisConnection
    group_id::AbstractString
    consumer_id::AbstractString
    state::ServiceState
    inbox::Channel{Event}
    outbox::Channel{Tuple{OutboxKind, Event}}
    claim_time::Dates.Period
end

function RedisEventDispatcher(group_id::AbstractString, consumer_id::AbstractString; host="127.0.0.1", port=6379, password="", db=0, outbox_size=32, inbox_size=32, claim_time=Dates.Minute(15))
    conn = Redis.RedisConnection(;host, port, password, db)
    inbox = Channel{Event}(inbox_size)
    outbox = Channel{Tuple{OutboxKind, Event}}(outbox_size)
    RedisEventDispatcher(conn, group_id, consumer_id, stopped, inbox, outbox, claim_time)
end

function is_connected(dispatcher::RedisEventDispatcher)
    Redis.is_connected(dispatcher.conn)
end

function is_running(dispatcher::RedisEventDispatcher)
    dispatcher.state == running
end

function _redis_data_to_event(stream, eventarr)
    publish_id = eventarr[1]
    eventdata = eventarr[2]
    recorddata = Dict{String,Any}()
    for ind = 1:2:size(eventdata, 1)
        recorddata[eventdata[ind]] = eventdata[ind + 1]
    end

    id = UUID(recorddata["id"])
    correlation_id = UUID(recorddata["meta.corr_id"])
    causation_id = UUID(recorddata["meta.caus_id"])
    datetime = Dates.DateTime(recorddata["meta.datetime"])
    data = CBOR.decode(Vector{UInt8}(get(recorddata, "data", "\xa0")))
    custom = CBOR.decode(Vector{UInt8}(get(recorddata, "custom", "\xa0")))
    aggregate_id_str = get(recorddata, "aggregate.id", nothing)
    aggregate_id = if isnothing(aggregate_id_str)
        nothing
    else
        UUID(aggregate_id_str)
    end
    stream_info = parse(StreamInfo, stream)

    Event(
        stream_info.context_name,
        stream_info.aggregate_name,
        stream_info.event_name,
        type=stream_info.type,
        id=id,
        publish_id=publish_id,
        aggregate_id=aggregate_id,
        datetime=datetime,
        correlation_id=correlation_id,
        causation_id=causation_id,
        published=true,
        custom=custom,
        data=data,
    )
end

function _run(dispatcher::RedisEventDispatcher, streams::Vector{String})
    # Create a group for each stream
    for stream in streams
        try
            redis_args = ["XGROUP", "CREATE", stream, dispatcher.group_id, "\$", "MKSTREAM"]
            Redis.execute_command(dispatcher.conn, redis_args)
        catch _err
            # TODO: Check the error type is correct
        end
    end

    last_claim = Dates.Date(0)

    # Endless loop to poll for items
    while dispatcher.state == running
        while isready(dispatcher.outbox)
            try
                (kind, event) = take!(dispatcher.outbox)
                redis_args = if kind == acknowledge
                    ["XACK", string(stream_info(event)), dispatcher.group_id, event.publish_id]
                elseif kind == send
                    eventid = if isnothing(event.publish_id) "*" else event.publish_id end
                    redis_args = [
                        "XADD",
                        string(stream_info(event)),
                        eventid,
                        "id", string(event.id),
                        "meta.corr_id", string(event.metadata.correlation_id),
                        "meta.caus_id", string(event.metadata.causation_id),
                        "meta.datetime", string(event.metadata.datetime),
                        "data", CBOR.encode(event.data),
                        "custom", CBOR.encode(event.custom),
                    ]

                    if !isnothing(event.aggregate.id)
                        redis_args =  [redis_args; ["aggregate.id", event.aggregate.id]]
                    end

                    redis_args
                else
                    nothing
                end

                if !isnothing(redis_args)
                    Redis.execute_command(dispatcher.conn, redis_args)
                end

                # Prevent blocking on the channel put!
                # If the inbox is full, the process is probably overwhelmed
                max_count = dispatcher.inbox.sz_max - length(dispatcher.inbox.data)

                # If there's no space we don't need to fetch more items
                if max_count != 0
                    redis_args = ["XREADGROUP", "GROUP", dispatcher.group_id, dispatcher.consumer_id, "COUNT", max_count, "STREAMS", streams..., ">"]
                    entries = Redis.execute_command(dispatcher.conn, redis_args)

                    if !isnothing(entries)
                        for streamrecord in entries
                            stream = streamrecord[1]
                            events = streamrecord[2]
                            for eventdata in events
                                try
                                    event = _redis_data_to_event(stream, eventdata)
                                    put!(dispatcher.inbox, event)
                                catch err
                                    println("Receive error: ", err)
                                end
                            end
                        end
                    end
                end

                if last_claim < (Date.now() + dispatcher.claim_time)
                    # Claim old events
                    last_claim = Date.now()
                    for stream in Random.shuffle(streams)
                        max_count = dispatcher.inbox.sz_max - length(dispatcher.inbox.data)
                        if max_count == 0
                            break
                        end

                        redis_args = [
                            "XAUTOCLAIM",
                            stream,
                            dispatcher.group_id,
                            dispatcher.consumer_id,
                            Dates.Millisecond(dispatcher.claim_time).value,
                            "0-0",
                            "COUNT",
                            max_count,
                        ]

                        result = Redis.execute_command(dispatcher.conn, redis_args)

                        if !isnothing(result) && length(result) > 1
                            for eventdata in result[2]
                                try
                                    event = _redis_data_to_event(stream, eventdata)
                                    put!(dispatcher.inbox, event)
                                catch err
                                    println("Receive error: ", err)
                                end
                            end
                        end
                    end
                end
            catch err
                println("$(kind) error: ", err)
            end
        end

        if length(dispatcher.inbox.data) == 0
            sleep(0.1)
        end
    end
end

function run!(dispatcher::RedisEventDispatcher, streams::Vector{StreamInfo})
    if dispatcher.state == running
        throw(EventException("The dispatcher is already running"))
    end

    # Validate the streams

    streams = map(streams) do stream
        try
            if !(stream.type in EventTypes)
                throw(ArgumentError("stream has invalid type $(repr(stream))"))
            end

            string(stream)
        catch _err
            throw(ArgumentError("invalid stream $(repr(stream))"))
        end
    end
    dispatcher.state = running

    @async _run(dispatcher, streams)
end

function stop!(dispatcher::RedisEventDispatcher)
    if dispatcher.state == stopped
        throw(EventException("The dispatcher is not running"))
    end

    dispatcher.state = stopped
end

function ack!(dispatcher::RedisEventDispatcher, event::Event)
    put!(dispatcher.outbox, (acknowledge, event))
end

function Base.put!(dispatcher::RedisEventDispatcher, event::Event)
    put!(dispatcher.outbox, (send, event))
end

function Base.take!(dispatcher::RedisEventDispatcher)
    take!(dispatcher.inbox)
end

function Base.isready(dispatcher::RedisEventDispatcher)
    isready(dispatcher.inbox)
end


text:
"""
    LessOrEqualConstant(x::SeaPearl.AbstractIntVar, v::Int)

Inequality constraint, `x <= v`
"""
struct LessOrEqualConstant <: OnePropagationConstraint
    x       ::AbstractIntVar
    v       ::Int
    active  ::StateObject{Bool}
    function LessOrEqualConstant(x, v, trailer)
        constraint = new(x, v, StateObject(true, trailer))
        addOnDomainChange!(x, constraint)
        return constraint
    end
end

# Base.show(io::IO, c::LessOrEqualConstant) = write(io, "LessOrEqualConstant constraint: ", c.x, " <= ", c.v)

"""
    propagate!(constraint::LessOrEqualConstant, toPropagate::Set{Constraint}, prunedDomains::CPModification)

`LessOrEqualConstant` propagation function. Basically remove the values above `v`.
"""
function propagate!(constraint::LessOrEqualConstant, toPropagate::Set{Constraint}, prunedDomains::CPModification)
    setValue!(constraint.active, false)

    addToPrunedDomains!(prunedDomains, constraint.x, removeAbove!(constraint.x.domain, constraint.v))
    triggerDomainChange!(toPropagate, constraint.x)
    if constraint in toPropagate
        pop!(toPropagate, constraint)
    end
    return !isempty(constraint.x.domain)
end

variablesArray(constraint::LessOrEqualConstant) = [constraint.x]

function Base.show(io::IO, ::MIME"text/plain", con::LessOrEqualConstant)
    println(io, typeof(con), ": ", con.x.id, " <= ", con.v, ", active = ", con.active)
    println(io, "   ", con.x)
end

function Base.show(io::IO, con::LessOrEqualConstant)
    print(io, typeof(con), ": ", con.x.id, " <= ", con.v)
end

struct LessOrEqual <: Constraint
    x       ::AbstractIntVar
    y       ::AbstractIntVar
    active  ::StateObject{Bool}

    """
        LessOrEqual(x::AbstractIntVar, y::AbstractIntVar, trailer::Trailer)

    Inequality between variables constraint, states that `x <= y`
    """
    function LessOrEqual(x, y, trailer)
        constraint = new(x, y, StateObject(true, trailer))
        addOnDomainChange!(x, constraint)
        addOnDomainChange!(y, constraint)
        return constraint
    end
end

"""
    propagate!(constraint::LessOrEqual, toPropagate::Set{Constraint}, prunedDomains::CPModification)

`LessOrEqual` propagation function.
"""
function propagate!(constraint::LessOrEqual, toPropagate::Set{Constraint}, prunedDomains::CPModification)

    if maximum(constraint.x.domain) <= minimum(constraint.y.domain)
        setValue!(constraint.active, false)
    end

    prunedX = removeAbove!(constraint.x.domain, maximum(constraint.y.domain))
    if !isempty(prunedX)
        addToPrunedDomains!(prunedDomains, constraint.x, prunedX)
        triggerDomainChange!(toPropagate, constraint.x)
    end

    prunedY = removeBelow!(constraint.y.domain, minimum(constraint.x.domain))
    if !isempty(prunedY)
        addToPrunedDomains!(prunedDomains, constraint.y, prunedY)
        triggerDomainChange!(toPropagate, constraint.y)
    end
    return !isempty(constraint.x.domain) && !isempty(constraint.y.domain)
end

variablesArray(constraint::LessOrEqual) = [constraint.x, constraint.y]

function Base.show(io::IO, ::MIME"text/plain", con::LessOrEqual)
    println(io, typeof(con), ": ", con.x.id, " ≤ ", con.y.id, ", active = ", con.active)
    println(io, "   ", con.x)
    println(io, "   ", con.y)
end

function Base.show(io::IO, con::LessOrEqual)
    print(io, typeof(con), ": ", con.x.id, " ≤ ", con.y.id)
end


text:
# TODO

struct GEFile <: MRIFile
  filename::String
end

function MRIBase.RawAcquisitionData(f::GEFile)
  error("Not yet implemented!")
end


text:
function get_earliest(obj::TimeStampedCache)
    return jcall(obj, "getEarliest", TimeStamped, ())
end

function get_latest(obj::TimeStampedCache)
    return jcall(obj, "getLatest", TimeStamped, ())
end

function get_neighbors(obj::TimeStampedCache, arg0::AbsoluteDate)
    return jcall(obj, "getNeighbors", Stream, (AbsoluteDate,), arg0)
end

function get_neighbors_size(obj::TimeStampedCache)
    return jcall(obj, "getNeighborsSize", jint, ())
end


text:
filename = pwd()*"/testinputsave.json"
@show pwd()
input = MultilevelInput()
saveinput(input, filename)
input2 = loadinput(MultilevelInput, filename)
rm(filename)
@test input == input2