text:
"""
Analytical reports generated during training, for debugging and
hyperparameters tuning.
"""
module Report

# Some imports for docstrings cross-references
import ..MCTS, ..Network, ..MemAnalysisParams

"""
    Report.Loss

Decomposition of the loss in a sum of terms (all have type `Float32`).

- `L` is the total loss: `L == Lp + Lv + Lreg + Linv`
- `Lp` is the policy cross-entropy loss term
- `Lv` is the average value mean square error
- `Lreg` is the L2 regularization loss term
- `Linv` is the loss term penalizing the average weight put by the
  network on invalid actions
"""
struct Loss
  L :: Float32
  Lp :: Float32
  Lv :: Float32
  Lreg :: Float32
  Linv :: Float32
end

"""
    Report.LearningStatus

Statistics about the performance of the neural network on a subset of the
memory buffer.

- `loss`: detailed loss on the samples, as an object of type
    [`Report.Loss`](@ref)
- `Hp`: average entropy of the ``π`` component of samples (MCTS policy);
    this quantity is independent of the network and therefore constant
    during a learning iteration
- `Hpnet`: average entropy of the network's prescribed policy on the samples
"""
struct LearningStatus
  loss :: Loss
  Hp :: Float32 # property of the memory, constant during a learning iteration
  Hpnet :: Float32
end

"""
    Report.Checkpoint

Report generated after a checkpoint evaluation.

- `batch_id`: number of batches after which the checkpoint was computed
- `status_after`: learning status at the checkpoint, as an object of type
   [`Report.LearningStatus`](@ref)
- `reward`: average reward collected by the contender network
- `redundancy`: ratio of duplicate positions encountered during the evaluation,
   not counting the initial position. If this number is too high, you may
   want to increase the move selection temperature.
- `nn_replaced`: true if the current best neural network was updated after
   the checkpoint
"""
struct Checkpoint
  batch_id :: Int
  status_after :: LearningStatus
  reward :: Float64
  redundancy :: Float64
  nn_replaced :: Bool
end

"""
    Report.Learning

Report generated at the end of the learning phase of an iteration.

- `time_convert`, `time_loss`, `time_train` and `time_eval` are the
    amounts of time (in seconds) spent at converting the samples,
    computing losses, performing gradient updates and evaluating checkpoints
    respectively
- `initial_status`: status before the learning phase, as an object of type
    [`Report.LearningStatus`](@ref)
- `losses`: loss value on each minibatch
- `checkpoints`: vector of [`Report.Checkpoint`](@ref) reports
- `nn_replaced`: true if the best neural network was replaced
"""
struct Learning
  time_convert :: Float64
  time_loss :: Float64
  time_train :: Float64
  time_eval :: Float64
  initial_status :: LearningStatus
  losses :: Vector{Float32}
  checkpoints :: Vector{Checkpoint}
  nn_replaced :: Bool
end

"""
    Report.Samples

Statistics about a set of samples, as collected during memory analysis.

- `num_samples`: total number of samples
- `num_boards`: number of distinct board positions
- `Wtot`: total weight of the samples
- `status`: [`Report.LearningStatus`](@ref) statistics of the current network
    on the samples
"""
struct Samples
  num_samples :: Int
  num_boards :: Int
  Wtot :: Float32
  status :: LearningStatus
end

"""
    Report.StageSamples

Statistics for the samples corresponding to a particular game stage,
as collected during memory analysis.

The samples whose statistics are collected in the
[`samples_stats`](@ref Report.Samples) field correspond to historical positions
where the number of remaining moves until the end of the game was in the range
defined by the `min_remaining_length` and `max_remaining_length` fields.
"""
struct StageSamples
  min_remaining_length :: Int
  max_remaining_length :: Int
  samples_stats :: Samples
end

"""
    Report.Memory

Report generated by the memory analysis phase of an iteration. It features
statistics for
  - the whole memory buffer (`all_samples::Report.Samples`)
  - the samples collected during the last self-play iteration
     (`latest_batch::Report.Samples`)
  - the subsets of the memory buffer corresponding to different game stages:
     (`per_game_stage::Vector{Report.StageSamples}`)

See [`MemAnalysisParams`](@ref).
"""
struct Memory
  latest_batch :: Samples
  all_samples :: Samples
  per_game_stage :: Vector{StageSamples}
end

"""
    Report.SelfPlay

Report generated after the self-play phase of an iteration.

- `samples_gen_speed`: average number of samples generated per second
- `average_exploration_depth`: see [`MCTS.average_exploration_depth`](@ref)
- `mcts_memory_footprint`: estimation of the maximal memory footprint of the
    MCTS tree during self-play, as computed by
    [`MCTS.approximate_memory_footprint`](@ref)
- `memory_size`: number of samples in the memory buffer at the end of the
    self-play phase
- `memory_num_distinct_boards`: number of distinct board positions in the
    memory buffer at the end of the self-play phase
"""
struct SelfPlay
  samples_gen_speed :: Float64
  average_exploration_depth :: Float64
  mcts_memory_footprint :: Int
  memory_size :: Int
  memory_num_distinct_boards :: Int
end

"""
    Report.Perfs

Performances report for a subroutine.
- `time`: total time spent, in seconds
- `allocated`: amount of memory allocated, in bytes
- `gc_time`: total amount of time spent in the garbage collector
"""
struct Perfs
  time :: Float64
  allocated :: Int64
  gc_time :: Float64
end

"""
    Report.Iteration

Report generated after each training iteration.
- Fields `self_play`, `memory`, `learning` have types [`Report.SelfPlay`](@ref),
    [`Report.SelfPlay`](@ref) and [`Report.Learning`](@ref) respectively
- Fields `perfs_self_play`, `perfs_memory_analysis` and `perfs_learning` are
    performance reports for the different phases of the iteration,
    with type [`Report.Perfs`](@ref)
"""
struct Iteration
  perfs_self_play :: Perfs
  perfs_memory_analysis :: Perfs
  perfs_learning :: Perfs
  self_play :: SelfPlay
  memory :: Union{Memory, Nothing}
  learning :: Learning
end

"""
    Report.Initial

Report summarizing the configuration of an agent before training starts.
- `num_network_parameters`: see [`Network.num_parameters`](@ref)
- `num_network_regularized_parameters`:
    see [`Network.num_regularized_parameters`](@ref)
- `mcts_footprint_per_node`: see [`MCTS.memory_footprint_per_node`](@ref)
"""
struct Initial
  num_network_parameters :: Int
  num_network_regularized_parameters :: Int
  mcts_footprint_per_node :: Int
end

#####
##### Profiling utilities
#####

macro timed(e)
  quote
    local v, t, mem, gct = Base.@timed $(esc(e))
    v, Perfs(t, mem, gct)
  end
end

end


