text:
module Axes
import ..Cubes: caxes, Cubes
using Dates
using Base.Iterators: take, drop

"""
    abstract CubeAxis{T}

Supertype of all axes. Every `CubeAxis` is an 1D Cube itself and can be passed
to mapCube operations.
"""
abstract type CubeAxis{T,S} end

Base.size(x::CubeAxis)=(length(x.values),)
Base.size(x::CubeAxis,i)=i==1 ? length(x.values) : error("Axis has only a single dimension")
Base.ndims(x::CubeAxis)=1

"""
    CategoricalAxis{T,S}

To represent axes that are categorical, where `T` is the element type.
The type parameter `S` denotes the axis name (a symbol).
The default constructor is:

    CategoricalAxis(axname::String,values::Vector{T})

"""
struct CategoricalAxis{T,S,RT} <: CubeAxis{T,S}
  values::RT
end

CategoricalAxis(s::Symbol,v)=CategoricalAxis{eltype(v),s,typeof(v)}(v)
CategoricalAxis(s::AbstractString,v)=CategoricalAxis(Symbol(s),v)

struct RangeAxis{T,S,R<:AbstractVector{T}} <: CubeAxis{T,S}
  values::R
end

"""
    RangeAxis{T,S,R}

To represent axes that are categorical, where `T` is the element type.
The type parameter `S` denotes the axis name (a symbol) and `R` the type of the
range which is used to represent the axis values.
The default constructor is:

    RangeAxis(axname::String,values::Range{T})

"""
RangeAxis(s::Symbol,v::AbstractVector{T}) where T = RangeAxis{T,s,typeof(v)}(v)
RangeAxis(s::AbstractString,v)=RangeAxis(Symbol(s),v)

Base.length(a::CubeAxis)=length(a.values)

"""
    axcopy(x,vals)

Makes a copy of a `CubeAxis` with the values `vals`
"""
axcopy(ax::RangeAxis,vals) = RangeAxis(axname(ax),vals)
axcopy(ax::CategoricalAxis,vals) = CategoricalAxis(axname(ax),vals)
axcopy(ax::RangeAxis) = RangeAxis(axname(ax),copy(ax.values))
axcopy(ax::CategoricalAxis) = CategoricalAxis(axname(ax),copy(ax.values))

Base.show(io::IO,a::RangeAxis)=print(io,rpad(Axes.axname(a),20," "),"Axis with ",length(a)," Elements from ",first(a.values)," to ",last(a.values))
function Base.show(io::IO,a::CategoricalAxis)
  print(io,rpad(Axes.axname(a),20," "), "Axis with ", length(a), " elements: ")
  if length(a.values)<10
    for v in a.values
      print(io,v," ")
    end
  else
    for v in take(a.values,2)
      print(io,v," ")
    end
    print(io,".. ")
    for v in drop(a.values,length(a.values)-2)
      print(io,v," ")
    end
  end
end

caxes(x::CubeAxis)=CubeAxis[x]

axname(::Type{<:CubeAxis{<:Any,U}}) where U = string(U)
axname(::CubeAxis{<:Any,U}) where U = string(U)
axsym(::CubeAxis{<:Any,S}) where S = S

get_step(r::AbstractRange)=step(r)
get_step(r::AbstractVector)=length(r)==0 ? zero(eltype(r)) : r[2]-r[1]

axVal2Index_ub(a::RangeAxis, v; fuzzy=false)=axVal2Index(a,v-abshalf(get_step(a.values)),fuzzy=fuzzy)
axVal2Index_lb(a::RangeAxis, v; fuzzy=false)=axVal2Index(a,v+abshalf(get_step(a.values)),fuzzy=fuzzy)

axVal2Index_ub(a::RangeAxis, v::Date; fuzzy=false)=axVal2Index(a,DateTime(v)-abshalf(get_step(a.values)),fuzzy=fuzzy)
axVal2Index_lb(a::RangeAxis, v::Date; fuzzy=false)=axVal2Index(a,DateTime(v)+abshalf(get_step(a.values)),fuzzy=fuzzy)

abshalf(a) = abs(a/2)
abshalf(a::Day) = abs(Millisecond(a)/2)
abshalf(a::Month) = iseven(Dates.value(a)) ? a/2 : Month(a÷2) + Day(15)

get_bb(ax::RangeAxis) = first(ax.values)-abshalf(get_step(ax.values)), last(ax.values)+abshalf(get_step(ax.values))
function axisfrombb(name,bb,n)
  offs = (bb[2]-bb[1])/(2*n)
  RangeAxis(name,range(bb[1]+offs,bb[2]-offs,length=n))
end

function axVal2Index(a::RangeAxis{<:Any,<:Any,<:AbstractRange},v;fuzzy=false)
  dt = v-first(a.values)
  r = round(Int,dt/step(a.values))+1
  return max(1,min(length(a.values),r))
end

convert_time(T::Type{<:TimeType}, v::TimeType) = T(year(v), month(v), day(v), hour(v), minute(v), second(v))
convert_time(T::Type{<:TimeType}, v::Date) = T(year(v), month(v), day(v), 0, 0, 0)
convert_time(::Type{Date},v::TimeType) = Date(year(v),month(v),day(v))
convert_time(::Type{Date},v::Date) = Date(year(v),month(v),day(v))

function axVal2Index(a::RangeAxis{T},v;fuzzy=false) where T<:TimeType
  vconverted = convert_time(T,v)
  dd = map(i->abs((i-vconverted)),a.values)
  mi,ind = findmin(dd)
  return ind
end
function axVal2Index(a::RangeAxis{T,<:Any,<:AbstractRange},v;fuzzy=false) where T<:TimeType
  vconverted = convert_time(T,v)
  dd = map(i->abs((i-vconverted)),a.values)
  mi,ind = findmin(dd)
  return ind
end
function axVal2Index(axis::CategoricalAxis{String},v::String;fuzzy::Bool=false)
  r=findfirst(isequal(v),axis.values)
  if r===nothing
    if fuzzy
      r=findall(axis.values) do i
        startswith(lowercase(i),lowercase(v[1:min(length(i),length(v))]))
      end
      if length(r)==1
        return(r[1])
      else
        error("Could not find unique value of $v in $axis")
      end
    else
      error("$v not found in $axis")
    end
  end
  r
end
axVal2Index(x,v::CartesianIndex{1};fuzzy::Bool=false)=min(max(v.I[1],1),length(x))
function axVal2Index(x,v;fuzzy::Bool=false)
  i = findfirst(isequal(v),x.values)
  if isa(i,Nothing)
      if fuzzy==true
          dd = map(i->abs(i-v),x.values)
          mi,ind = findmin(dd)
          return ind
      else
          error("Value $v not found in x")
      end
  else
      return i
  end
end
abstract type AxisDescriptor end
struct ByName <: AxisDescriptor
  name::String
end
struct ByInference <: AxisDescriptor end
struct ByValue <: AxisDescriptor
  v::CubeAxis
end
struct ByFunction <: AxisDescriptor
  f::Function
end

const VecOrTuple{S} = Union{Vector{<:S},Tuple{Vararg{<:S}}} where S

#findAxis(a::Any,c::VecOrTuple{<:CubeAxis}) = findAxis(get_descriptor(desc),c)
get_descriptor(a::String)=ByName(a)
get_descriptor(a::Symbol)=ByName(String(a))
get_descriptor(a::CubeAxis)=ByValue(a)
get_descriptor(a::Function)=ByFunction(a)
get_descriptor(a)=error("$a is not a valid axis description")
get_descriptor(a::AxisDescriptor)=a

function findAxis(bs::ByName,axlist::VecOrTuple{CubeAxis})
  matchstr=bs.name
  ism=findall(i->startswith(lowercase(axname(i)),lowercase(matchstr)),axlist)
  isempty(ism) && return nothing
  if length(ism)>1
    f = axlist[ism[1]]
    all(i->i==f,ism) || error("Multiple axes found matching string $matchstr")
    return f
  else
    return ism[1]
  end
end
function findAxis(bv::ByValue,axlist::VecOrTuple{CubeAxis})
  v=bv.v
  return findfirst(i->i==v,axlist)
end
function getAxis(desc,axlist::VecOrTuple{CubeAxis})
  i = findAxis(desc,axlist)
  if isa(i,Nothing)
    return nothing
  else
    return axlist[i]
  end
end
getOutAxis(desc,axlist,incubes,pargs,f) = getAxis(desc,unique(axlist))
function getOutAxis(desc::ByFunction,axlist,incubes,pargs,f)
  outax = desc.f(incubes,pargs)
  isa(outax,CubeAxis) || error("Axis Generation function $(desc.f) did not return an axis")
  outax
end
import DataStructures: counter
function getOutAxis(desc::Tuple{ByInference},axlist,incubes,pargs,f)
  inAxes = map(caxes,incubes)
  inAxSmall = map(i->filter(j->in(j,axlist),i) |>collect,inAxes)
  inSizes = map(i->(map(length,i)...,),inAxSmall)
  intypes = map(eltype, incubes)
  testars = map((s,it)->zeros(it,s...),inSizes, intypes)
  map(testars) do ta
    ta .= rand(Base.nonmissingtype(eltype(ta)),size(ta)...)
    if eltype(ta) >: Missing
      # Add some missings
      randind = rand(1:length(ta),length(ta)÷10)
      ta[randind] .= missing
    end
  end
  resu = f(testars...,pargs...)
  isa(resu,AbstractArray) || isa(resu,Number) || isa(resu,Missing) || error("Function must return an array or a number")
  (isa(resu,Number) || isa(resu,Missing)) && return ()
  outsizes = size(resu)
  outaxes = map(outsizes,1:length(outsizes)) do s,il
    if s>2
      i = findall(i->i==s,length.(axlist))
      if length(i)==1
        return axlist[i[1]]
      elseif length(i)>1
        @info "Found multiple matching axes for output dimension $il"
      end
    end
    return RangeAxis("OutAxis$(il)",1:s)
  end
  if !allunique(outaxes)
    #TODO: fallback with axis renaming in this case
    error("Could not determine unique output axes from output shape")
  end
  return (outaxes...,)
end
"""
    getAxis(desc::String, c)

Given the string of an axis name and a cube, returns this axis of the cube.
"""
getAxis(desc,c)=getAxis(desc,caxes(c))
getAxis(desc::ByValue,axlist::Vector{T}) where {T<:CubeAxis}=desc.v

"Fallback method"
findAxis(desc,c)=findAxis(desc,caxes(c))
findAxis(a,axlist::VecOrTuple{CubeAxis}) = findAxis(get_descriptor(a),axlist)

renameaxis(r::RangeAxis{T,<:Any,V}, newname) where {T,V} = RangeAxis{T,Symbol(newname),V}(r.values)
renameaxis(r::CategoricalAxis{T,<:Any,V}, newname) where {T,V} = CategoricalAxis{T,Symbol(newname),V}(r.values)

import Base.==
import Base.isequal
==(a::CubeAxis,b::CubeAxis)=(a.values==b.values) && (axname(a)==axname(b))
isequal(a::CubeAxis, b::CubeAxis) = a==b

using YAXArrayBase: YAXArrayBase
#Implement yaxarray interface
YAXArrayBase.dimname(x::CubeAxis,_) = axname(x)
YAXArrayBase.dimvals(x::CubeAxis,_) = x.values
YAXArrayBase.iscontdim(::RangeAxis,_) = true
YAXArrayBase.iscontdim(::CategoricalAxis,_) = false
end


text:
using YaoBase, YaoArrayRegister
export GeneralMatrixBlock, matblock

"""
    GeneralMatrixBlock{M, N, MT} <: PrimitiveBlock{N}

General matrix gate wraps a matrix operator to quantum gates. This is the most
general form of a quantum gate. `M` is the hilbert dimension (first dimension),
`N` is the hilbert dimension (second dimension) of current quantum state. For
most quantum gates, we have ``M = N``.
"""
struct GeneralMatrixBlock{M,N,T,MT<:AbstractMatrix{T}} <: PrimitiveBlock{N}
    mat::MT

    function GeneralMatrixBlock{M,N}(m::MT) where {M,N,T,MT<:AbstractMatrix{T}}
        (1 << M, 1 << N) == size(m) ||
            throw(DimensionMismatch("expect a $(1<<M) x $(1<<N) matrix, got $(size(m))"))

        return new{M,N,T,MT}(m)
    end
end

GeneralMatrixBlock(m::AbstractMatrix) = GeneralMatrixBlock{log2i.(size(m))...}(m)

"""
    matblock(m::AbstractMatrix)

Create a [`GeneralMatrixBlock`](@ref) with a matrix `m`.

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> matblock(ComplexF64[0 1;1 0])
matblock(...)
```

!!!warn

    Instead of converting it to the default data type `ComplexF64`,
    this will return its contained matrix when calling `mat`.
"""
matblock(m::AbstractMatrix) = GeneralMatrixBlock(m)

"""
    matblock(m::AbstractMatrix)

Create a [`GeneralMatrixBlock`](@ref) with a matrix `m`.
"""
matblock(m::AbstractBlock) = GeneralMatrixBlock(mat(m))

cache_key(m::GeneralMatrixBlock) = hash(m.mat)

"""
    mat(A::GeneralMatrixBlock)

Return the matrix of general matrix block.

!!!warn

    Instead of converting it to the default data type `ComplexF64`,
    this will return its contained matrix.
"""
mat(A::GeneralMatrixBlock) = A.mat

function mat(::Type{T}, A::GeneralMatrixBlock) where {T}
    if eltype(A.mat) == T
        return A.mat
    else
        # this errors before, but since we allow one to specify T in mat
        # this should be allowed but with a suggestion
        @warn "converting $(eltype(A.mat)) to eltype $T, consider create another matblock with eltype $T"
        return copyto!(similar(A.mat, T), A.mat)
    end
end

Base.:(==)(A::GeneralMatrixBlock, B::GeneralMatrixBlock) = A.mat == B.mat
Base.copy(A::GeneralMatrixBlock) = GeneralMatrixBlock(copy(A.mat))
Base.adjoint(x::GeneralMatrixBlock) = Daggered(x)

ext: jl


text:
using LinearAlgebra
using SparseArrays

include("CHOLMOD.jl")
include("../src/SparseFactorizations.jl")
import IncompleteLU

using PyPlot

function laplacian_1d(n)
    return spdiagm(
        -1 => fill(-1.0,n-1),
         0 => fill( 2.0,n),
         1 => fill(-1.0,n-1)
    )
end

function laplacian_2d(n)
    Δ = laplacian_1d(n)
    Id = sparse(I,n,n)
    return kron(Δ,Id) + kron(Id,Δ)
end

function laplacian_3d(n)
    Δ = laplacian_1d(n)
    Id = sparse(I,n,n)
    return kron(Δ,Id,Id) + kron(Id,Δ,Id) + kron(Id,Id,Δ)
end

function benchmark()
    nreps = 3
    if true
        n = round.(Int, 10.0.^LinRange(1,2.2,20))
        mat = laplacian_2d
    else
        n = round.(Int, 10.0.^LinRange(log10(5),log10(30),10))
        mat = laplacian_3d
    end
    funs = (
        ("CHOLMOD, simplicial", A->CHOLMOD.Sparse(A,-1), SC_A -> CHOLMOD.simplicial_factorize(SC_A)),
        ("CHOLMOD, supernodal", A->CHOLMOD.Sparse(A,-1), SC_A -> CHOLMOD.supernodal_factorize(SC_A)),
        # ("SparseFactorizations, analysis", identity, A -> SparseFactorizations.ldlt_structure(A)),
        ("SparseFactorizations, factorization", A->(A,SparseFactorizations.ldlt_structure(A)), ((A,P),) -> SparseFactorizations.ildlt(A,P)),
        ("SparseFactorizations", identity, A -> SparseFactorizations.ldlt(A)),
        ("SparseFactorizations, lof", identity, A -> SparseFactorizations.ildlt_lof(A,size(A,2))),
        ("SparseFactorizations, tol", identity, A -> SparseFactorizations.ildlt_tol(A,0.0)),
        ("IncompleteLU", identity, A -> IncompleteLU.ilu(A,τ=0.0)),
    )

    timings = [zeros(size(n)) for f in funs]

    # Warm up
    A = mat(4)
    for (name,setup,run) in funs
        run(setup(A))
    end

    for kn = 1:length(n)
        println("n = $(n[kn])")
        A = mat(n[kn])
        p = CHOLMOD.amd(A)
        A = A[p,p]
        for (kf,(name,setup,run)) in enumerate(funs)
            args = setup(A)
            timings[kf][kn] = minimum(@elapsed run(args) for i = 1:nreps)
        end
    end
    return (n,funs,timings)
end

function plot((n,funs,timings))
    fig = figure(figsize=(8,4))
    try
        for ((name,_),timing) in zip(funs,timings)
            loglog(n,timing,label=name)
        end
        xlabel(L"n")
        ylabel("runtime [seconds]")
        legend(
            frameon=false,
            loc="center left",
            bbox_to_anchor=(1,0.5)
        )
        tight_layout()
        savefig(@__DIR__()*"/benchmark.png")
    finally
        close(fig)
    end
end

max_stars_repo_name: ettersi/SparseFactorizations

max_issues_repo_name: ettersi/SparseFactorizations

max_forks_repo_name: ettersi/SparseFactorizations


text:
# This routine ABSOLUTELY should not allocate any memory, called inside time loop.
@inbounds @fastmath function compute_gradient!(issp::Int64, pac, pap)
	# aliases
	p1=pap[1].w1[:t][:p]
	p1p=pap[1].w1[:tp][:p]
	p1pp=pap[1].w1[:tpp][:p]
	p2p=pap[2].w1[:tp][:p]
	fc1=pac.fc[:δtI]

	dpdx1=pap[1].w1[:dx][:p]
	dpdx2=pap[2].w1[:dx][:p]
	dpdz1=pap[1].w1[:dz][:p]
	dpdz2=pap[2].w1[:dz][:p]

	gKI=pap[1].ss[issp].grad_mod[:KI]
	grhovxI=pap[1].ss[issp].grad_mod[:rhovxI]
	grhovzI=pap[1].ss[issp].grad_mod[:rhovzI]

	gmodKI!(gKI,p1,p1p,p1pp,p2p,pac.ic[:nx],pac.ic[:nz],fc1)
	gmodrhovxI!(grhovxI,dpdx1,dpdx2,pac.ic[:nx],pac.ic[:nz])
	gmodrhovzI!(grhovzI,dpdz1,dpdz2,pac.ic[:nx],pac.ic[:nz])
end

@inbounds @fastmath function gmodKI!(gKI,p1,p1p,p1pp,p2p,nx,nz,δtI)
	for ix=1:nx
		@simd for iz=1:nz
			# p at [it], pp at [it-1]	# dpdx and dpdz at [it]	# gradients w.r.t inverse of modKI, i.e., 1/rho/c^2
			@inbounds gKI[iz,ix] += ((-1.0 * (p1pp[iz, ix] + p1[iz, ix] - 2.0 * p1p[iz, ix]) * δtI * δtI) *  p2p[iz,ix])
		end
	end
end

@inbounds @fastmath function gmodrhovxI!(grhovxI,dpdx1,dpdx2,nx,nz)
	for ix=1:nx
		@simd for iz=1:nz
			# gradient w.r.t. inverse of rho on vx and vz grids
			@inbounds grhovxI[iz,ix] += (- dpdx2[iz,ix,1]*dpdx1[iz,ix,1])
		end
	end

end
@inbounds @fastmath function gmodrhovzI!(grhovzI,dpdz1,dpdz2,nx,nz)
	for ix=1:nx
		@simd for iz=1:nz
			# gradient w.r.t. inverse of rho on vx and vz grids
			@inbounds grhovzI[iz,ix] += (- dpdz2[iz,ix,1]*dpdz1[iz,ix,1])
		end
	end
end

@inbounds @fastmath function scale_gradient!(issp::Int64,pap,δ)
	gKI=pap[1].ss[issp].grad_mod[:KI]
	grhovxI=pap[1].ss[issp].grad_mod[:rhovxI]
	grhovzI=pap[1].ss[issp].grad_mod[:rhovzI]
	"gradient is formed by intergration over time, hence multiply with δt, but why not?"
	"I don't completely understand where the factors δx and δz are coming from..."
	"probably the source term should not be multiplied by δxI and δzI during adjoint propagation"
	rmul!(gKI,δ)
	rmul!(grhovxI,δ)
	rmul!(grhovzI,δ)
end

function grad_modrr!(pac)
	grad_modrr_sprayrr!(pac.grad_mod[:rhoI],pac.grad_mod[:rhovxI],pac.grad_mod[:rhovzI])
	grad_modrr_sprayrhovxI!(pac.grad_mod[:rhoI],pac.grad_mod[:rhovxI])
	grad_modrr_sprayrhovzI!(pac.grad_mod[:rhoI],pac.grad_mod[:rhovzI])
end

function stack_grads!(pac, pap)
	# theses are SharedArrays
	gmodKI=pac.grad_mod[:KI]
	gmodrhovxI=pac.grad_mod[:rhovxI]
	gmodrhovzI=pac.grad_mod[:rhovzI]
	for issp in 1:length(pap[1].ss)
		gs=pap[1].ss[issp].grad_mod[:KI]
		for i in eachindex(gmodKI)
			gmodKI[i]+=gs[i]  # only update gmodKI
		end
		gs=pap[1].ss[issp].grad_mod[:rhovxI]
		for i in eachindex(gmodrhovxI)
			gmodrhovxI[i]+=gs[i]
		end
		gs=pap[1].ss[issp].grad_mod[:rhovzI]
		for i in eachindex(gmodrhovzI)
			gmodrhovzI[i]+=gs[i]
		end
	end
end

function update_gradient!(pac)
	nx, nz=pac.ic[:nx], pac.ic[:nz]
	nznxd = prod(length.(pac.model.mgrid))

	# combine rhovxI and rhovzI
	grad_modrr!(pac)

	gradient=pac.gradient
	# truncate
	gmodKI=view(pac.grad_mod[:KI],npml+1:nz-npml,npml+1:nx-npml)
	gmodrr=view(pac.grad_mod[:rhoI],npml+1:nz-npml,npml+1:nx-npml)
	for i in 1:nznxd
		# parameterization is  [:KI, :ρI, :null]
		gradient[i]=gmodKI[i]  # update gmodKI
		gradient[nznxd+i]=gmodrr[i]  # update gmodrr
	end
end


text:
module MatchV5 # LOLTools

using ..LOLTools: AbstractDTO, lol_api_server, http_action, call_api, nothing_in_event
using HTTP

"""
    Get a list of match ids by puuid
    matches_by_puuid(api_key::String,
                     region::String,
                     puuid::String ;
                     endpoint::HTTP.URI = lol_api_server(region),
                     action::Function = http_action,
                     event::Function = nothing_in_event)::Vector{String}
"""
function matches_by_puuid(api_key::String,
                          region::String,
                          puuid::String ;
                          endpoint::HTTP.URI = lol_api_server(region),
                          action::Function = http_action,
                          event::Function = nothing_in_event,
                          query::Union{Nothing,Dict{String,String}} = nothing)::Vector{String}
    call_api(Vector{String}, api_key, action, endpoint, "/lol/match/v5/matches/by-puuid/$puuid/ids", event, matches_by_puuid; query=query)
end

struct MatchDto <: AbstractDTO
    metadata
    info
end

"""
    Get a match by match id
    match_by_match_id(api_key::String,
                      region::String,
                      matchId::String ;
                      endpoint::HTTP.URI = lol_api_server(region),
                      action::Function = http_action,
                      event::Function = nothing_in_event)::MatchDto
"""
function match_by_match_id(api_key::String,
                           region::String,
                           matchId::String ;
                           endpoint::HTTP.URI = lol_api_server(region),
                           action::Function = http_action,
                           event::Function = nothing_in_event)::MatchDto
    call_api(MatchDto, api_key, action, endpoint, "/lol/match/v5/matches/$matchId", event, match_by_match_id)
end

struct MatchTimelineDto <: AbstractDTO
    metadata
    info
end

"""
    Get a match timeline by match id
    match_timeline_by_match_id(api_key::String,
                               region::String,
                               matchId::String ;
                               endpoint::HTTP.URI = lol_api_server(region),
                               action::Function = http_action,
                               event::Function = nothing_in_event)::MatchTimelineDto
"""
function match_timeline_by_match_id(api_key::String,
                                    region::String,
                                    matchId::String ;
                                    endpoint::HTTP.URI = lol_api_server(region),
                                    action::Function = http_action,
                                    event::Function = nothing_in_event)::MatchTimelineDto
    call_api(MatchTimelineDto, api_key, action, endpoint, "/lol/match/v5/matches/$matchId/timeline", event, match_timeline_by_match_id)
end

end # module LOLTools.MatchV5


text:
pushfirst!(LOAD_PATH, joinpath(@__DIR__, "..")) # add LAMMPS to environment stack

using LAMMPS
using Documenter
using DocumenterCitations
using Literate

DocMeta.setdocmeta!(LAMMPS, :DocTestSetup, :(using LAMMPS); recursive=true)

bib = CitationBibliography(joinpath(@__DIR__, "citations.bib"))

##
# Generate examples
##

const EXAMPLES_DIR = joinpath(@__DIR__, "..", "examples")
const OUTPUT_DIR   = joinpath(@__DIR__, "src/generated")

examples = [
    "Basic SNAP" => "snap",
    "Fitting SNAP" => "fitting_snap",
]

for (_, name) in examples
    example_filepath = joinpath(EXAMPLES_DIR, string(name, ".jl"))
    Literate.markdown(example_filepath, OUTPUT_DIR, documenter=true)
end

examples = [title=>joinpath("generated", string(name, ".md")) for (title, name) in examples]

makedocs(bib;
    modules=[LAMMPS],
    authors="CESMIX-MIT",
    repo="https://github.com/cesmix-mit/LAMMPS.jl/blob/{commit}{path}#{line}",
    sitename="LAMMPS.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://cesmix-mit.github.io/LAMMPS.jl",
        assets=String[],
        mathengine = MathJax3(),
    ),
    pages = [
        "Home" => "index.md",
        "Examples" => examples,
        "API" => "api.md",
    ],
    doctest = true,
    linkcheck = true,
    strict = true,
)

deploydocs(;
    repo="github.com/cesmix-mit/LAMMPS.jl",
    devbranch = "main",
    push_preview = true,
)


text:
# Features added version of new_model4.jl
# Seems like the this version is faster than the previous version
wordW(m)=m[1];
postagv(m)=m[2];
bufferm(m)=m[3];
stackm(m)=m[4];
actm(m)=m[5];
actembed(m)=m[6];
parserv(m)=m[7];
deprelv(m)=m[8]
treeNN(m)=m[9]
fmatrix(m)=m[10]