text:
"""
    struct ComplexSpace <: EuclideanSpace{ℂ}

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

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

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

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

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

Base.min(V1::ComplexSpace, V2::ComplexSpace) = isdual(V1) == isdual(V2) ?
    ComplexSpace(min(dim(V1), dim(V2)), isdual(V1)) :
    throw(SpaceMismatch("Mininum of space and dual space does not exist"))
Base.max(V1::ComplexSpace, V2::ComplexSpace) = isdual(V1) == isdual(V2) ?
    ComplexSpace(max(dim(V1), dim(V2)), isdual(V1)) :
    throw(SpaceMismatch("Maximum of space and dual space does not exist"))

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

ext: jl


text:
# Author: Ritchie Lee, ritchie.lee@sv.cmu.edu
# Date: 11/11/2014

module GenericCoordImpl

export
    initialize,
    update,

    setRecord,
    getRecord,
    getAll,

    GenericCoord,
    GenericCoordRecord

using AbstractCASCoordImpl
using AbstractCASCoordInterfaces
using CommonInterfaces

import CommonInterfaces.initialize
import CommonInterfaces.update
import AbstractCASCoordInterfaces.setRecord
import AbstractCASCoordInterfaces.getRecord
import AbstractCASCoordInterfaces.getAll

type GenericCoord <: AbstractCASCoord

  num_records::Int64
  records::Vector{Any}

  function GenericCoord(num_records::Int)
    obj = new()

    obj.num_records = num_records
    obj.records = [nothing for i=1:num_records]

    return obj
  end
end

function setRecord(coord::GenericCoord, record_number::Int, record)

    coord.records[record_number] = record
end

update(coord::GenericCoord, record_number::Int, data) = setRecord(coord, record_number, record)

function initialize(coord::GenericCoord)

  fill!(coord.records, nothing)
end

getRecord(coord::GenericCoord,record_number::Int) = coord.records[record_number]

getAll(coord::GenericCoord) = coord.records

end


text:
#
# This file is part of the Actors.jl Julia package, MIT license
#
# Paul Bayer, 2020
#

using Actors, Distributed, Test
import Actors: spawn

length(procs()) == 1 && addprocs(1)

@everywhere using Actors
@everywhere function ident(id, from)
    id == from ?
        ("local actor",  id, from) :
        ("remote actor", id, from)
end
call(Actors._REG, empty!)

#
# test registry with one local actor
#
@test register(:act1, spawn(Bhv(ident, 1)))
a1 = Actors.diag(:act1, :act)
@test a1.name == :act1  # does it have its name ?
@test call(:act1, myid()) == ("local actor", 1, 1)
act1 = whereis(:act1)
unregister(:act1)
sleep(0.1)
@test isnothing(a1.name) # did it delete the name ?
@test isempty(registered())
@test register(:act1, act1)
sleep(0.1)
@test length(registered()) == 1
exit!(:act1)
sleep(0.1)
@test act1.chn.state == :closed
@test isempty(registered()) # did it unregister at exit ?
#
# test registry across workers
#
@test register(:act1, spawn(Bhv(ident, 1)))
@test register(:act2, spawn(Bhv(ident, 2), pid=2))
@test call(:act2, myid()) == ("remote actor", 2, 1)
@test fetch(@spawnat 2 call(:act1, myid())) == ("remote actor", 1, 2)
@test fetch(@spawnat 2 call(:act2, myid())) == ("local actor", 2, 2)
@test whereis(:act1).pid == 1
@test whereis(:act2).pid == 2
@test fetch(@spawnat 2 whereis(:act1)).pid == 1
@test fetch(@spawnat 2 whereis(:act2)).pid == 2
let
    r = registered()
    l = [i[1] for i in r]
    @test length(r) == 2
    @test :act1 in l
    @test :act2 in l
    r = fetch(@spawnat 2 registered())
    @test length(r) == 2
    @test all([i[2].chn for i in r]) do x
        x isa RemoteChannel
    end
end
unregister(:act1)
unregister(:act2)
@test isempty(registered())
#
# test API functions with registered actors
#
f(a, b) = a + b
@test register(:act1, spawn(Bhv(f, 1)))
@test !register(:act1, spawn(f, 2))
a1 = Actors.diag(:act1, :act)
send(:act1, Actors.Cast((1,)))
sleep(0.1)
@test a1.res == 2
@test request(:act1, Actors.Call, 2) == 3
become!(:act1, f, 0)
@test call(:act1, 1) == 1
cast(:act1, 2)
sleep(0.1)
@test a1.res == 2
@test exec(:act1, Bhv(f, 5, 5)) == 10
@test query(:act1, :res) == 2
init!(:act1, cos, 2pi)
sleep(0.1)
@test a1.init.f == cos
term!(:act1, cos, 2pi)
sleep(0.1)
@test a1.term.f == cos
sleep(0.1)
update!(:act1, 10)
sleep(0.1)
@test a1.sta == 10
unregister(:act1)


text:
####
#### Stan interface
####

# NOTE:
#   Type fields
#     fieldnames(CmdStan.Sample)
#       num_samples, num_warmup, save_warmup, thin, adapt, algorithm

#     fieldnames(CmdStan.Hmc)
#       engine, metric, stepsize, stepsize_jitter

#     fieldnames(CmdStan.Adapt)
#       engaged, gamma, delta, kappa, t0, init_buffer, term_buffer, window

#   Ref
#     http://goedman.github.io/Stan.jl/latest/index.html#Types-1

function sample(mf::T, ss::CmdStan.Sample) where T
    return sample(mf, ss.num_samples, ss.num_warmup,
                    ss.save_warmup, ss.thin, ss.adapt, ss.alg)
end

function sample(mf::T,
    num_samples::Int,
    num_warmup::Int,
    save_warmup::Bool,
    thin::Int,
    ss::CmdStan.Sample
) where T
    return sample(mf, num_samples, num_warmup, save_warmup, thin, ss.adapt, ss.alg)
end

function sample(mf::T,
    num_samples::Int,
    num_warmup::Int,
    save_warmup::Bool,
    thin::Int,
    adapt::CmdStan.Adapt,
    alg::CmdStan.Hmc
) where T
    if alg.stepsize_jitter != 0.0
        @warn("[Turing.sample] Turing does not support adding noise to stepsize yet.")
    end
    if adapt.engaged == false
        if isa(alg.engine, CmdStan.Static)   # hmc
            stepnum = Int(round(alg.engine.int_time / alg.stepsize))
            sample(mf, HMC(num_samples, alg.stepsize, stepnum); adaptor=NUTSAdaptor(adapt))
        elseif isa(alg.engine, CmdStan.Nuts) # error
            error("[Turing.sample] CmdStan.Nuts cannot be used with adapt.engaged set as false")
        end
    else
        if isa(alg.engine, CmdStan.Static)   # hmcda
            sample(mf, HMCDA(num_samples, num_warmup, adapt.delta, alg.engine.int_time);
                    adaptor=NUTSAdaptor(adapt))
        elseif isa(alg.engine, CmdStan.Nuts) # nuts
            if isa(alg.metric, CmdStan.diag_e)
                sample(mf, NUTS(num_samples, num_warmup, adapt.delta);
                        adaptor=NUTSAdaptor(adapt))
            else # TODO: reove the following since Turing support this feature now.
                @warn("[Turing.sample] Turing does not support full covariance matrix for pre-conditioning yet.")
            end
        end
    end
end

function AHMCAdaptor(adaptor::CmdAdaptorType) where CmdAdaptorType
    if :engaged in fieldnames(typeof(adaptor)) # CmdStan.Adapt
        adaptor.engaged ? spl.alg.n_adapts : 0,
        AHMC.Preconditioner(metric),
        AHMC.NesterovDualAveraging(adaptor.gamma,
            adaptor.t0, adaptor.kappa, adaptor.δ, init_ϵ),
            adaptor.init_buffer,
            adaptor.term_buffer,
            adaptor.window
    else # default adaptor
        @warn "Invalid adaptor type: $(typeof(adaptor)). Default adaptor is used instead."
        adaptor = AHMC.StanHMCAdaptor(
            spl.alg.n_adapts, AHMC.Preconditioner(:DiagEuclideanMetric),
            AHMC.NesterovDualAveraging(spl.alg.δ, init_ϵ)
        )
    end
end


text:
Random.seed!(0)

# facts("qraddcol") do

    # context("β = 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(undef, 0,0)
        r = Vector{Float64}(undef, m)

        for i in 1:m
            a = randn(m)
            Rout = Array{Float64}(undef, i, i)
            qraddcol!(A, R, a, Rout, r)
            global R = qraddcol(A, R, a)
            global A = [A a]
            @assert  norm( R'*R - A'*A ) < 1e-5
            @assert  norm( Rout'*Rout - A'*A ) < 1e-5 "$(norm( Rout'*Rout - A'*A )) is not less than 1e-5"
        end

        b = sum(A, dims=2)[:]
        (x, rtest) = csne(R, A, b)
        @assert norm(rtest) < 1e-5
        @assert norm(x .- 1) < 1e-5

        x = csne!(R, A, b, r)
        @assert norm(r) < 1e-5
        @assert norm(x .- 1) < 1e-5

    # end

    # context("β > 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(undef, 0,0)
        β = 0.1
        for i in 1:m
            a = randn(m)
            B = [A a]
            global R = qraddcol(A, R, a, β)
            @assert norm( R'*R - B'*B - β^2*I ) < 1e-5
        end
    # end

# end


text:
function driver_inductionless_MHD(;model=nothing, cutgeo=nothing,nx = 4,
  Re::Float64 = 10.0,
  Ha = (x) -> 10.0, fluid_dirichlet_tags = [], fluid_neumann_tags = [],
  c_w = 1.0, α=10.0, B = (x) -> VectorValue(0.0,1.0,0.0),
  magnetic_dirichlet_tags = [], magnetic_neumann_tags = [],
  magnetic_non_perfectly_conducting_walls_tag = [],
  fluid_dirichlet_conditions = (x) -> VectorValue(0.0,0.0,0.0),
  magnetic_dirichlet_conditions = (x) -> VectorValue(0.0,0.0,0.0),
  fluid_body_force = (x) -> VectorValue(0.0,0.0,0.0),
  constraint_presures::NTuple{2,Bool}=(false,false), max_nl_it=10,
  usegmres = true, precond_tau = 1e-9, linesearch=BackTracking(),
  resultsfile = nothing, verbosity::Verbosity=Verbosity(1),nsubcells=1,print_order=2)

  if typeof(Ha) == typeof(10.0)
    N = (Ha*Ha)/Re
  else
    N(x) = (Ha(x)*Ha(x))/Re
  end

  # Discretization
  order = 2
  if model === nothing
    domain = (-1.0,1.0,-1.0,1.0,0.0,0.1)
    map(x) = VectorValue(sign(x[1])*(abs(x[1])*0.5)^0.5,
                        sign(x[2])*(abs(x[2])*0.5)^0.5, x[3]*sqrt(2)/2)*2/sqrt(2)

    dirichlet_tags_u = append!(collect(1:20),[23,24,25,26])
    dirichlet_tags_j = append!(collect(1:20),[23,24,25,26])

    partition=(nx,ny,3)
    model = CartesianDiscreteModel(domain,partition;isperiodic=(false,false,true))

    labels = get_face_labeling(model)
    add_tag_from_tags!(labels,"dirichlet_u",dirichlet_tags_u)
    add_tag_from_tags!(labels,"dirichlet_j",dirichlet_tags_j)
    fluid_dirichlet_tags = ["dirichlet_u"]
    magnetic_dirichlet_tags = ["dirichlet_j"]
  end

  preprocessStepOutput(verbosity,"Preprocess step (0/3): Starting preprocess")

  Vu = FESpace(model, ReferenceFE(lagrangian,VectorValue{3,Float64},order);
      conformity=:H1, dirichlet_tags=fluid_dirichlet_tags)

  grid_topology = get_grid_topology(model)
  if is_n_cube(grid_topology.polytopes[1])
    p_conformity = :L2
  else
    p_conformity = :H1
  end

  if constraint_presures[1]
    Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
    conformity=p_conformity, constraint=:zeromean)
  else
    Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
    conformity=p_conformity)
  end

  if length(magnetic_dirichlet_tags) == 0
    Vj = FESpace(model, ReferenceFE(raviart_thomas,Float64,order-1);
        conformity=:Hdiv)
  else
    Vj = FESpace(model, ReferenceFE(raviart_thomas,Float64,order-1);
        conformity=:Hdiv, dirichlet_tags=magnetic_dirichlet_tags)
  end

  if constraint_presures[2]
    Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1);
      conformity=:L2, constraint=:zeromean)
  else
    Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1);
    conformity=:L2)
  end

  U = TrialFESpace(Vu,fluid_dirichlet_conditions)
  P = TrialFESpace(Vp)
  J = TrialFESpace(Vj,magnetic_dirichlet_conditions)
  Φ = TrialFESpace(Vφ)

  Y = MultiFieldFESpace([Vu, Vp, Vj, Vφ])
  X = MultiFieldFESpace([U, P, J, Φ])

  preprocessStepOutput(verbosity,"Preprocess step (1/3): Built FE spaces")

  trian = Triangulation(model)
  degree = 2*(order)
  dΩ = Measure(trian,degree)

  coord = get_physical_coordinate(trian)
  res_Ω(x,y) = InductionlessMHD.dimensionless_residual(x, y, Re, N, B, fluid_body_force, dΩ, coord)
  jac_Ω(x,dx,y) = InductionlessMHD.dimensionless_jacobian(x, dx, y, Re, N, B, dΩ, coord)

  if length(magnetic_non_perfectly_conducting_walls_tag) == 0
    op  = FEOperator(res_Ω,jac_Ω,X,Y)
  else
    if cutgeo === nothing
      btrian_j = BoundaryTriangulation(model,tags=magnetic_non_perfectly_conducting_walls_tag)
    else
      btrian_j = EmbeddedBoundary(cutgeo)
    end
    dΓ = Measure(btrian_j,degree)
    n = get_normal_vector(btrian_j)
    res_Γ(x,y) = InductionlessMHD.dimensionless_conducting_wall(x, y, n, c_w, dΓ, α=α)
    jac_Γ(x,dx,y) = InductionlessMHD.dimensionless_conducting_wall(dx, y, n, c_w, dΓ, α=α)

    res(x,y) = res_Ω(x,y) + res_Γ(x,y)
    jac(x,dx,y) = jac_Ω(x,dx,y) + jac_Γ(x,dx,y)
    op  = FEOperator(res,jac,X,Y)
  end

  preprocessStepOutput(verbosity,"Preprocess step (2/3): Defined triangulation and formulation")

  if usegmres
    nls = NLSolver(GmresSolver(verbose=linearSolverOutput(verbosity), preconditioner=ilu,τ=precond_tau);
      show_trace=nonlinearSolverOutput(verbosity), method=:newton, linesearch=linesearch, iterations=max_nl_it)
  else
    nls = NLSolver(;
      show_trace=nonlinearSolverOutput(verbosity), method=:newton, linesearch=linesearch, iterations=max_nl_it)
  end
  solver = FESolver(nls)

  preprocessStepOutput(verbosity,"Preprocess step (3/3): Configured solver. Starting to solve.")

  xh = solve(solver,op)

  if resultsfile ≠ nothing
    uh, ph, jh, φh = xh
    writevtk(trian, resultsfile, nsubcells=nsubcells, order=print_order,
      cellfields=["uh"=>uh, "ph"=>ph, "jh"=>jh, "phih"=>φh])
  end

  (xh, trian, dΩ)
end


text:
module TestRandom
using MXNet
using Base.Test

function test_uniform()
  dims = (100, 100, 2)
  info("random::uniform::dims = $dims")

  low = -10; high = 10
  seed = 123
  mx.srand!(seed)
  ret1 = mx.rand(low, high, dims)

  mx.srand!(seed)
  ret2 = mx.empty(dims)
  mx.rand!(low, high, ret2)

  @test copy(ret1) == copy(ret2)
  @test abs(mean(copy(ret1)) - (high+low)/2) < 0.1
end

function test_gaussian()
  dims = (80, 80, 4)
  info("random::gaussian::dims = $dims")

  μ = 10; σ = 2
  seed = 456
  mx.srand!(seed)
  ret1 = mx.randn(μ, σ, dims)

  mx.srand!(seed)
  ret2 = mx.empty(dims)
  mx.randn!(μ, σ, ret2)

  @test copy(ret1) == copy(ret2)
  @test abs(mean(copy(ret1)) - μ) < 0.1
  @test abs(std(copy(ret1)) - σ) < 0.1
end

test_uniform()
test_gaussian()

end


text:
md = fill("metadata", 4)
an1 = ArrayNode(rand(3, 4))
b = BagNode(an1, [1:4, 0:-1], md)
an2 = ArrayNode(randn(5, 4))
wb = WeightedBagNode(an2, [1:2, 3:4], rand(4), md)
pn = ProductNode((b=b, wb=wb))
an3 = ArrayNode(rand(10, 2))
x = ProductNode((pn, an3))
m = reflectinmodel(x)

@testset "list_lens" begin
    ls = list_lens(x)
    all_nodes = NodeIterator(x) |> collect
    all_fields = vcat(all_nodes, [md, an1.data, b.bags, an2.data, wb.bags, wb.weights, an3.data])
    all_fields = vcat(all_fields, Mill.metadata.(all_nodes))

    @test all(l -> get(x, l) in all_fields, ls)
    @test all(n -> n in all_fields, [walk(x, t) for t in list_traversal(x)])

    ls = list_lens(m)
    all_nodes = NodeIterator(m) |> collect
    all_fields = vcat(all_nodes, [m.m, m.m.m,
                      m.ms[1].m, m.ms[1].m.m,
                      m.ms[1].ms.b.im.m, m.ms[1].ms.b.a, m.ms[1].ms.b.bm, m.ms[1].ms.b.bm.m,
                      m.ms[1].ms.wb.im.m, m.ms[1].ms.wb.a, m.ms[1].ms.wb.bm, m.ms[1].ms.wb.bm.m,
                      m.ms[2].m])

    @test all(l -> get(m, l) in all_fields, ls)
    @test all(n -> n in all_fields, [walk(m, t) for t in list_traversal(m)])
end

@testset "findnonempty_lens" begin
    @test all(nobs.([get(x, l) for l in findnonempty_lens(x)]) .> 0)
end

@testset "find_lens" begin
    for t in list_traversal(x)
        ls = find_lens(x, x[t])
        @test all(l -> get(x, l) === x[t], ls)
    end

    for t in list_traversal(m)
        ls = find_lens(m, m[t])
        @test all(l -> get(m, l) === m[t], ls)
    end
end

@testset "code2lens & lens2code" begin
    for t in list_traversal(x)
        @test t == lens2code(x, code2lens(x, t))
    end
    for t in list_traversal(m)
        @test t == lens2code(m, code2lens(m, t))
    end
end

@testset "model_lens & data_lens" begin
    for n in NodeIterator(x)
        l = find_lens(x, n) |> only
        @test l == data_lens(x, model_lens(m, l))
    end
    for n in NodeIterator(m)
        l = find_lens(m, n) |> only
        @test l == model_lens(m, data_lens(x, l))
    end
end

@testset "replacein" begin
    for t in list_traversal(x)
        x2 = deepcopy(x)
        x3 = replacein(x2, x2[t], x[t])
        @test x3[t] === x[t]
    end

    for t in list_traversal(m)
        m2 = deepcopy(m)
        m3 = replacein(m2, m2[t], m[t])
        @test m3[t] === m[t]
    end
end


text:
using DifferentialEquations
using Plots

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

δt = 0.1
tmax = 80.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax;

u0 = [990.0,10.0,0.0]; # S,I.R

p = [0.05,10.0,0.25]; # β,c,γ

lockdown_times = [10.0, 20.0]
condition(u,t,integrator) = t ∈ lockdown_times
function affect!(integrator)
    if integrator.t < lockdown_times[2]
        integrator.p[1] = 0.01
    else
        integrator.p[1] = 0.05
    end
end
cb = PresetTimeCallback(lockdown_times, affect!);

prob_ode = ODEProblem(sir_ode!,u0,tspan,p)

sol_ode = solve(prob_ode, callback = cb);

plot(sol_ode, label = ["S" "I" "R"], title = "Lockdown in a SIR model")
vline!(lockdown_times, c = :red, w = 2, label = "")


text:
using Test, Unitful, Plots
using Unitful: m, s, cm, DimensionError
using UnitfulRecipes

# Some helper functions to access the subplot labels and the series inside each test plot
xguide(plt, idx=length(plt.subplots)) = plt.subplots[idx].attr[:xaxis].plotattributes[:guide]
yguide(plt, idx=length(plt.subplots)) = plt.subplots[idx].attr[:yaxis].plotattributes[:guide]
zguide(plt, idx=length(plt.subplots)) = plt.subplots[idx].attr[:zaxis].plotattributes[:guide]
xseries(plt, idx=length(plt.series_list)) = plt.series_list[idx].plotattributes[:x]
yseries(plt, idx=length(plt.series_list)) = plt.series_list[idx].plotattributes[:y]
zseries(plt, idx=length(plt.series_list)) = plt.series_list[idx].plotattributes[:z]

@testset "plot(y)" begin
    y = rand(3)m

    @testset "no keyword argument" begin
        @test yguide(plot(y)) == "m"
        @test yseries(plot(y)) ≈ ustrip.(y)
    end

    @testset "ylabel" begin
        @test yguide(plot(y, ylabel="hello")) == "hello (m)"
        @test yguide(plot(y, ylabel=P"hello")) == "hello"
        @test yguide(plot(y, ylabel="")) == ""
    end

    @testset "yunit" begin
        @test yguide(plot(y, yunit=cm)) == "cm"
        @test yseries(plot(y, yunit=cm)) ≈ ustrip.(cm, y)
    end

    @testset "ylims" begin # Using all(lims .≈ lims) because of uncontrolled type conversions?
        @test all(ylims(plot(y, ylims=(-1,3))) .≈ (-1,3))
        @test all(ylims(plot(y, ylims=(-1m,3m))) .≈ (-1,3))
        @test all(ylims(plot(y, ylims=(-100cm,300cm))) .≈ (-1,3))
        @test all(ylims(plot(y, ylims=(-100cm,3m))) .≈ (-1,3))
    end

    @testset "keyword combinations" begin
        @test yguide(plot(y, yunit=cm, ylabel="hello")) == "hello (cm)"
        @test yseries(plot(y, yunit=cm, ylabel="hello")) ≈ ustrip.(cm, y)
        @test all(ylims(plot(y, yunit=cm, ylims=(-1,3))) .≈ (-1,3))
        @test all(ylims(plot(y, yunit=cm, ylims=(-1,3))) .≈ (-1,3))
        @test all(ylims(plot(y, yunit=cm, ylims=(-100cm,300cm))) .≈ (-100,300))
        @test all(ylims(plot(y, yunit=cm, ylims=(-100cm,3m))) .≈ (-100,300))
    end
end

@testset "plot(x,y)" begin
    x, y = randn(3)m, randn(3)s

    @testset "no keyword argument" begin
        @test xguide(plot(x,y)) == "m"
        @test xseries(plot(x,y)) ≈ ustrip.(x)
        @test yguide(plot(x,y)) == "s"
        @test yseries(plot(x,y)) ≈ ustrip.(y)
    end

    @testset "labels" begin
        @test xguide(plot(x, y, xlabel= "hello")) == "hello (m)"
        @test xguide(plot(x, y, xlabel=P"hello")) == "hello"
        @test yguide(plot(x, y, ylabel= "hello")) == "hello (s)"
        @test yguide(plot(x, y, ylabel=P"hello")) == "hello"
        @test xguide(plot(x, y, xlabel= "hello", ylabel= "hello")) == "hello (m)"
        @test xguide(plot(x, y, xlabel=P"hello", ylabel=P"hello")) == "hello"
        @test yguide(plot(x, y, xlabel= "hello", ylabel= "hello")) == "hello (s)"
        @test yguide(plot(x, y, xlabel=P"hello", ylabel=P"hello")) == "hello"
    end
end

@testset "With functions" begin
    x, y = randn(3), randn(3)
    @testset "plot(f, x) / plot(x, f)" begin
        f(x) = x^2
        @test plot(  f, x*m) isa Plots.Plot
        @test plot(x*m,   f) isa Plots.Plot
        g(x) = x*m # If the unit comes from the function only then it throws
        @test_throws DimensionError plot(x, g) isa Plots.Plot
        @test_throws DimensionError plot(g, x) isa Plots.Plot
    end
    @testset "plot(x, y, f)" begin
        f(x,y) = x*y
        @test plot(x*m, y*s, f) isa Plots.Plot
        @test plot(x*m,   y, f) isa Plots.Plot
        @test plot(  x, y*s, f) isa Plots.Plot
        g(x,y) = x*y*m # If the unit comes from the function only then it throws
        @test_throws DimensionError plot(x, y, g) isa Plots.Plot
    end
end

@testset "Moar plots" begin
    @testset "data as $dtype" for dtype in [:Vectors, :Matrices, Symbol("Vectors of vectors")]
        if dtype == :Vectors
            x, y, z = randn(10), randn(10), randn(10)
        elseif dtype == :Matrices
            x, y, z = randn(10,2), randn(10,2), randn(10,2)
        else
            x, y, z = [rand(10), rand(20)], [rand(10), rand(20)], [rand(10), rand(20)]
        end

        @testset "One array" begin
            @test plot(x*m)                    isa Plots.Plot
            @test plot(x*m, ylabel="x")        isa Plots.Plot
            @test plot(x*m, ylims=(-1,1))      isa Plots.Plot
            @test plot(x*m, ylims=(-1,1) .* m) isa Plots.Plot
            @test plot(x*m, yunit=u"km")       isa Plots.Plot
        end

        @testset "Two arrays" begin
            @test plot(x*m, y*s)                    isa Plots.Plot
            @test plot(x*m, y*s, xlabel="x")        isa Plots.Plot
            @test plot(x*m, y*s, xlims=(-1,1))      isa Plots.Plot
            @test plot(x*m, y*s, xlims=(-1,1) .* m) isa Plots.Plot
            @test plot(x*m, y*s, xunit=u"km")       isa Plots.Plot
            @test plot(x*m, y*s, ylabel="y")        isa Plots.Plot
            @test plot(x*m, y*s, ylims=(-1,1))      isa Plots.Plot
            @test plot(x*m, y*s, ylims=(-1,1) .* s) isa Plots.Plot
            @test plot(x*m, y*s, yunit=u"ks")       isa Plots.Plot
            @test scatter(x*m, y*s)                 isa Plots.Plot
            if dtype ≠ Symbol("Vectors of vectors")
                @test scatter(x*m, y*s, zcolor=z*(m/s)) isa Plots.Plot
            end
        end

        @testset "Three arrays" begin
            @test plot(x*m, y*s, z*(m/s))                        isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), xlabel="x")            isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), xlims=(-1,1))          isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), xlims=(-1,1) .* m)     isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), xunit=u"km")           isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), ylabel="y")            isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), ylims=(-1,1))          isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), ylims=(-1,1) .* s)     isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), yunit=u"ks")           isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), zlabel="z")            isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), zlims=(-1,1))          isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), zlims=(-1,1) .* (m/s)) isa Plots.Plot
            @test plot(x*m, y*s, z*(m/s), zunit=u"km/hr")        isa Plots.Plot
            @test scatter(x*m, y*s, z*(m/s))                     isa Plots.Plot
        end

        @testset "Unitful/unitless combinations" begin
            mystr(x::Array{<:Quantity}) = "Q"
            mystr(x::Array) = "A"
            @testset "plot($(mystr(xs)), $(mystr(ys)))" for xs in [x, x*m], ys in [y, y*s]
                @test plot(xs, ys) isa Plots.Plot
            end
            @testset "plot($(mystr(xs)), $(mystr(ys)), $(mystr(zs)))" for xs in [x, x*m], ys in [y, y*s], zs in [z, z*(m/s)]
                @test plot(xs, ys, zs) isa Plots.Plot
            end
        end
    end

    @testset "scatter(x::$(us[1]), y::$(us[2]))" for us in collect(Iterators.product(fill([1, u"m", u"s"], 2)...))
        x, y = rand(10)*us[1], rand(10)*us[2]
        @test scatter(x,y)  isa Plots.Plot
        @test scatter(x,y, markersize=x)  isa Plots.Plot
        @test scatter(x,y, line_z=x)  isa Plots.Plot
    end

    @testset "contour(x::$(us[1]), y::$(us[2]))" for us in collect(Iterators.product(fill([1, u"m", u"s"], 2)...))
        x, y = (1:0.01:2)*us[1], (1:0.02:2)*us[2]
        z = x' ./ y
        @test contour(x,y,z)  isa Plots.Plot
        @test contourf(x,y,z) isa Plots.Plot
    end

    @testset "ProtectedString" begin
        y = rand(10)*u"m"
        @test plot(y, label=P"meters") isa Plots.Plot
    end
end

@testset "Comparing apples and oranges" begin
    x1 = rand(10) * u"m"
    x2 = rand(10) * u"cm"
    x3 = rand(10) * u"s"
    plt = plot(x1)
    plt = plot!(plt, x2)
    @test yguide(plt) == "m"
    @test yseries(plt) ≈ ustrip.(x2) / 100
    @test_throws DimensionError plot!(plt, x3) # can't place seconds on top of meters!
end

@testset "Inset subplots" begin
    x1 = rand(10) * u"m"
    x2 = rand(10) * u"s"
    plt = plot(x1)
    plt = plot!(x2, inset=bbox(0.5, 0.5, 0.3, 0.3), subplot=2)
    @test yguide(plt,1) == "m"
    @test yguide(plt,2) == "s"
end


text:
module ChainPartitioners

using SparseArrays
using LinearAlgebra
using Suppressor
using Random
using DataStructures
using Requires

export DominanceSum
export StepwiseDominanceSum
export DominanceCount
export BinaryDominanceCount
export StepwiseDominanceCount
export RookSum
export RookCount
export BinaryRookCount
export dominancesum
export dominancesum!
export rooksum!
export dominancecount
export dominancecount!
export rookcount!

export EnvelopeMatrix
export rowenvelope

export NetCount
export netcount
export netcount!
export selfpincount
export selfpincount!
export selfnetcount
export selfnetcount!
export dianetcount
export dianetcount!

export partwise
export partwisecount!

export AbstractModel
export AbstractIncreasingModel
export AbstractDecreasingModel

export ConstrainedCost
export ConstrainedCostOracle

export AbstractWorkCost
export AffineWorkModel
export WorkOracle

export VertexCount

export AbstractConnectivityModel
export AffineConnectivityModel
export ConnectivityOracle
export AbstractSymmetricConnectivityModel
export AffineSymmetricConnectivityModel
export SymmetricConnectivityOracle
export AbstractMonotonizedSymmetricConnectivityModel
export AffineMonotonizedSymmetricConnectivityModel
export MonotonizedSymmetricConnectivityOracle
export AbstractPrimaryConnectivityModel
export AffinePrimaryConnectivityModel
export PrimaryConnectivityOracle
export AbstractSecondaryConnectivityModel
export AffineSecondaryConnectivityModel
export SecondaryConnectivityOracle
export AbstractSymmetricEdgeCutModel
export AffineSymmetricEdgeCutModel
export SymmetricEdgeCutOracle
export AbstractPrimaryEdgeCutModel
export AffinePrimaryEdgeCutModel
export PrimaryEdgeCutOracle
export AbstractSecondaryEdgeCutModel
export AffineSecondaryEdgeCutModel
export SecondaryEdgeCutOracle
export HyperedgeCutOracle
export AbstractHyperedgeCutModel
export AffineHyperedgeCutModel

export AbstractEnvelopeModel
export AffineEnvelopeModel
export EnvelopeOracle

export bottleneck_value
export total_value
export oracle_stripe
export bound_stripe

export partition_stripe
export partition_plaid
export EquiSplitter
export ReferenceBottleneckSplitter
export ReferenceTotalSplitter
export DynamicBottleneckSplitter
export DynamicTotalSplitter
export ConvexTotalSplitter
export ConcaveTotalSplitter
export BisectIndexBottleneckSplitter
export FlipBisectIndexBottleneckSplitter
export BisectCostBottleneckSplitter
export FlipBisectCostBottleneckSplitter
export LazyBisectCostBottleneckSplitter
export LazyFlipBisectCostBottleneckSplitter
export SymmetricPartitioner
export MagneticPartitioner
export GreedyBottleneckPartitioner

export pack_stripe
export pack_plaid
export DisjointPacker
export AlternatingPacker
export SymmetricPacker
export EquiChunker
export ReferenceTotalChunker
export DynamicTotalChunker
export ConvexTotalChunker
export ConcaveTotalChunker
export OverlapChunker
export StrictChunker

export BlockComponentCostModel
export ColumnBlockComponentCostModel

export MetisPartitioner
export metis_scheme_recursive_bisection
export metis_scheme_direct_k_way
export metis_objective_edge_cut
export metis_objective_communication_volume
export metis_coarsen_random
export metis_coarsen_sorted_heavy_edge
export metis_initialize_grow
export metis_initialize_random
export metis_initialize_edge
export metis_initialize_node
export metis_refine_fm
export metis_refine_greedy
export metis_refine_fm_two_sided
export metis_refine_fm_one_sided

export KaHyParPartitioner
export kahypar_configuration_communication_volume
export kahypar_configuration_hyperedge_cut

export DisjointPartitioner
export AlternatingPartitioner
export AlternatingNetPartitioner
export PermutingPartitioner

export SplitPartition
export DomainPartition
export MapPartition

export permute_stripe
export permute_plaid
export DomainPermutation
export MapPermutation
export IdentityPermuter
export CuthillMcKeePermuter
export GatesCuthillMcKeePermuter
export KryslCuthillMcKeePermuter
export SpectralPermuter
export MinDegreePermuter
export ColumnMinDegreePermuter
export perm

export pattern

include("util.jl")

abstract type AbstractHint end
struct NoHint <: AbstractHint end
struct RandomHint <: AbstractHint end
struct SparseHint <: AbstractHint end
struct StepHint <: AbstractHint end

include("Partitions.jl")

include("EnvelopeMatrices.jl")

include("Costs.jl")

include("SparsePrefixMatrices.jl")
include("SparseColorArrays.jl")

include("DynamicChunker.jl")
include("DynamicSplitter.jl")
include("ReferenceSplitter.jl")

include("WorkCosts.jl")
include("ConnectivityCosts.jl")
include("SymmetricEdgeCutCosts.jl")
include("SymmetricConnectivityCosts.jl")
include("MonotonizedSymmetricConnectivityCosts.jl")
include("PrimaryEdgeCutCosts.jl")
include("PrimaryConnectivityCosts.jl")
include("SecondaryEdgeCutCosts.jl")
include("SecondaryConnectivityCosts.jl")
include("HyperedgeCutCosts.jl")
include("EnvelopeCosts.jl")
include("BlockCosts.jl")

include("PartwiseCounts.jl")

include("EquiPartitioner.jl")
include("AlternatingPacker.jl")
include("ConvexTotalChunker.jl")
include("ConcaveTotalChunker.jl")
include("OverlapChunker.jl")
include("StrictChunker.jl")
include("BisectIndexBottleneckSplitter.jl")
include("BisectCostBottleneckSplitter.jl")
include("LazyBisectCostBottleneckSplitter.jl")
include("AlternatingPartitioner.jl")
include("MagneticPartitioner.jl")
include("Permutations.jl")
include("CuthillMcKeePermuter.jl")
include("PermutingPartitioner.jl")

@deprecate DynamicTotalChunker(f, w_max) DynamicTotalChunker(ConstrainedCost(f, VertexCount(), w_max))

function __init__()
    @require AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e" include("glue_AMD.jl")
    @require Laplacians = "6f8e5838-0efe-5de0-80a3-5fb4f8dbb1de" include("glue_Laplacians.jl")
    @require Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" include("glue_Metis.jl")
    @require KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880" include("glue_KaHyPar.jl")
    @require CuthillMcKee = "17f17636-5e38-52e3-a803-7ae3aaaf3da9" include("glue_CuthillMcKee.jl")
    @require SymRCM = "286e6d88-80af-4590-acc9-0001b223b9bd" include("glue_SymRCM.jl")
end

end # module


text:
# l,m -> plm array index
ip(l,m) = ( l*(l+1) >> 1 )+1+m  # bitshift >> 1 divides by 2

# l,m -> ylm array index
iq(l,m) = l*(l+1)+1+m

type SphericalHarmonicTable
    lmax::Int
    p::Array{Float64}  # associated legendre polynomial P_l^m(cos(θ)) / sin(θ)^m
    q::Array{Float64}  # spherical harmonic Y_{lm}(θ,ϕ)
    c::Array{Float64}  # zenith angular factor cos(mϕ) * sin(θ)^m
    s::Array{Float64}  # zenith angular factor sin(mϕ) * sin(θ)^m
end

function SphericalHarmonicTable(lmax)
    p = Array{Float64}(ip(lmax,lmax))
    q = Array{Float64}(iq(lmax,lmax))
    c = Array{Float64}(lmax+1)
    s = Array{Float64}(lmax+1)
    c[1] = 1.
    s[1] = 0.
    p[1] = 1.
    q[1] = 1/sqrt(4π)
    SphericalHarmonicTable(lmax,p,q,c,s)
end

function computesphharm(l,m,p,q,c,s)
    if m == 0
        q[iq(l,m)] = N(l,m)*p[ip(l,m)]
    else  # m > 0
        z = sqrt(2)*N(l,m)*p[ip(l,m)]
        q[iq(l,m)] = (-1)^m*c[m+1]*z
        q[iq(l,-m)] = -s[m+1]*z
    end
end

# compute sqrt((2l+1)*(l-m)! / (4π*(l+m)!))
function N(l::Int, m::Int)  # m > 0
    f = 1.
    for k in l-m+1:l+m; f *= k end  # f = factorial(l-m) / factorial(l+m)
    sqrt( (2l+1) / (4π*f) )
end

function Base.call(Y::SphericalHarmonicTable, x::Float64, y::Float64, z::Float64)
    lmax,p,q,c,s = Y.lmax,Y.p,Y.q,Y.c,Y.s

    for m in 0:lmax-1
        p[ip(m+1,m+1)] = -(2m+1)*p[ip(m,m)]  # P_{m+1,m+1} = -(2m+1)*P_{mm}
        s[m+2] = x*s[m+1] + y*c[m+1]
        c[m+2] = x*c[m+1] - y*s[m+1]
        computesphharm(m+1,m+1,p,q,c,s)
        p[ip(m+1,m)] = (2m+1)*z*p[ip(m,m)]  # P_{m+1,m} = (2m+1)*z*P_{mm}
        computesphharm(m+1,m,p,q,c,s)
        for l in m+1:lmax-1
            p[ip(l+1,m)] = ( (2l+1)*z*p[ip(l,m)] - (l+m)*p[ip(l-1,m)] ) / (l-m+1)
            computesphharm(l+1,m,p,q,c,s)
        end
    end
end

Base.call(Y::SphericalHarmonicTable, p::Point3{Float64}) = Y(p[1],p[2],p[3])

Base.getindex(Y::SphericalHarmonicTable, l::Int, m::Int) = Y.q[iq(l,m)]

max_stars_repo_name: johncsnyder/Element


text:
using Soss
using Test
using Weave
using Documenter

DocMeta.setdocmeta!(Soss, :DocTestSetup,
    quote
        using Soss
        using Random
        Random.seed!(3)
    end; recursive=true)

function buildREADME()
    weave("../README.jmd", doctype= "github", throw_errors=true, cache=:refresh, args=Dict(:seed => 6))
    return true
end

# write your own tests here
@testset "README" begin
    @test buildREADME()
end

@testset "Doctests" begin
    doctest(Soss)
end


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

@doc raw"""The node this Taint is attached to has the \"effect\" on any pod that does not tolerate the Taint.

    IoK8sApiCoreV1Taint(;
        effect=nothing,
        key=nothing,
        timeAdded=nothing,
        value=nothing,
    )

    - effect::String : Required. The effect of the taint on pods that do not tolerate the taint. Valid effects are NoSchedule, PreferNoSchedule and NoExecute.
    - key::String : Required. The taint key to be applied to a node.
    - timeAdded::IoK8sApimachineryPkgApisMetaV1Time : TimeAdded represents the time at which the taint was added. It is only written for NoExecute taints.
    - value::String : Required. The taint value corresponding to the taint key.
"""
mutable struct IoK8sApiCoreV1Taint <: SwaggerModel
    effect::Any # spec type: Union{ Nothing, String } # spec name: effect
    key::Any # spec type: Union{ Nothing, String } # spec name: key
    timeAdded::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgApisMetaV1Time } # spec name: timeAdded
    value::Any # spec type: Union{ Nothing, String } # spec name: value

    function IoK8sApiCoreV1Taint(;effect=nothing, key=nothing, timeAdded=nothing, value=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1Taint, Symbol("effect"), effect)
        setfield!(o, Symbol("effect"), effect)
        validate_property(IoK8sApiCoreV1Taint, Symbol("key"), key)
        setfield!(o, Symbol("key"), key)
        validate_property(IoK8sApiCoreV1Taint, Symbol("timeAdded"), timeAdded)
        setfield!(o, Symbol("timeAdded"), timeAdded)
        validate_property(IoK8sApiCoreV1Taint, Symbol("value"), value)
        setfield!(o, Symbol("value"), value)
        o
    end
end # type IoK8sApiCoreV1Taint

const _property_map_IoK8sApiCoreV1Taint = Dict{Symbol,Symbol}(Symbol("effect")=>Symbol("effect"), Symbol("key")=>Symbol("key"), Symbol("timeAdded")=>Symbol("timeAdded"), Symbol("value")=>Symbol("value"))
const _property_types_IoK8sApiCoreV1Taint = Dict{Symbol,String}(Symbol("effect")=>"String", Symbol("key")=>"String", Symbol("timeAdded")=>"IoK8sApimachineryPkgApisMetaV1Time", Symbol("value")=>"String")
Base.propertynames(::Type{ IoK8sApiCoreV1Taint }) = collect(keys(_property_map_IoK8sApiCoreV1Taint))
Swagger.property_type(::Type{ IoK8sApiCoreV1Taint }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCoreV1Taint[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1Taint }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1Taint[property_name]

function check_required(o::IoK8sApiCoreV1Taint)
    (getproperty(o, Symbol("effect")) === nothing) && (return false)
    (getproperty(o, Symbol("key")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1Taint }, name::Symbol, val)
end


text:
function _adapt_ϵ(logϵ, Hbar, logϵbar, da_stat, m, M_adapt, δ, μ;
    γ=0.05, t0=10, κ=0.75)

if m <= M_adapt

Hbar = (1.0 - 1.0 / (m + t0)) * Hbar + (1 / (m + t0)) * (δ - da_stat)
logϵ = μ - sqrt(m) / γ * Hbar
logϵbar = m^(-κ) * logϵ + (1 - m^(-κ)) * logϵbar

else

logϵ = logϵbar

end

return logϵ, Hbar, logϵbar

end


text:
module FractalTools

using DocStringExtensions
using Distributed
using PyCall
using LinearAlgebra
using GeometryBasics
using Clustering
using StaticArrays

import Makie
import Makie: @recipe
import GeometryBasics: Ngon
import Base: show, display
import StatsBase: sample, Weights

const MAXPREC = 1024    # Maximum precision for BigFloat arithmetic

function __init__()
    global spt = pyimport_conda("scipy.spatial", "scipy")
end

include("testfunctions.jl")
include("datagenerators.jl")
include("recipes.jl")
include("ifs.jl")
include("interpolation.jl")
include("integration.jl")

end # module


text:
@testset "iterable -> Tuple" begin
    t1 = pyconvert(Tuple, (1, 2))
    @test t1 === (1, 2)
    t2 = pyconvert(Tuple{Vararg{Int}}, (3, 4, 5))
    @test t2 === (3, 4, 5)
    t3 = pyconvert(Tuple{Int,Int}, (6, 7))
    @test t3 === (6, 7)
end


text:
#!/usr/bin/env julia

# publication11

using Base.Test
using MathProgBase
using Logging
using Drs

A = Float64[3 4; 6 1]
b = Float64[6, 3]
c = Float64[-2, -1]

s = linprog(c, A, '<', b, -Inf, Inf, DrsMathProgSolver())

@test s.status == :Optimal
@test s.objval == -13/7
@test s.sol == [2/7, 9/7]

max_stars_repo_name: abogushov/lplib

max_issues_repo_name: abogushov/lplib

max_forks_repo_name: abogushov/lplib


text:
using ClimaComms
using ClimaCommsMPI

include(joinpath(@__DIR__, "..", "..", "..", "test", "stencil.jl"))

stencil_test(ClimaCommsMPI.MPICommsContext())


text:
### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, KernelDensity, Plots
	Random.seed!(0)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing4.16"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	mu1, sigma1 = 10, 5
	mu2, sigma2 = 40, 12
	dist1, dist2 = Normal(mu1,sigma1), Normal(mu2,sigma2)
	p = 0.3
	mixRv() = (rand() <= p) ? rand(dist1) : rand(dist2)
	mixPDF(x) = p*pdf(dist1,x) + (1-p)*pdf(dist2,x)

	n = 2000
	data = [mixRv() for _ in 1:n]

	kdeDist = kde(data)
end

# ╔═╡ 45f535f0-250a-11eb-2a6b-03853e2d8eb1
begin
	xGrid = -20:0.1:80
	pdfKDE = pdf(kdeDist,xGrid)

	plot(xGrid, pdfKDE, c=:blue, label="KDE PDF")
	stephist!(data, bins=50, c=:black, normed=:true, label="Histogram")
	p1 = plot!(xGrid, mixPDF.(xGrid), c=:red, label="Underlying PDF",
		xlims=(-20,80), ylims=(0,0.035), legend=:topleft,
		xlabel="X", ylabel = "Density")
end

# ╔═╡ 45f5722a-250a-11eb-2840-a7d3dd59711b
begin
	hVals = [0.5,2,10]
	kdeS = [kde(data,bandwidth=h) for h in hVals]
	plot(xGrid, pdf(kdeS[1],xGrid), c = :green, label= "h=$(hVals[1])")
	plot!(xGrid, pdf(kdeS[2],xGrid), c = :blue, label= "h=$(hVals[2])")
	p2 = plot!(xGrid, pdf(kdeS[3],xGrid), c = :purple, label= "h=$(hVals[3])",
		xlims=(-20,80), ylims=(0,0.035), legend=:topleft,
		xlabel="X", ylabel = "Density")
	plot(p1,p2,size = (800,400))
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing4.16"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═45f535f0-250a-11eb-2a6b-03853e2d8eb1
# ╠═45f5722a-250a-11eb-2840-a7d3dd59711b
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


text:
function KS_solve_Emin_CG_smearing( pw::PWGrid, V_ionic, Focc, Nstates::Int;
                           psi0=nothing, Potentials0 = nothing, E_NN = 0.0,
                           α_t = 3e-5, NiterMax=1000, verbose=false, kT=0.01 )

    Ns = pw.Ns
    Npoints = prod(Ns)
    Ngwx = pw.gvecw.Ngwx

    if psi0 == nothing
        psi = randn(Ngwx,Nstates) + im*randn(Ngwx,Nstates)
        psi = ortho_gram_schmidt(psi)
    else
        psi = copy(psi0)
    end

    if Potentials0 == nothing
        Potentials = PotentialsT( V_ionic, zeros(Npoints), zeros(Npoints) )
        rho = calc_rho( pw, Focc, psi )
        Potentials.Hartree = real( G_to_R( Ns, Poisson_solve(pw, rho) ) )
        Potentials.XC = excVWN( rho ) + rho .* excpVWN( rho )
    else
        Potentials = PotentialsT( Potentials0.Ionic,
                                  Potentials0.Hartree,
                                  Potentials0.XC )
    end

    d = zeros(ComplexF64, Ngwx, Nstates)
    g_old = zeros(ComplexF64, Ngwx, Nstates)
    d_old = zeros(ComplexF64, Ngwx, Nstates)
    Kg = zeros(ComplexF64, Ngwx, Nstates)
    Kg_old = zeros(ComplexF64, Ngwx, Nstates)

    d_Haux = zeros(ComplexF64, Nstates, Nstates)
    g_Haux_old = zeros(ComplexF64, Nstates, Nstates)
    d_Haux_old = zeros(ComplexF64, Nstates, Nstates)
    Kg_Haux = zeros(ComplexF64, Nstates, Nstates)
    Kg_Haux_old = zeros(ComplexF64, Nstates, Nstates)

    β        = 0.0
    β_Haux   = 0.0
    Etot_old = 0.0
    Etot     = 0.0
    if( E_NN != nothing )
        Energies = EnergiesT( 0.0, 0.0, 0.0, 0.0, 0.0, E_NN )
    else
        Energies = EnergiesT( 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 )
    end

    Nelectrons = sum(Focc)
    @printf("Nelectrons = %18.10f\n", Nelectrons)

    # Get eigenvalues
    mu = psi' * op_H(pw, Potentials, psi)
    evals = sort(real(eigvals(mu)))
    #
    kT = 0.01
    Focc, E_fermi = calc_Focc(evals, Nelectrons, kT, is_spinpol=false)
    for ist = 1:Nstates
        @printf("%3d %18.10f %18.10f\n", ist, evals[ist], Focc[ist])
    end
    @printf("E_fermi = %18.10f\n", E_fermi)

    κ = 0.1

    Haux = zeros(ComplexF64,Nstates,Nstates)
    for ist = 1:Nstates
        Haux[ist,ist] = evals[ist]
    end

    for iter = 1:NiterMax

        g = calc_grad(pw, Potentials, Focc, psi)
        g_Haux = calc_grad_Haux(pw, Potentials, Focc, evals, psi, kT)

        Kg = Kprec(pw,g)
        Kg_Haux = κ*g_Haux

        if iter != 1
            #β = real(sum(conj(g).*Kg))/real(sum(conj(g_old).*Kg_old))
            β = real(sum(conj(g-g_old).*Kg))/real(sum(conj(g_old).*Kg_old))
            β_Haux = real(sum(conj(g_Haux-g_Haux_old).*Kg_Haux))/real(sum(conj(g_Haux_old).*Kg_Haux_old))
            #β = real(sum(conj(g-g_old).*Kg))/real(sum(conj(g-g_old).*d))
            #β = real(sum(conj(g).*Kg))/real(sum((g-g_old).*conj(d_old)))
        end
        if β < 0.0
            @printf("β is smaller than 0, setting it to zero\n")
            β = 0.0
        end
        if β_Haux < 0.0
            @printf("β_Haux is smaller than 0, setting it to zero\n")
            β_Haux = 0.0
        end

        d = -Kg + β * d_old
        d_Haux = -Kg_Haux + β_Haux * d_Haux_old

        psic = ortho_gram_schmidt(psi + α_t*d)
        rho = calc_rho( pw, Focc, psic )
        Potentials.Hartree = real( G_to_R( Ns, Poisson_solve(pw, rho) ) )
        Potentials.XC = excVWN( rho ) + rho .* excpVWN( rho )
        gt = calc_grad( pw, Potentials, Focc, psic )

        denum = real(sum(conj(g-gt).*d))
        if denum != 0.0
            α = abs( α_t*real(sum(conj(g).*d))/denum )
        else
            α = 0.0
        end

        # Update wavefunction
        psi = psi[:,:] + α*d[:,:]

        # Update potentials
        psi = ortho_gram_schmidt(psi)
        rho = calc_rho( pw, Focc, psi )

        Potentials.Hartree = real( G_to_R( Ns, Poisson_solve(pw, rho) ) )
        Potentials.XC = excVWN( rho ) + rho .* excpVWN( rho )

        Energies = calc_Energies(pw, Potentials, Focc, psi, Energies.NN)
        Entropies = calc_entropy(Focc, kT, is_spinpol=false)
        Etot = Energies.Total + Entropies

        diffE = abs(Etot-Etot_old)
        @printf("CG step %8d = %18.10f %10.7e\n", iter, Etot, diffE)
        if diffE < 1e-6
            @printf("CONVERGENCE ACHIEVED\n")
            break
        end

        g_old = copy(g)
        d_old = copy(d)
        Kg_old = copy(Kg)
        Etot_old = Etot
    end
    return psi, Energies, Potentials
    #
end


text:
#=
atom:
- Julia version:
- Author: Dan
- Date: 2021-06-01
=#

export
    Atom, getBonds, collectBonds, setFormalCharge, getElementSymbol, coords
using Distances: euclidean, sqeuclidean

"""
Standard representation for an Atom (e.g. in a Molecule). See also [`AbstractComposite`](@ref).
"""
mutable struct Atom <: AbstractAtom    #AbstractAtom inherits from AbstractComposite
    name_           ::Union{String,Nothing}
    parent_         ::Union{Nothing,AbstractComposite}
    previous_       ::Union{Nothing,AbstractComposite}
    next_           ::Union{Nothing,AbstractComposite}
    first_child_    ::Union{AbstractComposite, Nothing}
    last_child_     ::Union{AbstractComposite, Nothing}
    type_name_      ::Union{String,Nothing}
    element_        ::Union{Element,Nothing}
    radius_         ::Union{Float64,Nothing}
    type_           ::Union{UInt8,Nothing}
    number_of_bonds_::Union{UInt8,Nothing}
    formal_charge_  ::Union{Int64,Nothing}
    position_       ::Union{SVector{3,Float64},Nothing}
    charge_         ::Union{Float64,Nothing}
    velocity_       ::Union{SVector{3,Float64},Nothing}
    force_          ::Union{SVector{3,Float64},Nothing}
    occupancy_      ::Union{Float64, Nothing}
    serial_         ::Union{Int64,Nothing}
    temp_factor_    ::Union{Float64,Nothing}
    selected_       ::Bool
    bonds_          ::Dict{Atom, Bond}
    properties_     ::Vector{Tuple{String,UInt8}}

    Atom() = new(nothing,nothing,nothing,nothing,nothing,nothing,nothing,
    nothing,nothing,nothing,nothing,nothing,nothing,nothing,nothing,nothing,
    nothing,nothing,nothing,false,Dict{Atom, Bond}(),Vector{Tuple{String,UInt8}}())

        Atom(   name_           ,
                parent_         ,
                previous_       ,
                next_           ,
                first_child_    ,
                last_child_     ,
                type_name_      ,
                element_        ,
                radius_         ,
                type_           ,
                number_of_bonds_,
                formal_charge_  ,
                position_       ,
                charge_         ,
                velocity_       ,
                force_          ,
                occupancy_      ,
                serial_         ,
                temp_factor_

            )   = new(name_           ,
                    parent_         ,
                    previous_       ,
                    next_           ,
                    first_child_    ,
                    last_child_     ,
                    type_name_      ,
                    element_        ,
                    radius_         ,
                    type_           ,
                    number_of_bonds_,
                    formal_charge_  ,
                    position_       ,
                    charge_         ,
                    velocity_       ,
                    force_          ,
                    occupancy_      ,
                    serial_         ,
                    temp_factor_    ,
                    false,
                    BondsDict(),
                    Vector{Tuple{String,Int8}}()
            )
end

const BondsDict = Dict{Atom, Bond}

Atom(   atomname::String, x::Float64, y::Float64, z::Float64,
        elem::String, charge::Union{Float64,Nothing},
         occupancy::Float64,serial::Int64, temp_factor::Float64) = begin

        Atom(atomname,
            nothing,
            nothing,
            nothing,
            nothing,
            nothing,
            nothing,
            Element(capitalize(elem)),
            nothing,
            nothing,
            nothing,
            nothing,
            SA_F64[x,y,z],
            charge,
            nothing,
            nothing,
            occupancy,
            serial,
            temp_factor)
    end

getBonds(at::T) where T<:AbstractAtom = at.bonds_

"""
    collectBonds(node::AbstractComposite)