function initmodel1(o, s; ftype=Float32)
    model = Any[]

    # wordembedding # 1
    push!(model, initx(o[:wembed], 950))

    # postag-level initialization, KnetArray initialization
    dpostag = o[:posembed] # 2
    for (k, n, d) in ((:postag, 17, dpostag),)
        push!(model, Any[ initr(d, GPUFEATS=true) for i=1:n ])
    end

    # buffer-lstm initialization # 3
    r_b, wr_b = initbuf(o[:bufembed], o[:bufhidden]) # you may change lstm - gru with kwargs
    push!(model, [r_b, wr_b])

    # stack-lstm initialization # 4
    r_s, wr_s = initbuf(o[:stembed], o[:sthidden])
    push!(model, [r_s, wr_s])

    # action-lstm initialization # 5
    r_a, wr_a = initbuf(o[:actembed], o[:acthidden])
    push!(model, [r_a, wr_a])

    # action embeddings # 6
    p = o[:arctype](s)
    push!(model, initx(o[:actembed], p.nmove))

    # decision module initialization # 7
    mlpdims = (r_s.hiddenSize + r_b.hiddenSize + r_a.hiddenSize, o[:hidden]..., p.nmove)
    info("mlpdims: $mlpdims")
    decider = Any[]
    for i=2:length(mlpdims)
        push!(decider, initx(mlpdims[i], mlpdims[i-1]))
        push!(decider, initx(mlpdims[i], 1))
    end
    push!(model, decider)

    # deprel embeddings # 8
    depreldim = o[:deprel]
    push!(model, Any[ initr(depreldim, GPUFEATS=true) for i=1:p.nmove-1 ])

    # tree representer # 9
    rt, wrt = rnninit(o[:deprel]+o[:wembed], o[:wembed], rnnType=o[:treeType])
    push!(model, [rt, wrt])

    # fmatrix # 10
    push!(model, initx(o[:fembed], o[:featdim]))

    optims = optimizers(model, o[:optimization])
    return model, optims

end

function cache_wvecs(model, parsers)
    W = wordW(model)
    B = length(parsers)
    walls = Array{Any}(B)

    for i in 1:B
        p = parsers[i]
        x = (gpu()>=0 ? KnetArray(hcat(p.sentence.cavec...)) : Array(hcat(p.sentence.cavec...)))
        X = W * x # TODO : try one-by-one multiplication
        T = length(p.sentence); w_i = Array{Any}(T)
        for t in 1:T
            temp   = X[:, t]
            w_i[t] = reshape(temp, length(temp), 1) # do we need that? or just temp
        end
        walls[i] = w_i
    end
    return walls
end

####### Buffer-Related

# youts[i] -> output of taking wi as an input
# y2use -> youts[p.wptr]
function scan_buffer(model, wvecs, parser, hidden)
    t0 = zeros(Float32, hidden, 1) # takes single instance
    t0 = (gpu()>=0 ? KnetArray(t0): Array(t0))
    if length(parser.sentence) == 1 # ugly test time!
        y0 = t0;
        c0 = reshape(t0, length(t0), 1, 1)
        return y0, c0
    end

    bufmodel, pvecs, featm = bufferm(model), postagv(model), fmatrix(model)
    rbuf, wbuf = bufmodel[1], bufmodel[2]
    nrows = size(AutoGrad.getval(featm), 1)
    _t0f = zeros(Float32, nrows,1)
    t0f = (gpu() >= 0 ? KnetArray(_t0f) : Array(_t0f))

    sentence = parser.sentence
    T = length(sentence)
    seqs = parser.wptr
    range_ = T:-1:seqs

    # There is no way to be done in single rnnforw call, cell required!!!
    # No need to store hiddens, reshape youts
    # isapprox(reshape(hyouts[i], 128,), youts[i]) is true
    h0 = c0 = reshape(t0, hidden, 1, 1)
    youts, cyouts = Array{Any}(T), Array{Any}(T)
    for i in range_
        _i1 = cat1d(wvecs[i], pvecs[sentence.postag[i]])
        _feats = parser.sentence.feats[i]
        _i2 = (isempty(_feats) ? t0f : sum(featm[:, _feats], 2))
        i1  = cat1d(_i1, _i2)
        y1, h0, c0, = rnnforw(rbuf, wbuf, i1, h0, c0, hy=true, cy=true)
        youts[i], cyouts[i] = y1, c0 # fill the ith words output hidden
    end
    return youts, cyouts
end

function cache_bufbatch(model, walls, parsers, hidden)
    yalls = Array{Any}(length(parsers)) # i for ith parser
    B = length(parsers)
    for i in 1:B
        wvecs = walls[i];
        p = parsers[i]
        youts, cyouts = scan_buffer(model, wvecs, p, hidden)
        yalls[i] = Any[youts, cyouts]
    end
    return yalls
end

function scan_bufbatch(yalls, parsers, hidden)
    B = length(parsers)
    instate = zeros(Float32, hidden, 1)
    ybuf = (gpu() >=0 ? KnetArray(instate) : Array(instate))

    ybatch = Any[]
    for i in 1:B
        p = parsers[i]
        if p.wptr > p.nword # buffer empty
            push!(ybatch, ybuf)
        else
            youts, _ = yalls[i]
            push!(ybatch, youts[p.wptr]) # 2use
        end
    end
    ybatch = cat1d(ybatch...)
    nrow = div(length(ybatch), B)
    return reshape(ybatch, nrow, B)
end

####### Stack-related

# You are given new set of wvecs and re-calculate each hidden again
function scan_stack(model, wvecs, parser, hidden, embdim)
    sentence = parser.sentence
    T = parser.sptr
    if length(1:T) < 1 # empty stack
        instate = zeros(Float32, hidden, 1) # B 1
        t0 = (gpu() >= 0 ? KnetArray(instate) : Array(instate))
        return t0
    end

    stmodel, pvecs, featm = stackm(model), postagv(model), fmatrix(model)
    rst, wst = stmodel[1], stmodel[2]

    nrows = size(AutoGrad.getval(featm), 1)
    _t0f = zeros(Float32, nrows, 1)
    t0f = (gpu() >= 0 ? KnetArray(_t0f) : Array(_t0f))

    inputs = Any[]
    for i in 1:T
        wi = parser.stack[i]
        _i1 = cat1d(wvecs[wi], pvecs[sentence.postag[wi]])
        _feats = parser.sentence.feats[i]
        _i2 = (isempty(_feats) ? t0f : sum(featm[:, _feats], 2))
        i1  = cat1d(_i1, _i2)
        push!(inputs, i1)
    end
    allin = reshape(hcat(inputs...), embdim, 1, T)

    _, hyout, = rnnforw(rst, wst, allin, hy=true) # final-hidden
    return reshape(hyout, hidden, 1)
end

# Be careful walls and parsers should have the same order
function scan_stackbatch(model, walls, parsers, hidden, embdim)
    B = length(parsers)
    yalls = Any[]
    for i in 1:B
        p = parsers[i]; wvecs = walls[i];
        yst = scan_stack(model, wvecs, p, hidden, embdim)
        push!(yalls, yst)
    end
    yalls = cat1d(yalls...)
    ncols = B
    nrows = div(length(yalls), ncols)
    return reshape(yalls, nrows, ncols)
end

###### Action related

# compute batched actions
function scan_action(model, actions, state)
    acemb, actmodel = actembed(model), actm(model)
    yout0, hout0, cout0 = state[1], state[2], state[3]
    input = acemb[:, actions]
    yout, hout, cout = rnnforw(actmodel[1], actmodel[2], input, hout0, cout0, hy=true, cy=true)
    return (yout, hout, cout)
end

####### Embedding-related
# 1. Left-move : even move-numbers, pair, hw = p.wptr;dw = p.stack[p.sptr]
# 2. Right-move: odd move-numbers, pair,  hw = p.stack[p.sptr-1]; dw = p.stack[p.sptr]
function update_cache!(model, wvecs, cacbufs, p, m)
    depvecs, trNN = deprelv(model), treeNN(model)
    rt, wt = trNN[1], trNN[2]

    if m == 1
        return wvecs
    elseif m % 2 == 0 # left
        hw, dw = leftmovepair(p)
        wold = wvecs[hw]
        r  = depvecs[m-1] # -1 to reduce shift
        x1 = cat1d(wvecs[dw], r) # [x;r]
        input = reshape(x1, length(x1), 1) # batchsize 1

        # update wvec
        wnew, = rnnforw(rt, wt, input, wold)
        wvecs[hw] = wnew

        # update cachedbufs
        bufmodel, pvecs = bufferm(model), postagv(model)
        rbuf, wbuf = bufmodel[1], bufmodel[2]
        ybuf, cbuf = cacbufs
        yold, cold = cacbufs
        hyin, cyin = yold[hw], cold[hw]

        # No way to get rid of hidden returning
        # updates the hidden vectors
        yold[hw], _, cold[hw], = rnnforw(rbuf, wbuf, wnew, hyin, cyin, cy=true)

    else # right
        hw, dw = rightmovepair(p)
        wold = wvecs[hw]
        r  = depvecs[m-1] # -1 to reduce shift
        x1 = cat1d(wvecs[dw], r)
        input = reshape(x1, length(x1), 1)
        wnew, = rnnforw(rt, wt, input, wold)
        wvecs[hw] = wnew
    end
    #return wvecs # no need to return it only updates related .caches
end

function oracleloss(allmodel, sentences, arctype, hiddens, embdim;losses=nothing, pdrop=(0.0, 0.0))
    hidstack, hidact, hidbuf = hiddens
    parsers = map(arctype, sentences)
    mcosts = Array{Cost}(parsers[1].nmove)
    parserdone  = falses(length(parsers))
    mlpmodel = parserv(allmodel)

    # cache
    walls = cache_wvecs(allmodel, parsers)
    cachedbufs = cache_bufbatch(allmodel, walls, parsers, hidbuf)

    # init
    B = length(parsers) # bathcsize
    t0 = zeros(Float32, hidact, B, 1)
    h0 = c0 = (gpu()>=0 ? KnetArray(t0) : Array(t0))
    y0 = reshape(h0, hidact, B)
    totalloss = 0.0
    nstate = (y0, h0, c0)
    inacts = map(t->1, 1:B)

    while !all(parserdone)
        yact, hact, cact = scan_action(allmodel, inacts, nstate)
        nstate = (yact, hact, cact);
        ysts   = scan_stackbatch(allmodel, walls, parsers, hidstack, embdim)
        ybufs  = scan_bufbatch(cachedbufs, parsers, hidbuf)
        encoded = vcat(yact, ysts, ybufs)
        scores  = mlp(mlpmodel, encoded, pdrop=pdrop)
        logprobs = logp(scores, 1)

        inacts = Int[] # to get the next step's actions
        # iterative loss-val
        for (i, p) in enumerate(parsers)
            if parserdone[i]
                push!(inacts, 1) # not to change batchsize
                continue
            end
            movecosts(p, p.sentence.head, p.sentence.deprel, mcosts)
            goldmove = indmin(mcosts)

            if mcosts[goldmove] == typemax(Cost)
                parserdone[i] = true
                p.sentence.parse = p
                push!(inacts, 1) # not to change batchsize
            else
                totalloss -= logprobs[goldmove, i]
                # update before moving to the next step
                update_cache!(allmodel, walls[i], cachedbufs[i], p, goldmove)
                move!(p, goldmove); push!(inacts, goldmove);
                if losses != nothing
                    loss1 = -getval(logprobs)[goldmove,i]
                    losses[1] += loss1
                    losses[2] += 1
                    if losses[2] < 1000
                        losses[3] = losses[1]/losses[2]
                    else
                        losses[3] = 0.999 * losses[3] + 0.001 * loss1
                    end
                end
            end
        end
    end
    return totalloss / length(sentences)
end

oraclegrad = grad(oracleloss) # less is more :)

# Implement test related

# endofparse to use in test
endofparse(p)=(p.sptr == 1 && p.wptr > p.nword)

function oracletest(allmodel, corpus, arctype, hiddens, batchsize, embdim; pdrop=(0.0, 0.0))
    sentbatches = minibatch(corpus, batchsize)
    hidstack, hidact, hidbuf = hiddens
    mlpmodel = parserv(allmodel)

    for sentences in sentbatches
        parsers = map(arctype, sentences)
        B = length(parsers) # bathcsize
        mcosts  = Array{Cost}(parsers[1].nmove)
        parserdone = map(endofparse, parsers)

        if all(parserdone) # shows how careless they prepared the dataset!
            for p in parsers
                p.sentence.parse = p
            end
            continue
        end

        # cache
        walls = cache_wvecs(allmodel, parsers)
        cachedbufs = cache_bufbatch(allmodel, walls, parsers, hidbuf)

        t0 = zeros(Float32, hidact, B, 1)
        h0 = c0 = (gpu()>=0 ? KnetArray(t0) : Array(t0))
        y0 = reshape(h0, hidact, B)
        totalloss = 0.0
        nstate = (y0, h0, c0)
        inacts = map(t->1, 1:B)

        while !all(parserdone)
            yact, hact, cact = scan_action(allmodel, inacts, nstate)
            nstate = (yact, hact, cact);
            ysts   = scan_stackbatch(allmodel, walls, parsers, hidstack, embdim)
            ybufs  = scan_bufbatch(cachedbufs, parsers, hidbuf)
            encoded = vcat(yact, ysts, ybufs)
            scores  = mlp(mlpmodel, encoded, pdrop=pdrop)
            logprobs = Array(logp(scores, 1))

            inacts = Int[]
            for (i, p) in enumerate(parsers)
                if parserdone[i]
                    if p.sentence.parse == nothing # to avoid from the first time
                        p.sentence.parse = p
                    end
                    push!(inacts, 1) # not to break batchsize
                    continue
                end
                isorted = sortperm(logprobs[:,i], rev=true) # ith col ith inst.
                for m in isorted # best&val action
                    if moveok(p, m)
                        update_cache!(allmodel, walls[i], cachedbufs[i], p, m)
                        move!(p, m)
                        push!(inacts, m) # not to break batchsize
                        break
                    end
                end
                if endofparse(p)
                    parserdone[i] = true
                    p.sentence.parse = p
                end
            end
        end
    end
    return las(corpus)