Returns a `Set` of all the `Bond`s of the atoms in the subtree rooted in `node`.
"""

"""
    coords(at::Atom)
Resturns a `SVector` with the Atoms coordinates.
"""
coords(at::Atom) = at.position_

collectBonds(node::AbstractComposite) = begin
    bonds = Set{Bond}()
    for at in collectAtoms(node)
        length(values(getBonds(at))) > 0 && push!(bonds, values(getBonds(at))...)
    end
    return collect(bonds)
end
viewBonds(node::AbstractComposite) = collectBonds(node)

"""
    getElementSymbol(at::Atom) -> String
Returns the Periodic Talble of Elements' symbol of the element of the atom.
"""
getElementSymbol(at::T) where T<:AbstractAtom = at.element_.symbol_

setFormalCharge(at::T, new_charge::Int64) where T<:AbstractAtom = begin at.formal_charge_ = new_charge end

Base.show(io::IO, at::Atom) = print(io, "Atom$(at.serial_)[",
    #"$( (isnothing(at.element_)) ? "-" : string(at.element_.symbol_) )|",
    "$( (isnothing(at.element_) || isnothing(at.element_.symbol_)) ? "-" : at.element_.symbol_)|",
    "$( (isnothing(at.parent_) || isnothing(at.parent_.name_)) ? "-" : at.parent_.name_ )]")


text:
using LOBPCG
using Base.Test

# write your own tests here


text:
using Pkg
# Add DiscreteAxis
Pkg.add(PackageSpec(url="https://github.com/xtalax/DiscreteAxis.jl", rev="master"))


text:
module Negotiations

# Dependencies
using YAML
using CSV
using DataFrames
using StatsBase
using Random
using Combinatorics
using SQLite
using ProgressMeter
using Chain
using Distributed

# Exports
export Agent
export ParameterSet
export Model
export Meeting
export initialize_db
export parameter_set_from_config
export load_database
export opinions_view
export setup_model
export simulate
export Rule
export BoundedConfidence
export ContinuousHomophily
export RealPolitics

# Source scripts
include("db_operations.jl")
include("data_views.jl")
include("config.jl")
include("model.jl")
include("rules.jl")
include("simulation.jl")
# include("convergence.jl")

end  # module


text:
##############################
# SELECTION SORT
##############################

function selection_sort(arr)
    for i = 1:length(arr), j = i:length(arr)
        if arr[j] < arr[i]
            arr[i], arr[j] = arr[j], arr[i]
        end
    end
end

# Get the array from user input
input = map(x->parse(Int, x), split(readline()))

# Sort the array
selection_sort(input)

# Print the array
for x in input
    print(x," ")
end

max_stars_repo_name: Nikola-Mircic/julia-cheatsheet

max_issues_repo_name: Nikola-Mircic/julia-cheatsheet

max_forks_repo_name: Nikola-Mircic/julia-cheatsheet


text:
"""
    This implements the Grids infrastructure