end


text:
# TODO: put all this into a new DynamicArray library.
# what is wrong with DynamicDenseCPU and DynamicDenseGPU?
# prevent resize from changing ndims

using CUDArt
import Base: isequal, similar, convert, copy, copy!, resize!, issparse
import Base: rand!, randn!, fill!
import CUDArt: to_host

### KUdense parametrized by array type, element type, and ndims:
# TODO: we should get rid of N, because it may change during a resize!

type KUdense{A,T,N}; arr; ptr; end

### CONSTRUCTORS

KUdense(a)=KUdense{atype(a),eltype(a),ndims(a)}(a, reshape(a, length(a)))

KUdense{A<:Array,T,N}(::Type{A}, ::Type{T}, d::NTuple{N,Int})=KUdense(Array(T,d))
KUdense{A<:CudaArray,T,N}(::Type{A}, ::Type{T}, d::NTuple{N,Int})=KUdense(CudaArray(T,d))

convert(::Type{KUdense}, a)=KUdense(a)
convert{A<:BaseArray}(::Type{A}, a::KUdense)=convert(A, a.arr)
convert{A,B}(::Type{KUdense{B}}, a::A)=KUdense(convert(B, a))
convert{A<:Array,B<:CudaArray,T,N}(::Type{KUdense{A,T,N}}, a::KUdense{B,T,N})=cpucopy(a)
convert{A<:CudaArray,B<:Array,T,N}(::Type{KUdense{A,T,N}}, a::KUdense{B,T,N})=gpucopy(a)

similar{A,T,N}(a::KUdense{A}, ::Type{T}, d::NTuple{N,Int})=KUdense(A,T,d)
similar{A,T,N}(a::KUdense{A,T}, d::NTuple{N,Int})=KUdense(A,T,d)
similar{A,T,N}(a::KUdense{A,T,N})=KUdense(A,T,size(a))

arr(a::Vector,d::Dims)=pointer_to_array(pointer(a), d)
arr(a::CudaVector,d::Dims)=CudaArray(a.ptr, d, a.dev)
arr(a)=(isa(a,KUdense) ? a.arr : a)

atype(::Array)=Array
atype(::SubArray)=Array
atype(::CudaArray)=CudaArray

isequal(a::KUdense,b::KUdense)=((typeof(a)==typeof(b)) && (sizeof(a)==sizeof(b)) && isequal(a.arr,b.arr))

### BASIC ARRAY OPS

atype{A}(::KUdense{A})=A

for fname in (:eltype, :length, :ndims, :size, :strides, :pointer, :isempty, :vecnorm)
    @eval (Base.$fname)(a::KUdense)=$fname(a.arr)
end

for fname in (:size, :stride)
    @eval (Base.$fname)(a::KUdense,n)=$fname(a.arr,n)
end

for fname in (:getindex, :setindex!, :sub)
    @eval (Base.$fname)(a::KUdense,n...)=$fname(a.arr,n...)
end

### BASIC COPY

copy!{A,B,T}(a::KUdense{A,T}, b::KUdense{B,T})=(resize!(a, size(b)); copy!(a.arr, 1, b.arr, 1, length(b)); a)
copy!{A,T}(a::KUdense{A,T}, b::Union{Array{T},CudaArray{T}})=(resize!(a, size(b)); copy!(a.arr, 1, b, 1, length(b)); a)
copy!{A,T}(a::Union{Array{T},CudaArray{T}}, b::KUdense{A,T})=(size(a)==size(b)||error(); copy!(a, 1, b.arr, 1, length(b)); a)
copy!{A,B,T}(a::KUdense{A,T}, ai::Integer, b::KUdense{B,T}, bi::Integer, n::Integer)=(copy!(a.arr, ai, b.arr, bi, n); a)
copy{A,T,N}(a::KUdense{A,T,N})=copy!(similar(a), a)

### EFFICIENT RESIZE

# Resize factor: 1.3 ensures a3 can be written where a0+a1 used to be
resizefactor(::Type{KUdense})=1.3

function resize!(a::KUdense, d::Dims)
    size(a)==d && return a
    n = prod(d)
    n > length(a.ptr) && resize!(a.ptr, round(Int,resizefactor(KUdense)*n+1))
    a.arr = arr(a.ptr, d)
    return a
end

resize!(a::KUdense, d::Int...)=resize!(a,d)

# Need to fix deepcopy so it does not create two arrays for arr and ptr:
# And atype changes correctly.

deepcopy_internal(x::KUdense,s::ObjectIdDict)=(haskey(s,x)||(s[x]=copy(x));s[x])
cpucopy_internal(x::KUdense,s::ObjectIdDict)=deepcopy_internal(x,s)
gpucopy_internal(x::KUdense,s::ObjectIdDict)=deepcopy_internal(x,s)

cpucopy_internal{A<:CudaArray}(x::KUdense{A},s::ObjectIdDict)=(haskey(s,x)||(s[x]=KUdense(to_host(x.arr)));s[x])
gpucopy_internal{A<:Array}(x::KUdense{A},s::ObjectIdDict)=(haskey(s,x)||(s[x]=KUdense(CudaArray(x.arr)));s[x])

randn!{A,T}(a::KUdense{A,T}, mean=zero(T), std=one(T))=(randn!(a.arr, mean, std); a)
rand!(a::KUdense)=(rand!(a.arr); a)
fill!{A,T}(a::KUdense{A,T},x)=(fill!(a.arr,convert(T,x)); a)

to_host{A<:CudaArray}(a::KUdense{A})=cpucopy(a)
issparse(a::KUdense)=false

# The final prediction output y should match the input x as closely as
# possible except for being dense.  These functions support obtaining
# the dense version of x.

dtype(x)=typeof(x)
dtype{T}(x::SparseMatrixCSC{T})=Array{T}
dtype{T}(x::CudaSparseMatrixCSC{T})=KUdense{CudaArray,T} # TODO: should this be CudaArray?
dtype{A,T}(x::KUdense{A,T})=KUdense{A,T}  # this is necessary, otherwise 1-dim x does not match 2-dim y.

dsimilar(x,d::Dims)=similar(x,d)
dsimilar{T}(x::SparseMatrixCSC{T},d::Dims)=Array(T,d)
dsimilar{T}(x::CudaSparseMatrixCSC{T},d::Dims)=KUdense(CudaArray,T,d)

function dsimilar!(l, n, x, dims=size(x))
    if (!isdefined(l,n) || !isa(l.(n), dtype(x)))
        l.(n) = dsimilar(x, dims)
    elseif (size(l.(n)) != dims)
        l.(n) = resize!(l.(n), dims)
    end
    return l.(n)
end


text:
# Test Function
function gauss_3D(p::Point{Float64}, p_p::Point{Float64})
	return @. exp(-(p.x)^ 2 - (p.y)^ 2 - (p.z)^ 2 - (p_p.x)^2 - (p_p.y)^2 - (p_p.z)^2)
end

# Jacobian of the translation
function jac_trasl(xl::Point{Float64}, xu::Point{Float64}, xl_p::Point{Float64}, xu_p::Point{Float64})
	return (xu.r-xl.r) * (xu.θ-xl.θ) * (xu.ϕ-xl.ϕ) * (xu_p.r-xl_p.r) * (xu_p.θ-xl_p.θ) * (xu_p.ϕ-xl_p.ϕ)
end

# A naif scalar product between two Point
function inner_prod(p_one::Point{Float64}, p_two::Point{Float64})
	return (p_one.x*p_two.x + p_one.y*p_two.y + p_one.z*p_two.z)
end

# The BEAM WAIST
function ω(p::Point{Float64})
	return (ω₀) * (1+(λ*inner_prod(get_versor(p), p) / (π * ω₀^2)  )^2)^0.5
end

# Atmospheric vertical temperature profile
function Tₚ(p::Point{Float64})
	return Tᵧ * ( 1 - (p.z / z_atm) )
end

# Distribution of turbulent structures on the atmosphere
function χ₁(p::Point{Float64}, p_p::Point{ Float64}, t_s::Float64)
	W        = Point{Float64}(45, 0, (π/2)+(π/(7.0)), undef, undef, undef)()
	W_incr   = Point{Float64}(0.0, 0.0, 0.0, W.x*t_s, W.y*t_s, 0)
	dr = diff_cart(p, p_p)
	dr_mod = ((dr.x-W_incr.x)^2 + (dr.y-W_incr.y)^2 + dr.z^2)^0.5
	return χ₁_0 * exp(- ( (dr_mod)^2 / (2*(L₀)^2) ) )
end

# Water Vapor vertical density profile
function χ₂(p::Point{Float64}, p_p::Point{Float64})
	return χ₂_0 * exp(- (p.z+p_p.z) / (2*z_0) )
end

# Effective area assiaciated to the beam projection in the sky
function B(p::Point{Float64})
	r_mod = (p.x^2 + p.y^2 +p.z^2)^0.5
	first_part = ( 2 * λ^2 * abs(inner_prod(get_versor(p), p))^2 ) / (π * ( ω(p) )^2 )
	#second_part = exp( - (r_mod^2 - (inner_prod(get_versor(p), p))^2) / (ω(p)^2 ) )
	return first_part #* second_part
end

# Jacobiano coordinate sferiche
function Jac_Sp(p::Point{Float64})
	return  cos(p.θ)
end

# Integrand of correlation coef

function corr(p::Point{Float64}, p_p::Point{Float64}, time_stamp::Float64)
 	return B(p) * B(p_p) * χ₁(p, p_p, time_stamp) * χ₂(p, p_p) * Tₚ(p) * Tₚ(p_p) * Jac_Sp(p) * Jac_Sp(p_p)
end

max_stars_repo_name: gruffoc/AFO

max_issues_repo_name: gruffoc/AFO

max_forks_repo_name: gruffoc/AFO


text:
using ConditionalJuMP

export Linear

"""
$(TYPEDEF)

Represents matrix multiplication.

`p(x)` is shorthand for [`matmul(x, p)`](@ref) when `p` is an instance of
`Linear`.

## Fields:
$(FIELDS)
"""
@auto_hash_equals struct Linear{T<:Real, U<:Real} <: Layer
    matrix::Array{T, 2}
    bias::Array{U, 1}

    function Linear{T, U}(matrix::Array{T, 2}, bias::Array{U, 1}) where {T<:Real, U<:Real}
        (matrix_width, matrix_height) = size(matrix)
        bias_height = length(bias)
        @assert(
            matrix_height == bias_height,
            "Number of output channels in matrix, $matrix_height, does not match number of output channels in bias, $bias_height."
        )
        return new(matrix, bias)
    end

end

function Linear(matrix::Array{T, 2}, bias::Array{U, 1}) where {T<:Real, U<:Real}
    Linear{T, U}(matrix, bias)
end

function Base.show(io::IO, p::Linear)
    input_size = size(p.matrix)[1]
    output_size = size(p.matrix)[2]
    print(io,
        "Linear($input_size -> $output_size)"
    )
end

function check_size(params::Linear, sizes::NTuple{2, Int})::Nothing
    check_size(params.matrix, sizes)
    check_size(params.bias, (sizes[end], ))
end

"""
$(SIGNATURES)

Computes the result of pre-multiplying `x` by the transpose of `params.matrix` and adding
`params.bias`.
"""
function matmul(
    x::Array{<:Real, 1},
    params::Linear)
    return transpose(params.matrix)*x .+ params.bias
end

"""
$(SIGNATURES)

Computes the result of pre-multiplying `x` by the transpose of `params.matrix` and adding
`params.bias`. We write the computation out by hand when working with `JuMPLinearType`
so that we are able to simplify the output as the computation is carried out.
"""
function matmul(
    x::Array{T, 1},
    params::Linear{U, V}) where {T<:JuMPLinearType, U<:Real, V<:Real}
    Memento.info(MIPVerify.LOGGER, "Applying $params ... ")
    (matrix_height, matrix_width) = size(params.matrix)
    (input_height, ) = size(x)
    @assert(matrix_height == input_height,
        "Number of values in input, $input_height, does not match number of values, $matrix_height that Linear operates on."
    )
    W = Base.promote_op(+, V, Base.promote_op(*, T, U))
    output = Array{W}(undef, matrix_width)

    for i in 1:matrix_width
        s::W = 0
        for j in 1:matrix_height
            s = increment!(s, x[j], params.matrix[j, i])
        end
        s += params.bias[i]
        ConditionalJuMP.simplify!(s)
        output[i] = s
    end

    return output
end

(p::Linear)(x::Array{<:JuMPReal}) = "Linear() layers work only on one-dimensional input. You likely forgot to add a Flatten() layer before your first linear layer." |> ArgumentError |> throw

(p::Linear)(x::Array{<:JuMPReal, 1}) = matmul(x, p)


text:
#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*-----------------------------------------------------------------------
#* Author: Tiago Os򱨯rio
#* $Id$
#*---------------------------------------------------------------------
type HiLoSelect
	HiLoSelect()=begin
		new(
			DanaSwitcher(Dict{Symbol,Any}(
				:Brief=>"HiLoSelect option",
				:Valid=>["higher","lower"],
				:Default=>"higher"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"input signal 1"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"input signal 2"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"output signal"
			)),
			[
				:(output = max([input1,input2])),
				:(output = min([input1,input2])),
			],
			[
				"Calculate output maximum","Calculate output minimum",
			],
			[:Select,],
			[:input1,:input2,:output,]
		)
	end
	Select::DanaSwitcher
	input1::DanaReal
	input2::DanaReal
	output::DanaReal
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export HiLoSelect
function setEquationFlow(in::HiLoSelect)
	let switch=Select
		if switch=="higher"
			addEquation(1)
		elseif switch=="lower"
			addEquation(2)
		end
	end