"""

# Grid infrastructure. Note that much of this is taken from Oceananigans

#####
##### Place-holder functions
#####

export
    Center, Face,
    AbstractTopology, Periodic, Bounded, Ghosted, topology,
    AbstractGrid,
    AbstractRectilinearGrid, RegularRectilinearCollocatedGrid,
    initialize_grid!, add_field

import Base: size, length, eltype, show, convert

#####
##### Abstract types
#####

"""
    AbstractField{X, Y, Z, A, G}

Abstract supertype for fields located at `(X, Y, Z)` with data stored in a container
of type `A`. The field is defined on a grid `G`.
"""
abstract type AbstractField{X, Y, Z, A, G} end

"""
    Center

A type describing the location at the center of a grid cell.
"""
struct Center end

"""
	xFace

A type describing the location at the x-face of a grid cell.
"""
struct xFace end

"""
	yFace

A type describing the location at the y-face of a grid cell.
"""
struct yFace end

"""
	zFace

A type describing the location at the z-face of a grid cell.
"""
struct zFace end

"""
	Vertex

A type describing the location at the vertex of a grid cell.
"""
struct Vertex end

"""
	Edge

A type describing the location at the edge of a grid cell.
"""
struct Edge end

"""
    AbstractTopology

Abstract supertype for grid topologies.
"""
abstract type AbstractTopology end

"""
    Periodic