end
function atributes(in::HiLoSelect,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Icon]="icon/HiLoSelect"
	fields[:Brief]="Model HiLo Select."
	fields[:Info]="== Inputs ==
* Two different input signals.

== Outputs ==
* One output signal.
"
	drive!(fields,_)
	return fields
end
HiLoSelect(_::Dict{Symbol,Any})=begin
	newModel=HiLoSelect()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
@testset "SS component system" begin
    SS = SourceSensor(name = :SS)

    @test length(freeports(SS)) == 1
    @test numports(SS) == 1
    @test length(parameters(SS)) == 0
    @test length(states(SS)) == 0
    @test length(equations(SS)) == 0
    @test length(constitutive_relations(SS)) == 0

    sys = ODESystem(SS)
    @test length(sys.systems) == 1
    @test sys.p1.E isa Num
    @test sys.p1.F isa Num
end

function find_subsys(sys,s)
    subsys = ModelingToolkit.get_systems(sys)
    return filter(x -> nameof(x) == s, subsys)[1]
end

@testset "Expose models" begin
    r = Component(:R)
    kcl = EqualFlow(name = :kcl)
    SSA = SourceSensor(name = :A)
    SSB = SourceSensor(name = :B)

    bg = BondGraph()
    add_node!(bg, [r, kcl, SSA, SSB])
    connect!(bg, kcl, r)
    connect!(bg, SSA, kcl)
    connect!(bg, kcl, SSB)

    bgn = BondGraphNode(bg)
    @test numports(bgn) == 2

    sys = ODESystem(bgn)
    expanded_sys = expand_connections(sys) # Note that the equations shouldn't simplify by default
    eqns = equations(expanded_sys)

    p1 = find_subsys(sys,:p1)
    p2 = find_subsys(sys,:p2)
    (E1,F1,E2,F2) = (p1.E, p1.F, p2.E, p2.F)

    Asys = find_subsys(sys,:A)
    Bsys = find_subsys(sys,:B)
    (AE,AF,BE,BF) = (Asys.p1.E, Asys.p1.F, Bsys.p1.E, Bsys.p1.F)

    @test (0 ~ E1 - AE) in eqns
    @test (0 ~ F1 + AF) in eqns
    @test (0 ~ E2 - BE) in eqns
    @test (0 ~ F2 + BF) in eqns
end

@testset "Modular RLC circuit" begin
    r = Component(:R)
    l = Component(:I)
    c = Component(:C)
    kvl = EqualEffort(name = :kvl)
    SS1 = SourceSensor(name = :SS1)
    SS2 = SourceSensor(name = :SS2)

    bg1 = BondGraph(:RC)
    add_node!(bg1, [r, c, kvl, SS1])
    connect!(bg1, r, kvl)
    connect!(bg1, c, kvl)
    connect!(bg1, SS1, kvl)
    bgn1 = BondGraphNode(bg1)

    bg2 = BondGraph(:L)
    add_node!(bg2, [l, SS2])
    connect!(bg2, l, SS2)
    bgn2 = BondGraphNode(bg2)

    bg = BondGraph()
    add_node!(bg, [bgn1, bgn2])
    connect!(bg, bgn1, bgn2)

    eqs = constitutive_relations(bg)
    @test length(eqs) == 2

    sys = ODESystem(bg)
    (R, C, L) = sys.ps
    (qC, pL) = sys.states
    e1 = D(qC) ~ -pL / L + (-qC / C / R)
    e2 = D(pL) ~ qC / C

    @test isequal(eqs[1].lhs, e1.lhs)
    @test isequal(simplify(eqs[1].rhs-e1.rhs), 0)
    @test isequal(eqs[2].lhs, e2.lhs)
    @test isequal(eqs[2].rhs, e2.rhs)
end

@testset "Modular reaction" begin
    bg1 = BondGraph(:R)
    re = Component(:re, :r)
    SSA = SourceSensor(name = :A)
    SSB = SourceSensor(name = :B)
    add_node!(bg1, [SSA, SSB, re])
    connect!(bg1, SSA, re; dstportindex = 1)
    connect!(bg1, re, SSB; srcportindex = 2)
    bgn1 = BondGraphNode(bg1)

    bg = BondGraph()
    A = Component(:ce, :A)
    B = Component(:ce, :B)
    add_node!(bg, [A, B, bgn1])
    connect!(bg, A, bgn1; dstportindex = 1)
    connect!(bg, bgn1, B; srcportindex = 2)

    sys = ODESystem(bg)
    eqs = constitutive_relations(bg)

    (xA, xB) = sys.states
    (KA, KB, r) = sys.ps
    e1 = D(xA) ~ -r * (KA * xA - KB * xB)
    e2 = D(xB) ~ r * (KA * xA - KB * xB)

    @test isequal(eqs[1].lhs, e1.lhs)
    @test isequal(eqs[1].rhs, e1.rhs)
    @test isequal(eqs[2].lhs, e2.lhs)
    @test isequal(eqs[2].rhs, e2.rhs)
end

max_stars_repo_name: mic-pan/BondGraphs

max_issues_repo_name: mic-pan/BondGraphs

max_forks_repo_name: mic-pan/BondGraphs


text:
using FluidScripts
using Base.Test

# write your own tests here


text:
immutable ExponentialSampler <: Sampleable{Univariate,Continuous}
    scale::Float64
end

rand(s::ExponentialSampler) = s.scale * randexp()

immutable ExponentialLogUSampler <: Sampleable{Univariate,Continuous}
    scale::Float64
end

rand(s::ExponentialLogUSampler) = -s.scale * log(rand())


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMCoupling.jl/blob/master/LICENSE
function get_C(refnode,slaves,dofs,ndofs,K_size)

    C_rows=size(slaves,1)*size(dofs,1)
    CC=spzeros(C_rows,K_size)
    g=spzeros(C_rows,1)
    D=spzeros(C_rows,C_rows)

    for i in 1:length(dofs),
        j in 0:length(slaves)-1
        CC[i+j*length(dofs) , refnode*ndofs-ndofs+dofs[i]]=1
        CC[i+j*length(dofs) , slaves[j+1]*ndofs-ndofs+dofs[i]]=-1
    end
    return CC,D,g
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

struct SnapshotsApi <: SwaggerApi
    client::Swagger.Client
end

"""

Creates or updates a snapshot.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: snapshotName::String (required)
Param: api_version::String (required)
Param: snapshot::Snapshot (required)
Return: Snapshot
"""
function snapshotsCreateOrUpdate(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, snapshotName::String, api_version::String, snapshot; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "PUT", Snapshot, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots/{snapshotName}", ["azure_auth"], snapshot)
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "snapshotName", snapshotName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Deletes a snapshot.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: snapshotName::String (required)
Param: api_version::String (required)
Return: OperationStatusResponse
"""
function snapshotsDelete(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, snapshotName::String, api_version::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "DELETE", OperationStatusResponse, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots/{snapshotName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "snapshotName", snapshotName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets information about a snapshot.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: snapshotName::String (required)
Param: api_version::String (required)
Return: Snapshot
"""
function snapshotsGet(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, snapshotName::String, api_version::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", Snapshot, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots/{snapshotName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "snapshotName", snapshotName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Grants access to a snapshot.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: snapshotName::String (required)
Param: api_version::String (required)
Param: grantAccessData::GrantAccessData (required)
Return: AccessUri
"""
function snapshotsGrantAccess(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, snapshotName::String, api_version::String, grantAccessData; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", AccessUri, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots/{snapshotName}/beginGetAccess", ["azure_auth"], grantAccessData)
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "snapshotName", snapshotName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Lists snapshots under a subscription.
Param: subscriptionId::String (required)
Param: api_version::String (required)
Return: SnapshotList
"""
function snapshotsList(_api::SnapshotsApi, subscriptionId::String, api_version::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", SnapshotList, "/subscriptions/{subscriptionId}/providers/Microsoft.Compute/snapshots", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Lists snapshots under a resource group.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: api_version::String (required)
Return: SnapshotList
"""
function snapshotsListByResourceGroup(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, api_version::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", SnapshotList, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Revokes access to a snapshot.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: snapshotName::String (required)
Param: api_version::String (required)
Return: OperationStatusResponse
"""
function snapshotsRevokeAccess(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, snapshotName::String, api_version::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", OperationStatusResponse, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots/{snapshotName}/endGetAccess", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "snapshotName", snapshotName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Updates (patches) a snapshot.
Param: subscriptionId::String (required)
Param: resourceGroupName::String (required)
Param: snapshotName::String (required)
Param: api_version::String (required)
Param: snapshot::SnapshotUpdate (required)
Return: Snapshot
"""
function snapshotsUpdate(_api::SnapshotsApi, subscriptionId::String, resourceGroupName::String, snapshotName::String, api_version::String, snapshot; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "PATCH", Snapshot, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/snapshots/{snapshotName}", ["azure_auth"], snapshot)
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "snapshotName", snapshotName)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

export snapshotsCreateOrUpdate, snapshotsDelete, snapshotsGet, snapshotsGrantAccess, snapshotsList, snapshotsListByResourceGroup, snapshotsRevokeAccess, snapshotsUpdate


text:
"""
    set_title!(axis, t)

Internal function to set the title of an Axis object.
"""
set_title!(a::Union{Axes,Axis}, t::String) = (a.title = Title(text=t))

"""
    set_title!(axis, s)

Internal function to set the base scale of an Axis object relative
to the parent font size.
"""
set_base!(a::Axis, s::Float64) = (a.base = s)

"""
    set_min!(axis, m)

Internal function to set the minimum value of the `axis`. Note that
`[x|y|x2|y2]lim!` is preferred.
"""
set_min!(a::Axis, m::Float64) = (a.min = m)

"""
    set_max!(axis, m)

Internal function to set the maximum value of the `axis`. Note that
`[x|y|x2|y2]lim!` is preferred.
"""
set_max!(a::Axis, m::Float64) = (a.max = m)

"""
    set_log!(axis, b)

Internal function to toggle the log option of the `axis` object.
Note that `[x|y|x2|y2]scale!` is peferred.
"""
set_log!(a::Axis, b::Bool) = (a.log = b)


text:
function acronym(phrase)
  uppercase(join(first.(filter.(isletter, (split(replace(phrase, '-' => " ")))))))
end

max_stars_repo_name: Rigante-pl/Exercism-Julia

max_issues_repo_name: Rigante-pl/Exercism-Julia

max_forks_repo_name: Rigante-pl/Exercism-Julia


text:
uval1 = UncertainValue(Normal, 1, 5)
uval2 = UncertainValue(Uniform, rand(1000))
uval3 = UncertainValue(rand(10000))

v = [uval1, uval2, uval3]

@test minimum(uval1) isa Float64
@test minimum(uval2) isa Float64
@test minimum(uval3) isa Float64

@test maximum(uval1) isa Float64
@test maximum(uval2) isa Float64
@test maximum(uval3) isa Float64

@test minimum(v) isa Float64
@test maximum(v) isa Float64


text:
"""
Structs and functions for implementing various pair interaction energies
 and fields, given a specific lattice to operate on.
Upon creation of a SpinSystem, all pair interactions get converted into their
 corresponding type here.
"""

"""
    BondTable{T}

Stores all of the bonds and associated data (e.g. exchange scalars/matrices) of type T
contained within a single pair interaction densely in memory, with quick access to bonds/Ts
on any sublattice.
"""
struct BondTable{T}
    bonds         :: Vector{Bond}      # All bonds, sorted on first basis index
    culled_bonds  :: Vector{Bond}      # Culled to minimal 1/2 bonds, sorted on first basis index
    data          :: Vector{T}         # Data T associated with each bond in `bonds`
    culled_data   :: Vector{T}         # Culled data T associated with each bond in `bonds`
    basis_indices :: Vector{Int}       # Indices to first bond of each basis index in `bonds`
end

"""Construct a `BondTable` given a `Crystal`, and a `Bond` along with its associated
    exchange matrix. Performs the symmetry transformations to obtain all other bonds
    and exchanges matrices.
"""
function BondTable(crystal::Crystal, bond::Bond, J::Mat3)
    bonds = Vector{Bond}()
    Js = Vector{Mat3}()
    basis_indices = [1]

    for i in 1:nbasis(crystal)
        (symbonds, symJs) = all_symmetry_related_couplings_for_atom(crystal, i, bond, J)
        append!(bonds, symbonds)
        append!(Js, symJs)
        push!(basis_indices, last(basis_indices) + length(symbonds))
    end

    # Compute the minimal set of the bonds/Js needed to specify all bonds
    # Bonds with i ↔ j and n ↔ -n are equivalent: only keep one of each pair
    culled_bonds = Vector{Bond}()
    culled_Js = Vector{Mat3}()
    for (bond, J) in zip(bonds, Js)
        neg_bond = Bond(bond.j, bond.i, -bond.n)
        if !(neg_bond in culled_bonds)
            push!(culled_bonds, bond)
            push!(culled_Js, J)
        end
    end

    BondTable{Mat3}(bonds, culled_bonds, Js, culled_Js, basis_indices)
end

function Base.show(io::IO, ::MIME"text/plain", bondtable::BondTable{T}) where {T}
    print(io, "$(length(bondtable))-element, $(length(bondtable.basis_indices)-1)-basis BondTable{$T}")
end

Base.length(bondtable::BondTable) = length(bondtable.bonds)
nbasis(bondtable::BondTable) = length(bondtable.basis_indices) - 1
## These functions generate iterators producing (bond, data) for various subsets of bonds ##
all_bonds(bondtable::BondTable) = zip(bondtable.bonds, bondtable.data)
culled_bonds(bondtable::BondTable) = zip(bondtable.culled_bonds, bondtable.culled_data)
function sublat_bonds(bondtable::BondTable, i::Int)
    @unpack bonds, data, basis_indices = bondtable
    @boundscheck checkindex(Bool, 1:length(basis_indices)-1, i) ? nothing : throw(BoundsError(bondtable, i))
    zip(bonds[basis_indices[i]:basis_indices[i+1]-1], data[basis_indices[i]:basis_indices[i+1]-1])
end

# This method `map`'s only the data within a BondTable, not the bonds themselves.
function Base.map(f, bondtable::BondTable)
    @unpack bonds, culled_bonds, data, culled_data, basis_indices = bondtable
    new_data = map(f, data)
    new_culled_data = map(f, culled_data)
    return BondTable{eltype(new_data)}(
        bonds, culled_bonds, new_data, new_culled_data, basis_indices
    )
end

abstract type AbstractPairIntCPU <: AbstractInteractionCPU end

"""
    HeisenbergCPU <: AbstractPairIntCPU

Implements an exchange interaction which is proportional to
the identity matrix.
"""
struct HeisenbergCPU <: AbstractPairIntCPU
    bondtable    :: BondTable{Float64}    # Bonds store effective J's = S_i J S_j (all identical)
    label        :: String
end

"""
    DiagonalCouplingCPU <: AbstractPairIntCPU

Implements an exchange interaction where matrices on all bonds
are diagonal.
"""
struct DiagonalCouplingCPU <: AbstractPairIntCPU
    bondtable     :: BondTable{Vec3}    # Bonds store diagonal of effective J's = S_i J S_j
    label         :: String
end

"""
    GeneralCouplingCPU <: AbstractPairIntCPU

Implements the most generalized interaction, where matrices on
all bonds are full 3x3 matrices which vary bond-to-bond.
"""
struct GeneralCouplingCPU <: AbstractPairIntCPU
    bondtable    :: BondTable{Mat3}    # Bonds store effective J's = S_i J S_j
    label        :: String
end

# Helper functions producing predicates checking if a matrix is approximately
# Heisenberg or diagonal
isheisen(J, tol) = isapprox(diagm(fill(J[1,1], 3)), J; atol=tol)
isdiag(J, tol) = isapprox(diagm(diag(J)), J; atol=tol)
isheisen(tol) = Base.Fix2(isheisen, tol)
isdiag(tol) = Base.Fix2(isdiag, tol)

# Figures out the correct maximally-efficient backend type for a quadratic interaction, and perform
#  all of the symmetry propagation.
function convert_quadratic(int::QuadraticInteraction, cryst::Crystal, sites_info::Vector{SiteInfo}; tol=1e-6)
    @unpack J, bond, label = int

    # Product S_i S_j between sites connected by the bond -- same for all bonds in a class
    # To get interactions quadratic in the unit vectors stored by `SpinSystem`, we internally
    #  store effective exchange matrices (Si J Sj).
    SiSj = sites_info[bond.i].S * sites_info[bond.j].S
    J *= SiSj
    bondtable = BondTable(cryst, bond, J)

    if all(isheisen(tol), bondtable.culled_data)
        # Pull out the Heisenberg scalar from each exchange matrix
        bondtable = map(J->J[1,1], bondtable)
        return HeisenbergCPU(bondtable, label)
    elseif all(isdiag(tol), bondtable.culled_data)
        # Pull out the diagonal of each exchange matrix
        bondtable = map(J->diag(J), bondtable)
        return DiagonalCouplingCPU(bondtable, label)
    else
        return GeneralCouplingCPU(bondtable, label)
    end
end

function energy(spins::Array{Vec3}, heisenberg::HeisenbergCPU)
    bondtable = heisenberg.bondtable
    effJ = first(bondtable.data)
    E = 0.0

    latsize = size(spins)[2:end]
    for (bond, _) in culled_bonds(bondtable)
        @unpack i, j, n = bond
        for cell in CartesianIndices(latsize)
            sᵢ = spins[i, cell]
            sⱼ = spins[j, offset(cell, n, latsize)]
            E += sᵢ ⋅ sⱼ
        end
    end
    return effJ * E
end

function energy(spins::Array{Vec3}, diag_coup::DiagonalCouplingCPU)
    bondtable = diag_coup.bondtable
    E = 0.0

    latsize = size(spins)[2:end]
    for (bond, J) in culled_bonds(bondtable)
        @unpack i, j, n = bond
        for cell in CartesianIndices(latsize)
            sᵢ = spins[i, cell]
            sⱼ = spins[j, offset(cell, n, latsize)]
            E += (J .* sᵢ) ⋅ sⱼ
        end
    end
    return E
end

function energy(spins::Array{Vec3}, gen_coup::GeneralCouplingCPU)
    bondtable = gen_coup.bondtable
    E = 0.0

    latsize = size(spins)[2:end]
    for (bond, J) in culled_bonds(bondtable)
        @unpack i, j, n = bond
        for cell in CartesianIndices(latsize)
            sᵢ = spins[i, cell]
            sⱼ = spins[j, offset(cell, n, latsize)]
            E += dot(sᵢ, J, sⱼ)
        end
    end
    return E
end

"Accumulates the local -∇ℋ coming from Heisenberg couplings into `B`"
@inline function _accum_neggrad!(B::Array{Vec3}, spins::Array{Vec3}, heisen::HeisenbergCPU)
    bondtable = heisen.bondtable
    effJ = first(bondtable.data)

    latsize = size(spins)[2:end]
    for (bond, _) in culled_bonds(bondtable)
        @unpack i, j, n = bond
        for cell in CartesianIndices(latsize)
            offsetcell = offset(cell, n, latsize)
            B[i, cell] = B[i, cell] - effJ * spins[j, offsetcell]
            B[j, offsetcell] = B[j, offsetcell] - effJ * spins[i, cell]
        end
    end
end

"Accumulates the local -∇ℋ coming from diagonal couplings into `B`"
@inline function _accum_neggrad!(B::Array{Vec3}, spins::Array{Vec3}, diag_coup::DiagonalCouplingCPU)
    bondtable = diag_coup.bondtable

    latsize = size(spins)[2:end]
    for (bond, J) in culled_bonds(bondtable)
        @unpack i, j, n = bond
        for cell in CartesianIndices(latsize)
            offsetcell = offset(cell, n, latsize)
            B[i, cell] = B[i, cell] - J .* spins[j, offsetcell]
            B[j, offsetcell] = B[j, offsetcell] - J .* spins[i, cell]
        end
    end
end

"Accumulates the local -∇ℋ coming from general couplings into `B`"
@inline function _accum_neggrad!(B::Array{Vec3}, spins::Array{Vec3}, gen_coup::GeneralCouplingCPU)
    bondtable = gen_coup.bondtable

    latsize = size(spins)[2:end]
    for (bond, J) in culled_bonds(bondtable)
        @unpack i, j, n = bond
        for cell in CartesianIndices(latsize)
            offsetcell = offset(cell, n, latsize)
            B[i, cell] = B[i, cell] - J * spins[j, offsetcell]
            B[j, offsetcell] = B[j, offsetcell] - J * spins[i, cell]
        end
    end
end


text:
#testing for julia functions :)

using Statistics
using LinearAlgebra
#using InvertedIndices
#operations:
#const LiefiesMatrix = Union{Matrix{Missing}, Matrix{<:Real}}
LiefiesMatrix{T <: Union{Missing, Real}} = AbstractMatrix{T}
#function f=some_function(x::LiefiesMatrix) :D

# function closure(data::Union{Array{Union{Float64,Missing},2},Array{Union{Int64, Missing},2},Vector{Float64},Vector{Int64}}) #Union{Vector{Regex}, NTuple{N, Regex}}
#::LiefiesMatrix) #Union{Vector{Regex}, NTuple{N, Regex}}
function closure(data::T) where T
    cvec = Array{Union{Missing, Float64}}(undef, size(data)...)
    for i in 1:size(data,1)
      cvec[i,:] = data[i,:] ./ sum(skipmissing(data[i,:])) #work around for sum not accepting skipmissing with dims argument
    end
    return cvec
end

function center(data::T) where T #inherit type T from data :D
    #data must be compositions!
    gvec = Matrix{Union{Missing, Float64}}(undef, 1, size(data,2))
    data = closure(data) #just in case
    data = BigFloat.(data)
    for j in 1:size(gvec,2) #by columns
      gvec[j] = prod(skipmissing((data[:,j]))) ^ (1/size(collect(skipmissing(data[:,j])),1)) #work around for sum not accepting skipmissing with dims argument
    end
    gvec=closure(gvec)
    return gvec
end

function ait_var(data::T) where T
    cen = center(data)
    avar = zeros(Float64, size(data,2), size(data,2))
    #avar = Matrix{Float64}(undef, size(data,2), size(data,2))
    #avar = Matrix{Union{Missing, Float64}}(undef, size(data,2), size(data,2))
    for i=1:size(data,2)
        for j=i:size(data,2)
            n=min(length(collect(skipmissing(data[:,i]))),length(collect(skipmissing(data[:,j]))))
            avar[i,j]=(1/(n-1))*sum(skipmissing((log.(data[:,i]./data[:,j]) .- log(cen[i]/cen[j])).^2))
        end
    end
    #avar = Symmetric(avar) #to make symmetric :)
    return avar
end

function geo(x)
    geo_mean = prod(x) .^ (1 ./ length(x))
    return geo_mean
end

function ait_inner(x,y)
    ai=sum( log.(x ./ geo(x)) .* log.(y ./ geo(y))   )
    return ai

    # ai=0.0        #alternative calculation:
    # for i=1:length(x)
    #     ai = ai + sum((log.(x[i] ./ x) .* (log.(y[i] ./ y))))
    # end
    # ai = 1/(2*length(x)) * ai
    #return ai
end

function ait_norm(x)
    an=0.0
    for i=1:length(x)
        an= an + sum((log.(x[i] ./ x) .^2))
    end
    an = sqrt(1/(2*length(x)) * an)
    return an
end

function square_ait_norm(x)
    an=0.0
    for i=1:length(x)
        an= an + sum((log.(x[i] ./ x) .^2))
    end
    an = (1/length(x)) * 1/(2*length(x)) * an
    return an
end

function ait_dist(x,y)
    if length(x) == length(y)
        ad = 0.0
        for i=1:length(x)
            ad = ad + sum((log.(x[i] ./ x) .- log.(y[i] ./ y)).^2)
        end
        ad = sqrt(1/(2*length(x)) * ad)
        return ad
    else
        println("x and y must be the same length")
    end
end

function perturb(x,y)
    p=closure(x) .* closure(y) #works even if x is mulitple rows!
    p=closure(p)
    return p
end

function perturb_diff(x,y)
    p=closure(x) .* inverse(y) #works even if x is mulitple rows!
    p=closure(p)
    return p
end

function power(x,a)
    p=closure(x) .^ a
    p=closure(p)
    return p
end

function inverse(x)
    ix = closure(1 ./ x)
    return ix
end

function clr(data)
    cvec = Matrix{Union{Missing, Float64}}(undef, size(data))
    for i=1:size(data,1) #by rows
        gmean = prod(skipmissing(data[i,:])) ^ (1/size(collect(skipmissing(data[i,:])),1)) #work around for sum not accepting skipmissing with dims argument
        cvec[i,:]=log.(data[i,:]./gmean)
    end
    return cvec
end

function ilr(x,sbp::Matrix{<:Real})

    bal=zeros(size(sbp)) #create orthonormal basis based on partition
    for i=1:size(sbp,1) #number of rows of sbp
        rr=findall(sbp[i,:] .== 1)
        ss=findall(sbp[i,:] .== -1)
        r=length(rr)
        s=length(ss)
        bal[i,rr] .= (1/r).*sqrt((r*s)/(r+s))
        bal[i,ss] .= -(1/s).*sqrt((r*s)/(r+s))
    end

    ilr_x=log.(x)*bal'
    return ilr_x, bal
end

function invilr(ilr_x,bal)
    #balance refers to contrast matrix (it is not the sbp!)
    x=invclr(ilr_x*bal);
    return x
end

function invclr(x_clr)
    x=closure(exp.(x_clr));
    return x
end

function amalga(x,col1,col2)

    A = zeros(size(x,2),3)
    A[col1,1] = 1
    A[col2,2] = 1
    A[setdiff(collect(1:size(x,2)),[col1;col2]),3] .= 1

    y = x * A
    return y, A

end

function bayes_multi_zero_replace(c,method,prior)
    #c is a count vector!
    #tj is the prior estimate
    #s is the strength parameter

    samples, D = size(c) #number of rows then columns
    n=sum(c,dims=2)

    if prior == "uniform"
        t=1/D .* ones(size(c))
    elseif prior == "modified"
        println("Assuming prior based on ratios found within dataset")
        t=zeros(size(c))
        for q=1:samples
            global t
            to_use = setdiff(1:samples, q)
            t[q,:]=sum(c[to_use,:],dims=1)
        end
        t=t./sum(t,dims=2)
    end

    #common imprecise Dirichlet models considered here: Bayes-Laplace and Square-root
    if method == "Bayes-Laplace"
        s = D * ones(samples,1)
    elseif method == "Square-root"
        s = sqrt.(sum(c,dims=2)) #total trials per sample!
    elseif method == "Geometric"
        #if these are very small numbers, can result in a zero multiplicative product
        println("accounting for very small values here...")
        temp = 10 .^ ((1/D) .* sum(log10.(t),dims=2))
        s = 1 ./ temp
        #s = 1 ./ prod(closure(t),dims=2).^(1/D)
    elseif method == "Jeffreys"
        s = D/2 * ones(samples,1)
    elseif method == "Perks"
        s = ones(samples,1)
    else
        println("Not a recognized method; please choose from Bayes-Laplace, Square-root, Geomteric, Jeffreys, or Perks")
    end

    println(method)
    #@show s
    #multiplicative part: only adjust rows for which there is a zero
    #fieldnames()...
    x=closure(c)
    r=deepcopy(x)
    global r, x, t, s
    for i=1:samples
        if any(x[i,:] .== 0) #row contains zeros
            zeroind=findall(x[i,:] .== 0) #work on the columns
            nonzero=findall(x[i,:] .!= 0)
            r[i,zeroind] = t[i,zeroind] .* s[i]./(n[i] .+ s[i])
            r[i,nonzero] = x[i,nonzero] .* (1 - sum(r[i,zeroind]))
        end
    end

    return r
    # x=closure(c)
    # r=x
    # zeroind=findall(x .== 0)
    # nonzero=findall(x .!= 0)
    # colind=map(x->x.I[2], zeroind)
    # rowind=map(x->x.I[1], zeroind)
    #
    # r[zeroind] = t[zeroind] .* s[colind]./(n[colind] .+ s[colind])
    # r[nonzero] = x[nonzero] .* (1 - sum(r[zeroind]))

end

#To compare with R package- that one seems to be doing strange thins with BDL values
# g=Matrix{Union{Missing, Float64}}(undef, 1, size(y,2))
# for j in 1:size(y,2) #by columns
#   g[j] = prod(skipmissing(y[:,j])) ^ (1/size(collect(skipmissing(y[:,j])),1)) #work around for sum not accepting skipmissing with dims argument
# end
#
# #(0.3 * x).^(1/2) = nrm
# g0=sum(g[1,[1,2,3]])
# #sum(g[1,[1,2,4]]) + q = nrm
# #nrm = y[1,4] + g0
#
# r_cen=[0.5116703 0.2558351 0.1193897 0.1131049]
# mm=((g[1]./r_cen[1] - g0).^2)./y[1,4]


text:
function binomial(n::Int,m::Int)
	# \Choose{n}{m}
	ans=1
	for i=1:m
		ans *=(n+1-i)
		ans/=i
		ans=Int(ans)
	end
	return ans

end
"""
# test
for i=0:10
println(binomial(10,i))
end
"""


text:
## Contains functions and objects for statistical analysis
using Statistics, StatsBase

TABLE_HEADER = ["Statistic", "N", "Mean", "Min", "25 Percentile", "Median", "75 Percentile", "Max", "Standard Deviation", "Skewness", "Kurtosis"]

"""
computes the summary statistics of a given data set
"""
function summary_data(name, data)
        n = length(data)
        μ, σ = mean_and_std(data)
        minim, q1, med, q3, maxim = StatsBase.nquantile(data, 4)
        sk = skewness(data)
        ku = kurtosis(data)
        return [name n μ minim q1 med q3 maxim σ sk ku]
end

max_stars_repo_name: zborffs/Delta

max_issues_repo_name: zborffs/Delta

max_forks_repo_name: zborffs/Delta


text:
"""

Predict and plot out vaccination impact in a county given initial condition and vaccination parameter inputs
"""

function predict_vacc_county_new(name::String,linelist_data_with_pos_neg,serology_data,
                    M_Kenya_ho,M_Kenya_other,M_Kenya_school,M_Kenya_work,N_kenya;
                    PCR_array,relative_testing_rate,
                    rapid, scenario,
                    prob_sus_w,
                    prob_sus,
                    prob_symp,
                    H_rate_by_age,
                    p_ID,p_IH, all_IFR_fits,ICU_fit,hosp_fit,
                    deaths_data,
                    priors,
                    T_to_full_efficacy,
                    vacc_end_date = Date(2022,6,30), # date vaccination program ends
                    startdate = Date(2020,12,31),
                    enddate = Date(2022,6,30),
                    obs_factor)

    fit_dict = load("modelfits/$(name)_model.jld2")
    model = fit_dict[first(keys(fit_dict))]
    up_county = uppercase(model.areaname)

    #Set times, all times after model_start_time are in reference to the model_start_time
    model_start_time = model.model_start_time # start in 1st Dec 2020
    model_end_time = model.model_end_time   # end of model run for fitting
    marchendpoint = model.marchendpoint # 10th March when roudn 3 serosurvey data ends
    alpha_variant_time = model.alpha_variant_time  # earliest time of alpha variant introduction
    delta_variant_time = model.delta_variant_time # earliest time of delta variant introduction
    janstarttime = model.janstarttime  # time t start including observations in likelihood
    vacc_end_point =  (vacc_end_date - model.startdate).value # End of model run
    pre_vacc_period = (Date(2021,3,6) - model.startdate).value # vaccination starts on 6th March 2021
    model_end_point =  (enddate - startdate).value # End of model run

    # Variants data
    kenya_variants = CSV.File("data/kenya_gisaid_variants.csv") |> DataFrame
    kenya_variants.date = [Date(d,DateFormat("dd/mm/yyyy")) for d in kenya_variants.date]
    kenya_variants.rel_time = [(d - Date(2020,2,20)).value for d in kenya_variants.date]

    idxs_alpha = kenya_variants.variant .== "Alpha"
    idxs_beta = kenya_variants.variant .== "Beta"
    idxs_delta = kenya_variants.variant .== "Delta"

    #Daily prediction of variant prevalence using linear interpolation
    interp_alpha = DataInterpolations.LinearInterpolation(kenya_variants.perc_sequences[idxs_alpha]./100,kenya_variants.rel_time[idxs_alpha])
    alpha_perc = [interp_alpha(t) for t in 1:maximum(kenya_variants.rel_time[idxs_alpha])]
    interp_beta = DataInterpolations.LinearInterpolation(kenya_variants.perc_sequences[idxs_beta]./100,kenya_variants.rel_time[idxs_beta])
    beta_perc = [interp_beta(t) for t in 1:maximum(kenya_variants.rel_time[idxs_beta])]
    interp_delta = DataInterpolations.LinearInterpolation(kenya_variants.perc_sequences[idxs_delta]./100,kenya_variants.rel_time[idxs_delta])
    delta_perc = [interp_delta(t) for t in 1:maximum(kenya_variants.rel_time[idxs_delta])]

    delta_perc = delta_perc[model_start_time:end]  # pick data from 1st Dec 2020
    alpha_perc = alpha_perc[model_start_time:end]
    beta_perc  = beta_perc[model_start_time:end]

    n = model_end_point

    delta_perc = [delta_perc;ones(n-length(delta_perc))]
    alpha_beta_perc = alpha_perc .+ beta_perc
    alpha_beta_perc = [alpha_beta_perc;zeros(n-length(alpha_beta_perc))]
    other_perc = 1.0 .- alpha_beta_perc .- delta_perc

    # fitted parameters for county
    MCMCdraws = get(model.MCMC_results.chain,[:β₀,:β_home,:β_school,:β_other,:β_work,:ϵ,:χ,:p_test,:E₀,:inc_R_αvar,:time_scale_αvar,:mid_point_αvar,:inc_R_δvar,:time_scale_δvar,:mid_point_δvar,:init_scale])
    prior = priors[countynames.==model.areaname][1]

    # Data sets to populate and save after loop
    crit_dis_inc_samples  = zeros(model_end_point,6,2000) # 6 age groups, 2000 samples from HMC
    severe_dis_inc_samples = zeros(model_end_point,6,2000)
    asymp_dis_inc_samples = zeros(model_end_point,6,2000)
    mild_dis_inc_samples = zeros(model_end_point,6,2000)
    death_inc_samples = zeros(model_end_point,6,2000)

    for k in 1:2000
        α= model.α  # 1/mean latent period
        ϵ= MCMCdraws.ϵ[k] #0.5    # Relative infectiousness of asymptomatic cases
        αP= model.αP # rate of movement from Pre symptomatic state
        γA = model.γA #Recovery rate from Asymptomatic state
        γM= model.γM  #Recovery rate from Mild disease

        relative_testing_rate = KenyaCoVaccines.relative_testing_rate
        PCR_array= KenyaCoVaccines.PCR_array #Relative Sensitivity of PCR test on days post-infection
        PCR_sensitivity = model.PCR_sensitivity #Base sensitivity for RT-PCR test
        PCR_specificity = model.PCR_specificity #Base specificity for RT_PCR test
        baseline_sero_array = KenyaCoVaccines.rel_sero_array_26days #Relative sensitivity of serology test on days post-infection
        sero_sensitivity = model.sero_sensitivity#Base sensitivity for serology test
        sero_specificity= model.sero_specificity #Base specificity for serology test
        M_BB= model.M_BB

        ω= model.ω # Waning of natural protection, rate per day
        σω= model.p_sus_w # relative suceptibility after first episode: 0.16 for all ages for now.
        σ = model.p_sus
        δ = model.p_symp
        υ = model.hosp_rate_by_age

        β₀ = MCMCdraws.β₀[k]
        β_home = MCMCdraws.β_home[k]
        β_school = MCMCdraws.β_school[k]
        β_other = MCMCdraws.β_other[k]
        β_work = MCMCdraws.β_work[k]

        #Set variance scalers
        clustering_factor_PCR = 0.5
        M_PCR = 30.
        T = eltype(β₀)

        inc_R_αvar= MCMCdraws.inc_R_αvar[k]
        time_scale_αvar= MCMCdraws.time_scale_αvar[k]
        mid_point_αvar= MCMCdraws.mid_point_αvar[k]
        inc_R_δvar= MCMCdraws.inc_R_δvar[k]
        time_scale_δvar= MCMCdraws.time_scale_δvar[k]
        mid_point_δvar= MCMCdraws.mid_point_δvar[k]
        init_scale = MCMCdraws.init_scale[k]

        p = convert.(T,[β₀,β_home,0.0,β_other,β_work,α,ϵ,αP,γA,γM,ω,inc_R_αvar,time_scale_αvar,alpha_variant_time + mid_point_αvar*30,inc_R_δvar,time_scale_δvar, delta_variant_time + mid_point_δvar*30,
         init_scale])  # start with schools closed
        p = vcat(p,model.VE_acquisition,model.VE_infectiousness,model.VE_severe_disease_risk) # add vaccine efficacies to the paramter vector

        init_sero = model.init_sero
        γV = γM

        N= model.N
        E₀ = MCMCdraws.E₀[k]

        θ = (β₀=β₀,β_home=β_home,β_school=β_school,β_other=β_other,β_work=β_work,E₀=E₀,ϵ=ϵ,init_scale=init_scale)
        u0 = KenyaCoVaccines.create_initial_conditions(θ,model)

        #Set up the callback for lockdown
        schoolsclosure_times = [Float64((Date(2021, 3, 19) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 7, 16) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 10, 1) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 12, 23) - Date(2020, 12, 1)).value)]
        function affect_schoolsclosed!(integrator)
            integrator.p[3] = 0.0 #Schools shut
        end
        schoolsclosed_cb = PresetTimeCallback(schoolsclosure_times, affect_schoolsclosed!,save_positions=(false,false))

        #Set up callback for reopening Schools
        schoolsreopening_times = [Float64((Date(2021, 1, 4) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 5, 10) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 7, 26) - Date(2020, 12, 1)).value),
            Float64((Date(2021, 10, 11) - Date(2020, 12, 1)).value)]
        function affect_schoolsopening!(integrator)
            integrator.p[3] = β_school # Schools open #
        end
        schoolsreopening_cb = PresetTimeCallback(schoolsreopening_times, affect_schoolsopening!,save_positions=(false,false))

        cb = CallbackSet(schoolsclosed_cb, schoolsreopening_cb)

        immune_escape_variant_time = Float64((Date(2022,1,1) - model.startdate).value) #  December 1 2021, time to implement variant immunity escape

        function immune_escape_variant!(integrator)
            # Reduce vaccine efficacies by 50%
            integrator.p[19:27] *= 0.5

            # Start with 'variant_E' exposed number of individuals
            variant_E = 1
            prop_vacc = sum(integrator.u[:,1:8,:],dims=2)[:,1,:] ./ sum(integrator.u[:,1:8,:],dims=[2,3])[:]
            integrator.u[:,2,:] .= variant_E .* prop_vacc # number exposed in each age group and vaccinatin status
            integrator.u[:,1,:] .= integrator.u[:,1,:] .- integrator.u[:,2,:] #substract E from S

            # move 50% from R and W -> S
            loss_susceptible_R = 0.5 .* integrator.u[:,7,:]
            loss_susceptible_W = 0.5 .* integrator.u[:,8,:]
            integrator.u[:,1,:] .+= loss_susceptible_R
            integrator.u[:,7,:] .-= loss_susceptible_R
            integrator.u[:,1,:] .+= loss_susceptible_W
            integrator.u[:,8,:] .-= loss_susceptible_W
        end
        immune_escape_variant_cb = PresetTimeCallback([immune_escape_variant_time],immune_escape_variant!,save_positions=(false,false))

        # Update callbacks based on scenario
        if  scenario==1 || scenario>4  # For Scenario 4-6
         cb = CallbackSet(cb, immune_escape_variant_cb)
        end

        county_model = KenyaCoVaccines.CoVAreaModel(name,KenyaCoVaccines.ll_onegroup_newvariant_infboost,prior;
                                PCR_cases = linelist_data_with_pos_neg,
                                sero_cases = serology_data,
                                average_sero_init = zeros(6),
                                deaths = deaths_data,
                                pop_data= N_kenya,
                                M_county_ho=M_Kenya_ho,
                                M_county_other=M_Kenya_other,
                                M_county_school=M_Kenya_school,
                                M_county_work=M_Kenya_work,
                                rapid = rapid,
                                scenario= scenario,
                                startdate=startdate,
                                enddate=enddate)

        #Solve dynamical system
        sol = solve(county_model.prob, BS3();tspan = (0,model_end_point),callback = cb,u0=u0, p=p,saveat = 1) #isoutofdomain=(u,p,t)->any(x->x<0,u)

        # Number of unscalled (for testing) cases over time by severity
        ι = KenyaCoVaccines.get_incidence_time_array_by_age_dose(sol,10)

        # Number of tests over time by age
        PCR_cases = sum(linelist_data_with_pos_neg.cases[:,linelist_data_with_pos_neg.areas .== up_county,:,:],dims = [2])[:,1,:,:]
        p_vec = [[1,2,3,4],[5,6,7,8,9,10],[11,12],[13,14],[15,16],[17]]
        PCR_cases_collaped =  zeros(Int64,size(PCR_cases,1),6,2)
        for a in 1:6
            PCR_cases_collaped[:,a:a,:] .= sum(PCR_cases[:,p_vec[a][1]:p_vec[a][end],:], dims=2)
        end
        PCR_cases = PCR_cases_collaped
        num_tests =  PCR_cases[model_start_time:end,:,1] .+ PCR_cases[model_start_time:end,:,2]
        num_tests_age = zeros(size(ι,1),6)
        for a in 1:6
            num_tests_age[:,a] = vcat(num_tests[:,a],fill(mean(num_tests[:,a]),(length(ι)-length(num_tests[:,a]))))  # use mean number of teste for forcast periods, could be changed
        end

        # scalling parameters
        p_test = MCMCdraws.p_test[k]
        χ = MCMCdraws.χ[k]
        p_test_scale = 5e-4

        # Apply vaccine efficacy against disease
        v = (1 .- model.VE_severe_disease_risk') .* υ  # efficacy

        l_inc = [inc .* v for inc in ι]
        cases_wt = similar(ι)
        cases_αβ = similar(ι)
        cases_δ  = similar(ι)
        crit_incidence = similar(ι)
        sev_incidence = similar(ι)

        for i in 1:size(ι,1)
            cases_wt[i] = sum(l_inc[i] .* other_perc[i],dims=2) # sum over dose
            cases_αβ[i] = sum(l_inc[i] .* alpha_beta_perc[i], dims=2)
            cases_δ[i]  = sum(l_inc[i] .* delta_perc[i], dims=2)

            PCR_num_wt = cases_wt[i] .* p_test_scale .* p_test
            PCR_denom_wt = PCR_num_wt .+ p_test_scale*((p_test/χ)*(N .- cases_wt[i]))
            p_PCR_pred_wt = PCR_num_wt./PCR_denom_wt
            cases_wt[i]  =  p_PCR_pred_wt .* num_tests_age[i,:]

            PCR_num_αβ = cases_αβ[i] .* p_test_scale .* p_test
            PCR_denom_αβ = PCR_num_αβ .+ p_test_scale*((p_test/χ)*(N .- cases_αβ[i]))
            p_PCR_pred_αβ = PCR_num_αβ./PCR_denom_αβ
            cases_αβ[i]  =  p_PCR_pred_αβ .* num_tests_age[i,:]

            PCR_num_δ = cases_δ[i] .* p_test_scale .* p_test
            PCR_denom_δ = PCR_num_wt .+ p_test_scale*((p_test/χ)*(N .- cases_δ[i]))
            p_PCR_pred_δ = PCR_num_δ./PCR_denom_δ
            cases_δ[i]  =  p_PCR_pred_δ .* num_tests_age[i,:]

            # need to edit structure on ICU and hosp fit to include kth sample fit rathert than point est (single value)
            crit_incidence[i] = ICU_fit[1]*(cases_wt[i] .+ ICU_fit[2].*cases_αβ[i] .+ ICU_fit[3].*cases_δ[i])
            sev_incidence[i] = hosp_fit[1]*(cases_wt[i] .+ hosp_fit[2].*cases_αβ[i] .+ hosp_fit[3].*cases_δ[i])
        end

        crit_incidence =  Matrix(VectorOfArray([inc[:] for inc in crit_incidence])')  # convert to a matrix
        sev_incidence =  Matrix(VectorOfArray([inc[:] for inc in sev_incidence])')

        for a in 1:6 # convolve with PCR array to get right timing of PCR detectable incidence
            crit_incidence[:,a] = KenyaCoVaccines.simple_conv(crit_incidence[:,a],KenyaCoVaccines.PCR_array)
            sev_incidence[:,a] = KenyaCoVaccines.simple_conv(sev_incidence[:,a],KenyaCoVaccines.PCR_array)
        end

        # Saving severe and critical disease incidence
        crit_dis_inc_samples[:,:,k]   .= crit_incidence
        severe_dis_inc_samples[:,:,k] .= sev_incidence

        # Asymptomatic cases
        asymp_dis_inc = [inc .* (ones(3)'.* (1 .- δ )) for inc in ι]
        asymp_dis_inc = Matrix(VectorOfArray([sum(inc,dims=2)[:] for inc in asymp_dis_inc])')  #sum over dose
        asymp_dis_inc_samples[:,:,k] .= asymp_dis_inc

        # Mild disease
        symp_dis_inc  =  [inc .* (ones(3)'.* δ) for inc in ι]
        symp_dis_inc =  Matrix(VectorOfArray([sum(inc,dims=2)[:] for inc in symp_dis_inc])')  #sum over dose
        for a in 1:6
            symp_dis_inc[:,a] =  KenyaCoVaccines.simple_conv(symp_dis_inc[:,a],p_IH)  # convolve with time to hsop function to get right timing of Hospitalisation
        end
        mild_dis_inc = symp_dis_inc .- sev_incidence .- crit_incidence
        mild_dis_inc_samples[:,:,k] .= mild_dis_inc

        # Number of unscalled deaths by age group over time
        IFR = all_IFR_fits[[nm[1] .== name for nm in all_IFR_fits]][1].IFR # EDIT to include HMC samples, so there shoudl be a 'k' index on the IFR
        d = (1 .- VE_severe_disease_risk') .* IFR # efficacy multiplied by risk of death by age, assumes only those with severe disease die

        death_inc =  [inc .* d for inc in ι]
        death_inc =  Matrix(VectorOfArray([sum(inc,dims=2)[:] for inc in death_inc])')  #sum over dose
        for a in 1:6
            death_inc[:,a] =  KenyaCoVaccines.simple_conv(death_inc[:,a],p_ID)  # convolve with time to death function to get right timing of deaths
        end
        death_inc = death_inc ./ obs_factor

    end # end loop over HMC chains

    # pick data from Jan for CEA analysis
    asymp_dis_inc_samples = asymp_dis_inc_samples[janstarttime:end,:,:]
    crit_dis_inc_samples = crit_dis_inc_samples[janstarttime:end,:,:]
    severe_dis_inc_samples = severe_dis_inc_samples[janstarttime:end,:,:]
    mild_dis_inc_samples = mild_dis_inc_samples[janstarttime:end,:,:]
    death_inc_samples = death_inc_samples[janstarttime:end,:,:]

    if rapid  # becasue we now have rapid scenario we save them in different directory: I have created this filing structure in the repository
        rapid_text = "rapid"
    else
        rapid_text = "non_rapid"
    end

    @save("scenario_$(scenario)/$(rapid_text)/mild_cases/severe_cases_$(name).jld2",mild_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/severe_cases/severe_cases_$(name).jld2",severe_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/critical_cases/critical_cases_$(name).jld2",crit_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/asymp_cases/asymp_cases_$(name).jld2",asymp_dis_inc_samples)
    @save("scenario_$(scenario)/$(rapid_text)/deaths/deaths_$(name).jld2",death_inc_samples)

    # Number of vaccines administered over time, this is the same for all 2000 chains
    doses_Kenya_baseline = 0.
    if scenario==1
        doses_Kenya = 0.
    else
        doses_Kenya = (1.25 + 9.76 +4.9) .* 1000000  #Number of doses planned for Kenya over the ynext one year under a 30% cover plan
    end

    age_dist_vacc, doses_daily = KenyaCoVaccines.age_dist([name],doses_Kenya,N_kenya,vacc_end_point,pre_vacc_period,rapid)
    vacc_rate_1,vacc_rate_2 = KenyaCoVaccines.vacc_rate_data(doses_daily,vacc_end_point,pre_vacc_period,age_dist_vacc,N_kenya,name,scenario,rapid)
    num_doses = sum(vacc_rate_1, dims=2)[:] .+ sum(vacc_rate_2, dims=2)[:]
    num_doses = num_doses[janstarttime:end]
    complete_doses_age = vacc_rate_2[janstarttime:end,:]  # number with complete doses

    age_dist_vacc, doses_daily = KenyaCoVaccines.age_dist([name],doses_Kenya_baseline,N_kenya,vacc_end_point,pre_vacc_period,rapid)
    vacc_rate_1,vacc_rate_2 = KenyaCoVaccines.vacc_rate_data(doses_daily,vacc_end_point,pre_vacc_period,age_dist_vacc,N_kenya,name,baseline_scenario,rapid)
    num_doses_baseline = sum(vacc_rate_1, dims=2)[:] .+ sum(vacc_rate_2, dims=2)[:]
    num_doses_baseline = num_doses_baseline[janstarttime:end]
    complete_doses_age_baseline = vacc_rate_2[janstarttime:end,:]

    @save("scenario_$(scenario)/$(rapid_text)/doses/doses_$(name).jld2",num_doses)
    @save("scenario_$(scenario)/$(rapid_text)/doses_age/doses_$(name).jld2",complete_doses_age)

    return nothing
end


text:
# ---
# title: Custom scales
# cover: assets/custom_scales.png
# description: Custom palettes and custom attributes.
# author: "[Pietro Vertechi](https://github.com/piever)"
# id: custom_scales
# ---

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

# A palette maps particular values to particular attribute specifications (e.g.
# 1 maps to green, 2 maps to red). Sometimes, there is no default palette for a
# specific attribute, and you will need to specify it manually, but there are
# sensible default palettes for many attributes. In either case, you can always
# manually specify the palette used for a particular attribute.

# TODO: allow legend to use custom attribute of plot, such as the arrowhead or
# the arrowcolor and pass correct legend symbol.

# !!! note
#     A related concept (from Makie) is a colormap, which maps a continuous space
#     of numbers to a sequence of colors. For discrete colors you will want to
#     employ a palette, not a colormap, because discrete values are mapped to
#     colors within AlgebraOfGraphics. AlgebraOfGraphics doesn't directly handle
#     continuous colors: this is a feature of the underlying Makie plots.

x=repeat(1:20, inner=20)
y=repeat(1:20, outer=20)
u=cos.(x)
v=sin.(y)
c=rand(Bool, length(x))
d=rand(Bool, length(x))
df = (; x, y, u, v, c, d)
colors = [colorant"#E24A33", colorant"#348ABD"]
heads = ['▲', '●']
plt = data(df) *
    mapping(:x, :y, :u, :v) *
    mapping(arrowhead=:c => nonnumeric) *
    mapping(arrowcolor=:d => nonnumeric) *
    visual(Arrows, arrowsize=10, lengthscale=0.3)
draw(plt; palettes=(arrowcolor=colors, arrowhead=heads))

# To associate specific attribute values to specific data values, use pairs.
# Missing keys will cycle over values that are not pairs.

x = rand(100)
y = rand(100)
z = rand(["a", "b", "c", "d"], 100)
df = (; x, y, z)
plt = data(df) * mapping(:x, :y, color=:z)
colors = ["a" => colorant"#E24A33", "c" => colorant"#348ABD", colorant"#988ED5", colorant"#777777"]
draw(plt; palettes=(color=colors,))

# Categorical color gradients can also be passed to `palettes`.

x = rand(200)
y = rand(200)
z = rand(["a", "b", "c", "d", "e", "f", "g", "h"], 200)
df = (; x, y, z)
plt = data(df) * mapping(:x, :y, color=:z)
colors = cgrad(:cividis, 8, categorical=true)
fg = draw(plt; palettes=(color=colors,))

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


text:
using JetPack, Jets, Test

n1dom,n2dom = 10,11
n1rng,n2rng = 29,31

@testset "JopInterp, linearity test" for T in (Float32,Float64)
    A = JopInterp(JetSpace(T,n1dom,n2dom), JetSpace(T,n1rng,n2rng))
    lhs,rhs = linearity_test(A)
    @test lhs ≈ rhs
end

@testset "JopInterp, dot product test, T=$(T)" for T in (Float32,Float64)
    A = JopInterp(JetSpace(T,n1dom,n2dom), JetSpace(T,n1rng,n2rng))
    lhs,rhs = dot_product_test(A, rand(domain(A)), rand(range(A)))
    @test lhs ≈ rhs
end

@testset "JopInter, test bugfix for specific dom,rng, T=$(T)" for T in (Float32,Float64)
    A = JopInterp(JetSpace(T,280,222), JetSpace(T,281,2221))
    lhs,rhs = dot_product_test(A, rand(domain(A)), rand(range(A)))
    @test lhs ≈ rhs
end