Grid topology for periodic dimensions.
"""
struct Periodic <: AbstractTopology end

"""
    Bounded

Grid topology for bounded dimensions. These could be wall-bounded dimensions
or dimensions
"""
struct Bounded <: AbstractTopology end

"""
    Ghosted

Grid topology for bounded dimensions where we set boundary conditions using ghost points. These could be wall-bounded dimensions
or dimensions
"""
struct Ghosted <: AbstractTopology end

"""
    AbstractGrid{FT, TX, TY, TZ}

Abstract supertype for grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractGrid{FT, DIM, B} end

"""
    AbstractCollocatedRectilinearGrid{FT, TX, TY, TZ}

Abstract supertype for rectilinear grids with collocated points, with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractRectilinearCollocatedGrid{FT, DIM,  B} <: AbstractGrid{FT, DIM, B} end

"""
    fields = add_field(fields::NamedTuple, name::Symbol, newfield)

Adds a new field to the `fields` NamedTuple
"""
function add_field(fields::NamedTuple, name::Symbol, newfield)

    fields = merge(fields, NamedTuple{(name,)}( (newfield,)) );

    return fields
end

Base.eltype(::AbstractGrid{FT}) where FT = FT
Base.size(grid::AbstractGrid) = (grid.Nx, grid.Ny, grid.Nz)
Base.length(grid::AbstractGrid) = (grid.Lx, grid.Ly, grid.Lz)

# Include the different grid types
include("regular_collocated_rectilinear_grid.jl")


text:
module TrackerAD

using ..ThArrays

using Tracker
using Tracker: Tracked, Call, data, track

## Tensor
ThArrays.Tensor(x::Tracker.TrackedArray) = Tensor(data(x), requires_grad=true)

## TrackedTensor
struct TrackedTensor{T, N} <: AbstractArray{T, N}
    tracker::Tracked{Tensor{T, N}}
    data::Tensor{T, N}
    grad::Tensor{T, N}

    TrackedTensor{T, N}(
        t::Tracked{Tensor{T, N}},
        data::Tensor{T, N}) where {T, N} = new(t, data)
    TrackedTensor{T, N}(
        t::Tracked{Tensor{T, N}},
        data::Tensor{T, N}, grad::Tensor{T, N}) where {T, N} = new(t, data, grad)
end

Tracker.data(x::TrackedTensor) = x.data
Tracker.tracker(x::TrackedTensor) = x.tracker
Tracker.track(c::Tracker.Call, x::Tensor) = TrackedTensor(c, x)
Tracker.track(c::Tracker.Call, x::TrackedTensor) = TrackedTensor(c, data(x))

function TrackedTensor(c::Tracker.Call, t::Tensor{T, N}) where {T, N}
    TrackedTensor{T, N}(Tracker.Tracked{Tensor{T, N}}(c), t)
end

function TrackedTensor(c::Tracker.Call, t::Tensor{T, N}, d::Tensor{T, N}) where {T, N}
    TrackedTensor{T, N}(Tracker.Tracked{Tensor{T, N}}(c, d), t, d)
end

function TrackedTensor(t::Tensor)
    TrackedTensor(Tracker.Call(), t, ThC.zeros_like(t))
end

Base.eltype(x::Type{<:TrackedTensor{T}}) where T <: Real = TrackedTensor{T, 0}

function Base.show(io::IO, x::TrackedTensor)
  show(io, data(x))
  print(io, "(tracked Tensor)")
end

Base.copy(x::TrackedTensor) = x
Base.setindex!(xs::TrackedTensor, v, i...; kwargs...) =
    error("Can't differentiate `setindex!`")

## Fallthrough methods
for f in :[Base.size, Base.ndims, Base.collect].args
    @eval @inline $f(x::TrackedTensor, a...) = $f(data(x), a...)
end

## patches to Tracker.jl

Tracker.param(x::Tensor) = TrackedTensor(ThC.requires_grad!(x, true))
Tracker.init_grad(x::Tensor) = ThC.zeros_like(x)
Tracker.zero_grad!(x::Tensor) = (x .= 0)

"""
const __FORWARD_RESULT = IdDict{Any, Any}()
Tracker.collectmemaybe(x::TrackedTensor) = begin
    __FORWARD_RESULT[Tracker.tracker(x)] = x
    x
end

function Tracker.back(g::Tracker.Grads, x::Tracker.Tracked{Tensor{T, N}}, Δ) where {T, N}
    if haskey(__FORWARD_RESULT, x)
        ThAD.backward(data(__FORWARD_RESULT[x]), Tensor(float.(Δ)))
        delete!(__FORWARD_RESULT, x)
    end
    x.isleaf && (Tracker.accum!(g, x, Δ); return)
    ref = x.ref -= 1
    if ref > 0 || haskey(g, x)
        Tracker.accum!(g, x, Δ)
        ref == 0 && Tracker.back_(g, x.f, g[x])
    else
        ref == 0 && Tracker.back_(g, x.f, Δ)
    end
    return
end

# we use `_tr` to instead of the above patch now
"""
Tracker.collectmemaybe(x::TrackedTensor) = _tr(x)

## Switches
_th(x) = track(_th, x)
Tracker.@grad function _th(x)
    r = TrackedTensor(Tensor(x, requires_grad=true))
    r, (d) -> begin
        (ThC.ones_like(data(r)) .* d,)
    end
end

_th(x::Tracker.TrackedArray) = track(_th, x)
Tracker.@grad function _th(x::Tracker.TrackedArray)
    r = TrackedTensor(Tensor(data(x), requires_grad=true))
    r, (d) -> begin
        (ThC.ones_like(data(r)) .* d,)
    end
end

_tr(x) = track(_tr, x)
Tracker.@grad function _tr(x)
    x, (d) -> begin
        (ones(size(x)) .* d,)
    end
end

_tr(x::TrackedTensor{T, 0}) where {T} = track(_tr, x)
Tracker.@grad function _tr(x::TrackedTensor{T, 0}) where {T}
    r = convert(T, data(x))
    r, (d) -> begin
        ThAD.backward(data(x), Tensor(float(d)))
        (float(d),)
    end
end

_tr(x::TrackedTensor{T, N}) where {T, N} = track(_tr, x)
Tracker.@grad function _tr(x::TrackedTensor{T, N}) where {T, N}
    r = convert(Array, data(x))
    r, (d) -> begin
        ThAD.backward(data(x), Tensor(d))
        (ones(size(r)) .* d,)
    end
end

## Methods and Grads

Base.Broadcast.broadcasted(f, t::TrackedTensor, args...) = track(Base.Broadcast.broadcasted, f, t, args...)
Tracker.@grad function Base.Broadcast.broadcasted(f, t::TrackedTensor, args...)
    r = Base.Broadcast.broadcasted(f, data(t), data.(args)...)
    r, (d) -> begin
        grads = map(args) do arg
            (arg isa TrackedTensor) ? ThAD.get_grad(data(arg), d) : nothing
        end
        (nothing, ThAD.get_grad(data(t), d), grads...)
    end
end

macro grad_for_tensor(name)
    esc(quote
        $name(t::TrackedTensor, args...) = track($name, t, args...)
        Tracker.@grad function $name(t::TrackedTensor, args...)
            r = $name(data(t), data.(args)...)
            r, (d) -> begin
                grads = map(args) do arg
                    (arg isa TrackedTensor) ? ThAD.get_grad(data(arg), d) : nothing
                end
                (ThAD.get_grad(data(t), d), grads...)
            end
        end
        end)
end

#
# Methods in src/thc/thc.jl, can be extracted by the command:
# perl -n -e \
#   'if(m/import (Base\..*)/){ $i++; print "$1, "; print "\n" unless $i % 5;}' \
#   src/thc/thc.jl
#
for name in :[Base.abs, Base.acos, Base.all, Base.angle, Base.any,
              Base.argmax, Base.argmin, Base.asin, Base.atan, Base.cat,
              Base.ceil, Base.clamp, Base.clamp!, Base.coalesce, Base.conj,
              Base.cos, Base.cosh, Base.cumprod, Base.cumsum, Base.detach,
              Base.empty, Base.exp, Base.expm1, Base.fill!, Base.floor,
              Base.imag, Base.isfinite, Base.isnan, Base.log, Base.log10,
              Base.log1p, Base.log2, Base.max, Base.min, Base.mv,
              Base.ones, Base.prod, Base.put!, Base.rand, Base.randn,
              Base.range, Base.real, Base.repeat, Base.reshape, Base.resize!,
              Base.round, Base.sign, Base.sin, Base.sinh, Base.sort,
              Base.split, Base.sqrt, Base.sum, Base.tan, Base.tanh,
              Base.transpose, Base.trunc, Base.values, Base.view, Base.zeros,
              ].args
    @eval @grad_for_tensor($name)
end

#
# Methods in src/tensor.jl
#
for name in :[Base.getindex, Base.cat, Base.hcat, Base.vcat].args
    @eval @grad_for_tensor($name)
end

#
# Methods in src/common-methods.jl
#
for name in :[Base.:+, Base.:-, Base.:*, Base.:/, Base.div, Base.:^,
              Base.adjoint,].args
    @eval @grad_for_tensor($name)
end

end


text:
using HDF5: h5read, h5readattr
using DataFrames


text:
import Base: zero, one, +, -, *, /, //, ^,==, inv, iszero, isone,convert, show,rand

struct NTL_INIT_zz_p end # type for init
const _init_zz_p=NTL_INIT_zz_p()
"""
struct zz_p{T} is the type of finite fields

    ## Examples:
    julia> zz_p{7} # The type of finite fields of cardinality 7
    julia> zz_p{3}(4) # Element 4 in zz_p{3} which equals 1 in zz_p{3}
    julia> zz_p{3}() # Element 0 in zz_p{3}
"""
mutable struct zz_p{T}
    _rep::Int
    function zz_p{T}() where {T}
        new(0)
    end
    function zz_p{T}(a::Int, ::NTL_INIT_zz_p) where {T}
        new(a)
    end
    function zz_p{T}(a::Int) where {T}
        new(mod(a, T))
    end
end

## info of type zz_p{T}
"""
Generate the type zz_p{T} of prime fields

    ## Examples:
    julia> zz_p(3)
"""
function zz_p(a::Int)
    return zz_p{a}
end
function modulus(::Type{zz_p{T}}) where {T}
    return T
end
function zero(::Type{zz_p{T}}) where {T}
    return zz_p{T}(0)
end
function one(::Type{zz_p{T}}) where {T}
    return zz_p{T}(1, _init_zz_p)
end
## convert
function convert(::Type{zz_p{T}}, a::Int) where {T}
    return zz_p{T}(a)
end
function convert(::Type{Int}, a::zz_p{T}) where {T}
    return a._rep
end
## representation
function modulus(::Type{zz_p{T}}) where {T}
    return T
end
@inline function swap!(x::zz_p{T}, y::zz_p{T}) where {T}
    z = x._rep;x._rep = y._rep; y._rep =z;
    nothing
end

## inline funcitons
#const NTL_BITS_PER_LONG=Sys.WORD_SIZE
# @inline function cast_unsigned(a::Int)
#     return Core.bitcast(UInt, a)
# end
# @inline function sp_SignMask(a::Int)
#    return -(cast_unsigned(a) >> (NTL_BITS_PER_LONG-1));
# end
@inline function sp_CorrectExcess(a::Int, n::Int)
   return a-n >= 0 ? a-n : a
end
@inline function sp_CorrectDeficit(a::Int, n::Int)
   return a >= 0 ? a : a+n
end

## add
@inline function AddMod0(a::Int,b::Int, n::Int)::Int
    return Int(mod(Int128(a)+Int128(b),n))
 end
@inline function AddMod(a::Int,b::Int, n::Int)::Int
    r = a+b;
    return sp_CorrectExcess(r, n)
 end
function add!(z::zz_p{T} , x::zz_p{T}, y::zz_p{T}) where {T}
    z._rep = AddMod(x._rep,y._rep, T)
    return z
end
function add(x::zz_p{T}, y::zz_p{T}) where {T}
    Z=AddMod(x._rep,y._rep, T)
    return zz_p{T}(Z, _init_zz_p)
end
+(x::zz_p{T}, y::zz_p{T}) where {T} = add(x,y)
add(x::zz_p{T}, Y::Int) where {T} = add(x,convert(zz_p{T},Y))
add(X::Int, y::zz_p{T}) where {T} = add(convert(zz_p{T},X),y)

## sub
@inline function SubMod0(a::Int,b::Int, n::Int)::Int
    return mod(a-b, n)
 end
@inline function SubMod(a::Int,b::Int, n::Int)::Int
    r = a-b;
    return sp_CorrectDeficit(r, n)
 end
function sub!(z::zz_p{T} , x::zz_p{T}, y::zz_p{T}) where {T}
    z._rep = SubMod(x._rep,y._rep, T)
    return z
end
function sub(x::zz_p{T}, y::zz_p{T}) where {T}
    Z=SubMod(x._rep,y._rep, T)
    return zz_p{T}(Z, _init_zz_p)
end
-(x::zz_p{T}, y::zz_p{T}) where {T} = sub(x,y)
-(x::zz_p{T}, Y::Int) where {T} = sub(x,convert(zz_p{T},Y))
-(X::Int, y::zz_p{T}) where {T} = sub(convert(zz_p{T},X),y)

## negate
# x = -a
@inline function NegateMod(a::Int, n::Int)::Int
    # mod((0- a) , n)
      return SubMod(0, a, n)
   end
function negate!(x::zz_p{T}, a::zz_p{T}) where {T}
    x._rep=NegateMod(a._rep,T)
    return x
end
function negate(a::zz_p{T}) where {T}
    Z = NegateMod(a._rep,T)
    return zz_p{T}(Z, _init_zz_p)
end
-(a::zz_p{T}) where {T} = negate(a)

## mul
@inline function MulMod(a::Int, b::Int, n::Int)::Int
    Int(mod(Int128(a)*Int128(b),n))
end
function mul!(z::zz_p{T} , x::zz_p{T}, y::zz_p{T}) where {T}
    z._rep = MulMod(x._rep,y._rep, T)
    return z
end
function mul(x::zz_p{T}, y::zz_p{T}) where {T}
    Z=MulMod(x._rep,y._rep, T)
    return zz_p{T}(Z, _init_zz_p)
end

*(x::zz_p{T}, y::zz_p{T}) where {T} = mul(x,y)
*(x::zz_p{T}, Y::Int) where {T} = mul(x, convert(zz_p{T},Y))
*(X::Int, y::zz_p{T}) where {T} = mul(convert(zz_p{T},X), y)

## inv
# 1/a mod n
function InvMod(a::Int, n::Int)::Int
    d,u,v = gcdx(a, n)
    if d!= 1
        error("InvMod: inverse undefined")
    elseif u< 0
        return u+n
    end
    return u
end
function inv!(z::zz_p{T}, x::zz_p{T}) where {T}
    z._rep=InvMod(x._rep, T)
    return z
end
function inv(x::zz_p{T}) where {T}
    Z = InvMod(x._rep, T)
    return zz_p{T}(Z, _init_zz_p)
end

##div
function div!(z::zz_p{T} , x::zz_p{T}, y::zz_p{T}) where {T}
    z._rep = MulMod(x._rep, InvMod(y._rep, T), T)
    return z
end
function div(x::zz_p{T}, y::zz_p{T}) where {T}
    Z=MulMod(x._rep, InvMod(y._rep, T), T)
    return zz_p{T}(Z, _init_zz_p)
end
/(x::zz_p{T}, y::zz_p{T}) where {T} = div(x,y)
/(x::zz_p{T}, Y::Int) where {T} = div(x, convert(zz_p{T},Y))
/(X::Int, y::zz_p{T}) where {T} = div(convert(zz_p{T},X), y)

## comparison
@inline function iszero(x::zz_p{T}) where {T}
    return x._rep == 0
end
@inline function isone(x::zz_p{T}) where {T}
    return x._rep == 1
end
@inline function ==(x::zz_p{T}, y::zz_p{T}) where {T}
    return x._rep == y._rep
end
@inline function ==(x::zz_p{T}, Y::Int) where {T}
    return x._rep == Y
end
@inline function ==(X::Int, y::zz_p{T}) where {T}
    return X == y._rep
end
## random numbers
function rand(::Type{zz_p{T}}) where {T}
    return zz_p{T}(rand(Int))
end

## power x = a^e
function PowerMod(a::Int, e::Int , n::Int)
   if e < 0
      return InvMod(PowerMod(a,-e,n),n)
   end
    x::Int = 1;
    y::Int = a;
    while (e != 0)
        if (e & 1) != 0
                x = MulMod(x, y, n)
        end
        y = MulMod(y, y, n);
        e = e >> 1;
    end
   return x;
end

## show
show(io::IO, x::zz_p{T}) where {T} = print(io, "zz_", T, "(", x._rep ,")" )

@inline function power!(z::zz_p{T}, x::zz_p{T}, e::Int) where {T}
    z._rep= PowerMod(x._rep, e, T)
    return z
end

@inline function power(x::zz_p{T}, e::Int) where {T}
    Z= PowerMod(x._rep, e, T)
    return zz_p{T}(Z, _init_zz_p)
end

^(x::zz_p{T}, e::Int) where {T} = power(x,e)


text:
#=
# DG solver
=#
module DGSolver

export solve, nonlinear_solve

# See finch_import_symbols.jl for a list of all imported symbols.
import ..Finch: @import_finch_symbols
@import_finch_symbols()

using LinearAlgebra, SparseArrays

include("fe_boundary.jl"); # Can we use the CG versions here?
include("nonlinear.jl");
include("cg_matrixfree.jl");

# Things to do before and after each step(or stage) ########
function default_pre_step() end
function default_post_step() end

pre_step_function = default_pre_step;
post_step_function = default_post_step;

function set_pre_step(fun) global pre_step_function = fun; end
function set_post_step(fun) global post_step_function = fun; end
############################################################

function init_solver()
    global pre_step_function = default_pre_step;
    global post_step_function = default_post_step;
end

function linear_solve(var, bilinear, linear, face_bilinear, face_linear, stepper=nothing)
    if config.dimension > 2
        printerr("DG solver only available for 1D and 2D. Surface quadrature under construction.")
        #return;
    end

    if config.linalg_matrixfree
        return solve_matrix_free_sym(var, bilinear, linear, stepper);
        #return solve_matrix_free_asym(var, bilinear, linear, stepper);
    end
    if prob.time_dependent && !(stepper === nothing)
        assemble_t = @elapsed((A, b) = assemble(var, bilinear, linear, face_bilinear, face_linear, 0, stepper.dt));
        log_entry("Initial assembly took "*string(assemble_t)*" seconds");
        # global Amat = A;
        # global bvec = b;
        # display(A);
        # display(b);

        log_entry("Beginning "*string(stepper.Nsteps)*" time steps.");
        t = 0;
        sol = sol = get_var_vals(var);

        start_t = Base.Libc.time();
        assemble_t = 0;
        linsolve_t = 0;
        last2update = 0;
        last10update = 0;
        print("Time stepping progress(%): 0");
        for i=1:stepper.Nsteps
            stupidjulia=sol;
            if stepper.stages > 1
                # LSRK4 is a special case, low storage
                if stepper.type == LSRK4
                    # Low storage RK4:
                    # p0 = u
                    #   ki = ai*k(i-1) + dt*f(p(i-1), t+ci*dt)
                    #   pi = p(i-1) + bi*ki
                    # u = p5

                    # store the original values, x, in a temp vector, so they aren't lost.
                    tmppi = get_var_vals(var);
                    tmpki = zeros(size(b));
                    for rki=1:stepper.stages
                        rktime = t + stepper.c[rki]*stepper.dt;
                        # p(i-1) is currently in u
                        pre_step_function();

                        assemble_t += @elapsed(b = assemble(var, nothing, linear, nothing, face_linear, rktime, stepper.dt; rhs_only = true));

                        linsolve_t += @elapsed(sol = A\b);
                        if rki == 1 # because a1 == 0
                            tmpki = stepper.dt .* sol;
                        else
                            tmpki = stepper.a[rki].*tmpki + stepper.dt .* sol;
                        end
                        tmppi = tmppi + stepper.b[rki].*tmpki

                        # Need to apply boundary conditions now to tmppi
                        tmppi = apply_boundary_conditions_rhs_only(var, tmppi, rktime);

                        place_sol_in_vars(var, tmppi, stepper);
                        post_step_function();
                    end

                else
                    # Explicit multi-stage methods:
                    # x = x + dt*sum(bi*ki)
                    # ki = rhs(t+ci*dt, x+dt*sum(aij*kj)))   j < i

                    # will hold the final result
                    result = get_var_vals(var);
                    # will be placed in var.values for each stage
                    tmpvals = result;
                    # Storage for each stage
                    tmpki = zeros(length(b), stepper.stages);
                    for stage=1:stepper.stages
                        stime = t + stepper.c[stage]*stepper.dt;
                        pre_step_function();

                        assemble_t += @elapsed(b = assemble(var, nothing, linear, nothing, face_linear, stime, stepper.dt; rhs_only = true));

                        linsolve_t += @elapsed(tmpki[:,stage] = A\b);

                        tmpvals = result;
                        for j=1:(stage-1)
                            if stepper.a[stage, j] > 0
                                tmpvals += stepper.dt * stepper.a[stage, j] .* tmpki[:,j];
                            end
                        end

                        # Need to apply boundary conditions now
                        tmpvals = apply_boundary_conditions_rhs_only(var, tmpvals, stime);

                        place_sol_in_vars(var, tmpvals, stepper);
                        post_step_function();
                    end
                    for stage=1:stepper.stages
                        result += stepper.dt * stepper.b[stage] .* tmpki[:, stage];
                    end

                    result = apply_boundary_conditions_rhs_only(var, result, t + stepper.dt);
                    place_sol_in_vars(var, result, stepper);
                    post_step_function();
                end
            else
                ##### TODO ###### Update needed
                pre_step_function();
                assemble_t += @elapsed(b = assemble(var, nothing, linear, nothing, face_linear, t, stepper.dt; rhs_only = true));
                linsolve_t += @elapsed(sol = A\b);

                place_sol_in_vars(var, sol, stepper);
                post_step_function();
            end

            #println("b(bdry): "*string(b[1])*", "*string(b[end]));

            t += stepper.dt;

            progressPercent = Int(floor(i*100.0/stepper.Nsteps));
            if progressPercent - last2update >= 2
                last2update = progressPercent;
                if progressPercent - last10update >= 10
                    print(progressPercent);
                    last10update = progressPercent;
                else
                    print(".");
                end
            end
        end
        println("");
        end_t = Base.Libc.time();

        log_entry("Stepping took "*string(end_t-start_t)*" seconds. ("*string(assemble_t)*" for assembly, "*string(linsolve_t)*" for linear solve)");
        #display(sol);
		# outfile = "linear_sol.txt"
		# open(outfile, "w") do f
  		# 	for ii in sol
    	# 		println(f, ii)
  		# 	end
		# end # the file f is automatically closed after this block finishes
        return sol;

    else
        assemble_t = @elapsed((A, b) = assemble(var, bilinear, linear, face_bilinear, face_linear));
        # uncomment to look at A
        # global Amat = A;

        sol_t = @elapsed(sol = A\b);

        log_entry("Assembly took "*string(assemble_t)*" seconds");
        log_entry("Linear solve took "*string(sol_t)*" seconds");
        #display(A);
        #display(b);
        #display(sol);
        return sol;
    end
end

function nonlinear_solve(var, nlvar, bilinear, linear, stepper=nothing)
    if prob.time_dependent && !(stepper === nothing)
        #TODO time dependent coefficients

        log_entry("Beginning "*string(stepper.Nsteps)*" time steps.");
        t = 0;
        start_t = Base.Libc.time();
        nl = nonlinear(100, 1e-9, 1e-9);
        init_nonlinear(nl, var, nlvar, bilinear, linear);
        for i=1:stepper.Nsteps
			println("solve for time step ", i);
			newton(nl, assemble, assemble_rhs_only, nlvar, t, stepper.dt);
            t += stepper.dt;
			nlvar[4].values = copy(nlvar[1].values);
			nlvar[5].values = copy(nlvar[2].values);
			if (i % 10 ==0)
				outfile = string("LD_u_",i);
				open(outfile, "w") do f
					for ii=1:length(nlvar[1].values)
						println(f, nlvar[1].values[ii])
					end
					close(f)
				end
				outfile = string("LD_v_",i);
				open(outfile, "w") do f
					for ii=1:length(nlvar[2].values)
						println(f, nlvar[2].values[ii])
					end
					close(f)
				end
			end
        end
        end_t = Base.Libc.time();

        log_entry("Solve took "*string(end_t-start_t)*" seconds");
        #display(sol);
        #return nlvar.values;
        return [];

    else
        start_t = Base.Libc.time();
        nl = nonlinear(100, 1e-12, 1e-12);
        init_nonlinear(nl, var, nlvar, bilinear, linear);
        newton(nl, assemble, assemble_rhs_only, nlvar);
        end_t = Base.Libc.time();

        log_entry("Solve took "*string(end_t-start_t)*" seconds");

        #return nlvar.values;
        return [];
    end
end

# ########################
# # tmp for testing
# function the_desired(args)

# end
# #########################

# assembles the A and b in Au=b
function assemble(var, bilinear, linear, face_bilinear, face_linear, t=0.0, dt=0.0; rhs_only = false)
    Np = refel.Np;
    nel = size(grid_data.loc2glb,2);
    N1 = size(grid_data.allnodes,2);
    multivar = typeof(var) <: Array;
    maxvarindex = 0;
    if multivar
        # multiple variables being solved for simultaneously
        dofs_per_node = 0;
        var_to_dofs = [];
        for vi=1:length(var)
            tmp = dofs_per_node;
            dofs_per_node += length(var[vi].symvar);
            push!(var_to_dofs, (tmp+1):dofs_per_node);
            maxvarindex = max(maxvarindex,var[vi].index);
        end
    else
        # one variable
        dofs_per_node = length(var.symvar);
        maxvarindex = var.index;
    end
    Nn = dofs_per_node * N1;

    b = zeros(Nn);

    if !rhs_only
        AI = zeros(Int, nel*dofs_per_node*Np*dofs_per_node*Np);
        AJ = zeros(Int, nel*dofs_per_node*Np*dofs_per_node*Np);
        AV = zeros(nel*dofs_per_node*Np*dofs_per_node*Np);
    end

    # Stiffness and mass are precomputed for uniform grid meshes
    precomputed_mass_stiffness = config.mesh_type == UNIFORM_GRID && config.geometry == SQUARE
    if precomputed_mass_stiffness
        wdetj = refel.wg .* geo_factors.detJ[1];
        J = geo_factors.J[1];
        if config.dimension == 1
            (RQ1, RD1) = build_deriv_matrix(refel, J);
            TRQ1 = RQ1';
            stiffness = [(TRQ1 * diagm(wdetj) * RQ1)];
        elseif config.dimension == 2
            (RQ1, RQ2, RD1, RD2) = build_deriv_matrix(refel, J);
            (TRQ1, TRQ2) = (RQ1', RQ2');
            stiffness = [(TRQ1 * diagm(wdetj) * RQ1) , (TRQ2 * diagm(wdetj) * RQ2)];
        else
            (RQ1, RQ2, RQ3, RD1, RD2, RD3) = build_deriv_matrix(refel, J);
            (TRQ1, TRQ2, TRQ3) = (RQ1', RQ2', RQ3');
            stiffness = [(TRQ1 * diagm(wdetj) * RQ1) , (TRQ2 * diagm(wdetj) * RQ2) , (TRQ3 * diagm(wdetj) * RQ3)];
        end
        mass = (refel.Q)' * diagm(wdetj) * refel.Q;
    else
        stiffness = 0;
        mass = 0;
    end

    loop_time = Base.Libc.time();
    # Elemental loop follows elemental ordering
    for ei=1:nel
        eid = elemental_order[ei];
        loc2glb = grid_data.loc2glb[:,eid];           # global indices of this element's nodes for extracting values from var arrays

        if !rhs_only
            Astart = (eid-1)*Np*dofs_per_node*Np*dofs_per_node + 1; # The segment of AI, AJ, AV for this element

        end

        volargs = (var, eid, 0, grid_data, geo_factors, refel, t, dt, stiffness, mass);

        if dofs_per_node == 1
            linchunk = linear.func(volargs);  # get the elemental linear part
            b[loc2glb] .+= linchunk;

            if !rhs_only
                bilinchunk = bilinear.func(volargs); # the elemental bilinear part
                #A[loc2glb, loc2glb] .+= bilinchunk;         # This will be very inefficient for sparse A
                for jj=1:Np
                    offset = Astart - 1 + (jj-1)*Np;
                    for ii=1:Np
                        AI[offset + ii] = loc2glb[ii];
                        AJ[offset + ii] = loc2glb[jj];
                        AV[offset + ii] = bilinchunk[ii, jj];
                    end
                end
            end

        elseif typeof(var) == Variable
            # only one variable, but more than one dof
            linchunk = linear.func(volargs);
            insert_linear!(b, linchunk, loc2glb, 1:dofs_per_node, dofs_per_node);

            if !rhs_only
                bilinchunk = bilinear.func(volargs);
                insert_bilinear!(AI, AJ, AV, Astart, bilinchunk, loc2glb, 1:dofs_per_node, dofs_per_node);
            end

        else
            linchunk = linear.func(volargs);
            insert_linear!(b, linchunk, loc2glb, 1:dofs_per_node, dofs_per_node);

            if !rhs_only
                bilinchunk = bilinear.func(volargs);
                insert_bilinear!(AI, AJ, AV, Astart, bilinchunk, loc2glb, 1:dofs_per_node, dofs_per_node);
            end
        end
    end

    # surface assembly for dg
    if !(face_linear===nothing) && (rhs_only || !(face_bilinear===nothing))
        for fid = 1:size(grid_data.face2glb,3)
            frefelind = [grid_data.faceRefelInd[1,fid], grid_data.faceRefelInd[2,fid]]; # refel based index of face in both elements

            face2glb = grid_data.face2glb[:,:,fid]; # local to global for face
            if frefelind[2] > 0 # not a boundary face
                flocal = [refel.face2local[frefelind[1]], refel.face2local[frefelind[2]]]; # local indices of face in both elements
            else # a boundary face
                face2glb[:,2] = face2glb[:,1]; # same global index for inside and out
                frefelind[2] = frefelind[1];
                flocal = [refel.face2local[frefelind[1]], refel.face2local[frefelind[1]]]; # local indices of face in both elements
            end
            eid = grid_data.face2element[1,fid]; # element on side 1

            faceBID = grid_data.facebid[fid]; # BID of face (0 for interior)

            surfargs = (var, eid, fid, grid_data, geo_factors, refel, t, dt);

            if dofs_per_node == 1
                linchunk = face_linear.func(surfargs);  # get the elemental linear part
                #linchunk = the_desired(rhsargs);
                # linchunk[1] = linchunk[1][flocal[1]];
                # linchunk[2] = linchunk[2][flocal[2]];
                if faceBID == 0
                    b[face2glb[:,1]] .+= linchunk[1];
                    b[face2glb[:,2]] .+= linchunk[2];
                elseif grid_data.face2element[1,fid] == 0
                    b[face2glb[:,2]] .+= linchunk[2];
                else
                    b[face2glb[:,1]] .+= linchunk[1];
                end

                if !rhs_only
                    bilinchunk = face_bilinear.func(surfargs); # the elemental bilinear part

                    if typeof(bilinchunk[1]) <: Number
                        bilinchunk = [[bilinchunk[1]], [bilinchunk[2]], [bilinchunk[3]], [bilinchunk[4]]];
                    else
                        # bilinchunk[1] = bilinchunk[1][flocal[1],flocal[1]];
                        # bilinchunk[2] = bilinchunk[2][flocal[1],flocal[2]];
                        # bilinchunk[3] = bilinchunk[3][flocal[2],flocal[1]];
                        # bilinchunk[4] = bilinchunk[4][flocal[2],flocal[2]];
                    end
                    #bilinchunk = temporary_lhs_func(lhsargs);
                    #println(bilinchunk)
                    nfp = size(face2glb,1);
                    for jj=1:nfp
                        append!(AI, face2glb[:,1]);
                        append!(AJ, face2glb[jj,1]*ones(Int, nfp));
                        append!(AV, bilinchunk[1][:,jj]);

                        append!(AI, face2glb[:,1]);
                        append!(AJ, face2glb[jj,2]*ones(Int, nfp));
                        append!(AV, bilinchunk[2][:,jj]);

                        append!(AI, face2glb[:,2]);
                        append!(AJ, face2glb[jj,1]*ones(Int, nfp));
                        append!(AV, bilinchunk[3][:,jj]);

                        append!(AI, face2glb[:,2]);
                        append!(AJ, face2glb[jj,2]*ones(Int, nfp));
                        append!(AV, bilinchunk[4][:,jj]);
                    end
                end
            end
        end
    end

    loop_time = Base.Libc.time() - loop_time;

    if !rhs_only
        # Build the sparse A. Uses default + to combine overlaps
        A = sparse(AI, AJ, AV);
    end

    # Boundary conditions
    bc_time = Base.Libc.time();
    if rhs_only
        b = apply_boundary_conditions_rhs_only(var, b, t);
    else
        (A, b) = apply_boundary_conditions_lhs_rhs(var, A, b, t);
    end
    bc_time = Base.Libc.time() - bc_time;

    if rhs_only
        return b;

    else
        log_entry("Elemental loop time:     "*string(loop_time));
        log_entry("Boundary condition time: "*string(bc_time));

        return (A, b);
    end
end

# Inset the single dof into the greater construct
function insert_linear!(b, bel, glb, dof, Ndofs)
    # group nodal dofs
    for d=1:length(dof)
        ind = glb.*Ndofs .- (Ndofs-dof[d]);
        ind2 = ((d-1)*length(glb)+1):(d*length(glb));

        b[ind] = b[ind] + bel[ind2];
    end
end

function insert_bilinear!(AI, AJ, AV, Astart, ael, glb, dof, Ndofs)
    Np = length(glb);
    # group nodal dofs
    for dj=1:length(dof)
        indj = glb.*Ndofs .- (Ndofs-dof[dj]);
        indj2 = ((dj-1)*Np+1):(dj*Np);
        for di=1:length(dof)
            indi = glb.*Ndofs .- (Ndofs-dof[di]);
            indi2 = ((di-1)*Np+1):(di*Np);

            #a[indi, indj] = a[indi, indj] + ael[indi2, indj2];

            for jj=1:Np
                offset = Astart + (jj-1 + Np*(dj-1))*Np*Ndofs + Np*(di-1) - 1;
                for ii=1:Np
                    AI[offset + ii] = indi[ii];
                    AJ[offset + ii] = indj[jj];
                    AV[offset + ii] = ael[indi2[ii], indj2[jj]];
                end
            end
        end
    end

end

function place_sol_in_vars(var, sol, stepper)
    # place the values in the variable value arrays
    if typeof(var) <: Array
        tmp = 0;
        totalcomponents = 0;
        for vi=1:length(var)
            totalcomponents = totalcomponents + length(var[vi].symvar);
        end
        for vi=1:length(var)
            components = length(var[vi].symvar);
            for compi=1:components
                if stepper.type == EULER_EXPLICIT
                    var[vi].values[compi,:] += sol[(compi+tmp):totalcomponents:end];
                else
                    var[vi].values[compi,:] = sol[(compi+tmp):totalcomponents:end];
                end
                tmp = tmp + 1;
            end
        end
    else
        components = length(var.symvar);
        for compi=1:components
            if stepper.type == EULER_EXPLICIT
                var.values[compi,:] += sol[compi:components:end];
            else
                var.values[compi,:] = sol[compi:components:end];
            end
        end
    end
end

function get_var_vals(var)
    # place the variable values in a vector
    if typeof(var) <: Array
        tmp = 0;
        totalcomponents = 0;
        for vi=1:length(var)
            totalcomponents = totalcomponents + length(var[vi].symvar);
        end
        vect = zeros(totalcomponents * length(var[1].values[1,:]));
        for vi=1:length(var)
            components = length(var[vi].symvar);
            for compi=1:components
                vect[(compi+tmp):totalcomponents:end] = var[vi].values[compi,:];
                tmp = tmp + 1;
            end
        end
    else
        components = length(var.symvar);
        vect = zeros(components * length(var.values[1,:]));
        for compi=1:components
            vect[compi:components:end] = var.values[compi,:];
        end
    end

    return vect;
end

end #module

max_stars_repo_name: paralab/Finch

max_issues_repo_name: paralab/Finch

max_forks_repo_name: paralab/Finch