text:
module SymbolServer

export SymbolServerProcess
export getstore

using Serialization, Pkg

include("clientprocess/from_static_lint.jl")

mutable struct SymbolServerProcess
    process::Base.Process

    function SymbolServerProcess(;environment=nothing, depot=nothing)
        jl_cmd = joinpath(Sys.BINDIR, Base.julia_exename())
        client_process_script = joinpath(@__DIR__, "clientprocess", "clientprocess_main.jl")

        env_to_use = copy(ENV)

        if depot!==nothing
            if depot==""
                delete!(env_to_use, "JULIA_DEPOT_PATH")
            else
                env_to_use["JULIA_DEPOT_PATH"] = depot
            end
        end

        p = if environment===nothing
            open(Cmd(`$jl_cmd $client_process_script`, env=env_to_use), read=true, write=true)
        else
            open(Cmd(`$jl_cmd --project=$environment $client_process_script`, dir=environment, env=env_to_use), read=true, write=true)
        end

        return new(p)
    end
end

function request(server::SymbolServerProcess, message::Symbol, payload)
    serialize(server.process, (message, payload))
    ret_val = deserialize(server.process)
    !(ret_val isa Tuple{Symbol,<:Any}) && error("Invalid response:\n", ret_val)
    return ret_val
end

function load_store_from_disc(file)
    io = open(file)
    store = deserialize(io)
    close(io)
    typeof(store) != ModuleStore && @info "WARNING: Type mismatch in loaded store"
    return store
end

# Public API

function getstore(server::SymbolServerProcess)
    storedir = abspath(joinpath(@__DIR__, "..", "store"))
    depot = Dict{String,ModuleStore}()

    if !isfile(joinpath(storedir, "Base.jstore")) || !isfile(joinpath(storedir, "Core.jstore"))
        get_core_package(server)
    end
    isfile(joinpath(storedir, "Base.jstore")) || error("Couldn't create Base store")
    isfile(joinpath(storedir, "Core.jstore")) || error("Couldn't create Core store")

    depot["Base"] = load_store_from_disc(joinpath(storedir, "Base.jstore"))
    depot["Core"] = load_store_from_disc(joinpath(storedir, "Core.jstore"))

    installed_pkgs_in_env = get_installed_packages_in_env(server)
    all_pkgs_in_env = get_all_packages_in_env(server)

    for (pkg_name, uuid) in installed_pkgs_in_env
        if isfile(joinpath(storedir, "$uuid.jstore"))
            depot[pkg_name] = load_store_from_disc(joinpath(storedir, "$uuid.jstore"))
        else
            load_package(server, (pkg_name => uuid))
            if isfile(joinpath(storedir, "$uuid.jstore"))
                depot[pkg_name] = load_store_from_disc(joinpath(storedir, "$uuid.jstore"))
            end
        end
    end
    for (pkg_name, uuids) in all_pkgs_in_env
        pkg_name in keys(depot) && continue
        uuid = first(uuids) # will need fix for multiple package versions within an env
        if isfile(joinpath(storedir, "$uuid.jstore"))
            depot[pkg_name] = load_store_from_disc(joinpath(storedir, "$uuid.jstore"))
        else
            # search for uuid within installed_pkgs_in_env dependencies
        end
    end

    return depot
end

function Base.kill(server::SymbolServerProcess)
    serialize(server.process, (:close, nothing))
    # kill(s.process)
end

function get_core_package(server::SymbolServerProcess)
    status, payload = request(server, :get_core_packages, nothing)
    if status == :success
        return payload
    else
        error(payload)
    end
end

function get_installed_packages_in_env(server::SymbolServerProcess)
    status, payload = request(server, :get_installed_packages_in_env, nothing)
    if status == :success
        return payload
    else
        error(payload)
    end
end

function get_all_packages_in_env(server::SymbolServerProcess)
    status, payload = request(server, :get_all_packages_in_env, nothing)
    if status == :success
        return payload
    else
        error(payload)
    end
end

function load_package(server::SymbolServerProcess, pkg)
    status, payload = request(server, :load_package, pkg)
    if status == :success
        return payload
    else
        error(payload)
    end
end

end # module


text:
"""
The `DataSet` type is a versatile container for storing data. Typically, it is constructed by passing it three vectors `x`, `y`, `sigma` where the components of `sigma` quantify the standard deviation associated with each y-value.
Alternatively, a full covariance matrix can be supplied for the `ydata` instead of a vector of standard deviations. The contents of a `DataSet` `DS` can later be accessed via `xdata(DS)`, `ydata(DS)`, `ysigma(DS)`.

Examples:

In the simplest case, where all data points are mutually independent and have a single ``x``-component and a single ``y``-component each, a `DataSet` consisting of four points can be constructed via
```julia
DataSet([1,2,3,4], [4,5,6.5,7.8], [0.5,0.45,0.6,0.8])
```
or alternatively by
```julia
using LinearAlgebra
DataSet([1,2,3,4], [4,5,6.5,7.8], Diagonal([0.5,0.45,0.6,0.8].^2))
```
where the diagonal covariance matrix in the second line is equivalent to the vector of standard deviations supplied in the first line.

For measurements with multiple components, it is also possible to enter them as a `Matrix` where the columns correspond to the respective components.
```julia
DataSet([0, 0.5, 1], [1 100; 2 103; 3 108], [0.5 8; 0.4 5; 0.6 10])
```
Note that if the uncertainty matrix is square, it may be falsely interpreted as a covariance matrix instead of as the columnwise specification of standard deviations.

More generally, if a dataset consists of ``N`` points where each ``x``-value has ``n`` many components and each ``y``-value has ``m`` many components, this can be specified to the `DataSet` constructor via a tuple ``(N,n,m)`` in addition to the vectors `x`, `y` and the covariance matrix.
For example:
```julia
X = [0.9, 1.0, 1.1, 1.9, 2.0, 2.1, 2.9, 3.0, 3.1, 3.9, 4.0, 4.1]
Y = [1.0, 5.0, 4.0, 8.0, 9.0, 13.0, 16.0, 20.0]
Cov = Diagonal([2.0, 4.0, 2.0, 4.0, 2.0, 4.0, 2.0, 4.0])
dims = (4, 3, 2)
DS = DataSet(X, Y, Cov, dims)
```
In this case, `X` is a vector consisting of the concatenated x-values (with 3 components each) for 4 different data points.
The values of `Y` are the corresponding concatenated y-values (with 2 components each) of said 4 data points. Clearly, the covariance matrix must therefore be a positive-definite ``(m \\cdot N) \\times (m \\cdot N)`` matrix.
"""
struct DataSet <: AbstractDataSet
    x::AbstractVector
    y::AbstractVector
    InvCov::AbstractMatrix
    dims::Tuple{Int,Int,Int}
    logdetInvCov::Real
    WoundX::Union{AbstractVector,Nothing}
    xnames::AbstractVector{String}
    ynames::AbstractVector{String}
    function DataSet(DF::Union{DataFrame,AbstractMatrix})
        size(DF,2) > 3 && throw("Unclear dimensions of input $DF.")
        DataSet(ToCols(Matrix(DF))...)
    end
    function DataSet(x::AbstractArray, y::AbstractArray)
        println("No uncertainties in the y-values were specified for given DataSet, assuming σ=1 for all y's.")
        DataSet(x, y, 1.0)
    end
    DataSet(x::AbstractArray, y::AbstractArray, allsigmas::Real) = DataSet(x, y, allsigmas*ones(length(y)*length(y[1])))
    # Also make a fancy version for DataFrames that infers the variable names?
    function DataSet(X::AbstractArray, Y::AbstractArray, Σ_y::AbstractArray)
        size(X,1) != size(Y,1) && throw("Inconsistent number of x-values and y-values given: $(size(X,1)) != $(size(Y,1)).")
        # If Σ_y not a square matrix, assume each column is vector of standard deviations associated with y:
        Σ_y = size(Σ_y,1) != size(Σ_y,2) ? Unwind(Σ_y) : Σ_y
        DataSet(Unwind(X), Unwind(Y), Σ_y, (size(X,1), ConsistentElDims(X), ConsistentElDims(Y)))
    end
    DataSet(x::AbstractVector{<:Number}, y::AbstractVector{<:Measurement}, args...) = DataSet(x,[y[i].val for i in 1:length(y)],[y[i].err for i in 1:length(y)], args...)
    ####### Only looking at sigma from here on out
    function DataSet(x::AbstractVector, y::AbstractVector, sigma::AbstractVector, dims::Tuple{Int,Int,Int})
        Sigma = Unwind(sigma)
        DataSet(Unwind(x), Unwind(y), Sigma, Diagonal([Sigma[i]^(-2) for i in 1:length(Sigma)]), dims)
    end
    DataSet(x::AbstractVector, y::AbstractVector, Σ::AbstractMatrix, dims::Tuple{Int,Int,Int}) = DataSet(Unwind(x), Unwind(y), Σ, inv(Σ), dims)
    function DataSet(x::AbstractVector{<:Number},y::AbstractVector{<:Number},sigma::AbstractArray{<:Number},InvCov::AbstractMatrix{<:Number},dims::Tuple{Int,Int,Int})
        !all(x->(x > 0), dims) && throw("Not all dims > 0: $dims.")
        !(Npoints(dims) == Int(length(x)/xdim(dims)) == Int(length(y)/ydim(dims)) == Int(size(InvCov,1)/ydim(dims))) && throw("Inconsistent input dimensions.")
        x = floatify(x);  y = floatify(y);  InvCov = HealthyCovariance(InvCov)
        if xdim(dims) == 1
            return DataSet(x, y, InvCov, dims, logdet(InvCov), nothing)
        else
            # return new(x,y,InvCov,dims,logdet(InvCov),collect(Iterators.partition(x,xdim(dims))))
            return DataSet(x, y, InvCov, dims, logdet(InvCov), [SVector{xdim(dims)}(Z) for Z in Windup(x,xdim(dims))])
        end
    end
    function DataSet(x::AbstractVector, y::AbstractVector, InvCov::AbstractMatrix, dims::Tuple{Int,Int,Int}, logdetInvCov::Real, WoundX::Union{AbstractVector,Nothing})
        DataSet(x, y, InvCov, dims, logdetInvCov, WoundX, CreateSymbolNames(xdim(dims),"x"), CreateSymbolNames(ydim(dims),"y"))
    end
    function DataSet(x::AbstractVector, y::AbstractVector, InvCov::AbstractMatrix, dims::Tuple{Int,Int,Int},
                            logdetInvCov::Real, WoundX::Union{AbstractVector,Nothing}, xnames::AbstractVector{String}, ynames::AbstractVector{String})
        new(x, y, InvCov, dims, logdetInvCov, WoundX, xnames, ynames)
    end
end

# For SciMLBase.remake
DataSet(;
x::AbstractVector=[0.],
y::AbstractVector=[0.],
InvCov::AbstractMatrix=Diagonal([1.]),
dims::Tuple{Int,Int,Int}=(1,1,1),
logdetInvCov::Real=-Inf,
WoundX::Union{AbstractVector,Nothing}=nothing,
xnames::AbstractVector{String}=["x"],
ynames::AbstractVector{String}=["y"]) = DataSet(x, y, InvCov, dims, logdetInvCov, WoundX, xnames, ynames)

# Specialized methods for DataSet
dims(DS::DataSet) = DS.dims
xdata(DS::DataSet) = DS.x
ydata(DS::DataSet) = DS.y
function ysigma(DS::DataSet)
    sig = !issparse(yInvCov(DS)) ? inv(yInvCov(DS)) : inv(convert(Matrix,yInvCov(DS)))
    return isdiag(sig) ? sqrt.(Diagonal(sig).diag) : sig
end
xsigma(DS::DataSet) = zeros(Npoints(DS)*xdim(DS))

yInvCov(DS::DataSet) = DS.InvCov

WoundX(DS::DataSet) = _WoundX(DS, DS.WoundX)

logdetInvCov(DS::DataSet) = DS.logdetInvCov

xnames(DS::DataSet) = DS.xnames
ynames(DS::DataSet) = DS.ynames

# function InformNames(DS::DataSet, xnames::AbstractVector{String}, ynames::AbstractVector{String})
#     @assert length(xnames) == xdim(DS) && length(ynames) == ydim(DS)
#     DataSet(xdata(DS), ydata(DS), yInvCov(DS), (Npoints(DS),xdim(DS),ydim(DS)), logdetInvCov(DS), WoundX(DS), xnames, ynames)
# end

ext: jl


text:
abstract MeasureOp{n} <: Measure
abstract UnaryOp{A} <: MeasureOp{1}
abstract ScalarOp{A} <: MeasureOp{2}
abstract BinaryOp{A, B} <: MeasureOp{2}

immutable Neg{A <: Measure} <: UnaryOp{A}
    a::A
end

immutable Add{A <: Measure, B <: Measure} <: BinaryOp{A, B}
    a::A
    b::B
end

immutable Min{A <: Measure, B <: Measure} <: BinaryOp{A, B}
    a::A
    b::B
end

immutable Max{A <: Measure, B <: Measure} <: BinaryOp{A, B}
    a::A
    b::B
end

immutable Div{A <: Measure} <: ScalarOp{A}
    a::A
    b::Number
end

immutable Mul{A <: Measure} <: ScalarOp{A}
    a::A
    b::Number
end

# Easy simplifications
# TODO: Add more simplifications
Neg(x::Neg) = x.a
iszero(x::Measure) = false

@compat Base.:+(a::Measure, b::Measure) = Add(a, b)
@compat Base.:-(a::Measure) = Neg(a)
@compat Base.:-(a::Neg) = a.value
@compat Base.:-(a::Measure, b::Measure) = Add(a, -b)
@compat Base.:/(a::Measure, b::Number) = Div(a, b)
@compat Base.:/{T <: Measure}(a::T, b::T) = Div(a, b)
@compat Base.:*(a::Measure, b::Number) = Mul(a, b)
@compat Base.:*(a::Number, b::Measure) = Mul(b, a)
Base.min(a::Measure, b::Measure) = Min(a, b)
Base.max(a::Measure, b::Measure) = Max(a, b)


text:
immutable TDist <: ContinuousUnivariateDistribution
    df::Float64 # non-integer degrees of freedom allowed
    function TDist(d::Real)
    	if d > 0.0
    		new(float64(d))
    	else
    		error("df must be positive")
    	end
    end
end

@_jl_dist_1p TDist t

function entropy(d::TDist)
	return ((d.df + 1.0) / 2.0) *
	       (digamma((d.df + 1.0) / 2.0) - digamma((d.df) / 2.0)) +
	       (1.0 / 2.0) * log(d.df) +
	       lbeta(d.df + 1.0, 1.0 / 2.0)
end

insupport(d::TDist, x::Number) = isreal(x) && isfinite(x)

mean(d::TDist) = d.df > 1 ? 0.0 : NaN

median(d::TDist) = 0.0

modes(d::TDist) = [0.0]

function pdf(d::TDist, x::Real)
	return 1.0 / (sqrt(d.df) * beta(0.5, 0.5 * d.df)) *
	       (1.0 + x^2 / d.df)^(-0.5 * (d.df + 1.0))
end

function var(d::TDist)
	if d.df > 2.0
		return d.df / (d.df - 2.0)
	elseif d.df > 1.0
		return Inf
	else
		return NaN
	end
end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "LibPQ"
version = v"14.1"
pg_version = string(version.major, '.', version.minor)
tzcode_version = "2021e"

# Collection of sources required to build LibPQ
sources = [
    ArchiveSource(
        "https://ftp.postgresql.org/pub/source/v$pg_version/postgresql-$pg_version.tar.gz",
        "b29030525e1314d676f41e6007a96d4489ba0d03fa93e67b477c1d5386790c8f"
    ),
    ArchiveSource(
        "https://data.iana.org/time-zones/releases/tzcode$tzcode_version.tar.gz",
        "584666393a5424d13d27ec01183da17703273664742e049d4f62f62dab631775",
    ),
]

# Bash recipe for building across all platforms
# NOTE: readline and zlib are not used by libpq
script = raw"""
cd $WORKSPACE/srcdir
make CC=$BUILD_CC VERSION_DEPS= zic
export ZIC=$WORKSPACE/srcdir/zic
cd postgresql-*/
if [[ "${target}" == i686-linux-musl ]]; then
    # Small hack: swear that we're cross-compiling.  Our `i686-linux-musl` is
    # bugged and it can run only a few programs, with the result that the
    # configure test to check whether we're cross-compiling returns that we're
    # doing a native build, but then it fails to run a bunch of programs during
    # other tests.
    sed -i 's/cross_compiling=no/cross_compiling=yes/' configure
fi
./configure --prefix=${prefix} --build=${MACHTYPE} --host=${target} --with-includes=$prefix/include --with-libraries=$prefix/lib --without-readline --without-zlib --with-openssl
make -C src/interfaces/libpq install
install_license COPYRIGHT
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms(; experimental=true)

# The products that we will ensure are always built
products = [
    LibraryProduct("libpq", :LIBPQ_HANDLE)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("OpenSSL_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "TerminalImageViewer"
version = v"1.0.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/stefanhaustein/TerminalImageViewer.git", "134ca7390986e6e61f29753f7a3f54794078ea67"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/TerminalImageViewer/src/main/cpp/
# We need to compile with GCC everywhere
export CXX=g++
make
make prefix=${prefix} install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Linux(:i686, libc=:glibc),
    Linux(:x86_64, libc=:glibc),
    Linux(:aarch64, libc=:glibc),
    Linux(:armv7l, libc=:glibc, call_abi=:eabihf),
    Linux(:powerpc64le, libc=:glibc),
    Linux(:i686, libc=:musl),
    Linux(:x86_64, libc=:musl),
    Linux(:aarch64, libc=:musl),
    Linux(:armv7l, libc=:musl, call_abi=:eabihf),
    MacOS(:x86_64),
    FreeBSD(:x86_64)
]

platforms = expand_cxxstring_abis(platforms)

# The products that we will ensure are always built
products = [
    ExecutableProduct("tiv", :tiv)
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; preferred_gcc_version = v"7.1.0")


text:
#------------------------------------------------------------------------------
#  Example.jl
#
#
#
#
#
#
#
#
#
#  Paul.Soderlind@unisg.ch, April 2002, to Julia Nov 2015
#------------------------------------------------------------------------------

using Dates, Optim

include("jlFiles/BondPricePs.jl")
include("jlFiles/BondNSxPs.jl")
include("jlFiles/printmat.jl")

ytmLoss  = 0           #0: mimimize squared price errors: 1: minimize squared ytm errors
  weightLoss = 1       #1: 1/mat as weights for (fitted price-price)^2 in loss fn (only if ytmLoss==0)
#------------------------------------------------------------------------------

                                 #Swedish bond data for 29 Dec 1993

#coupons in %/yr
c  = [ 0;        0;         0;         0;         11.5;      10.75;
       11;       13;        10.25;     6;         9 ]

#time to maturity in year
tm = [ 0.00274;  0.21096;   0.46027;   0.88219;   1.67397;   3.06849;
       5.06301;  7.46027;   9.34795;   11.11507;  15.30685 ]

#interest rates.  Bills: simple rates in %/yr; bonds: yield to maturity in %/yr
y  = [ 7.75;     6.835;    6.655;     6.41;      6.215;     6.195;
       6.41;     6.755;    7.01;      7.21;      7.325 ]

n = length(y)       #number of bonds
#----------------------------------------------------------------------------

#transform the data

c  = c/100 # -> column vector, coupons and yields as 0.05 rather than 5
y  = y/100

vvc    = c .== 0                #if bill, change from simple to effective rate
y[vvc] = (1 .+ tm[vvc].*y[vvc]).^(1.0./tm[vvc]) .- 1

P = fill(NaN,length(y))         #calculate bond prices
for i = 1:n
  local ti
  ti   = mod(tm[i],1):tm[i]
  P[i] = BondPricePs(y[i],c[i],ti)
end
#----------------------------------------------------------------------------

#estimating parameters in extended Nelson&Siegel model, restricted so b0 + b1 = log(1+y[1])

parX0 = [0.1045;-0.03;-0.0562;1.2;0;0.5]       #starting guess

if ytmLoss == 1                                #loss(ytm)
  NSXbR = BondNSxEstPs(parX0,y,tm,c,log(1+y[1]),1)
else
  if weightLoss == 1                           #loss(P/maturity)
    NSXbR = BondNSxEstPs(parX0,P,tm,c,log(1+y[1]),0,1.0./tm)
  else                                         #loss(P)
    NSXbR = BondNSxEstPs(parX0,P,tm,c,log(1+y[1]))
  end
end

println("\nParameter estimates: ")
printTable(NSXbR,[""],["b0","b1","b2","tau","b3","tau2"])
#----------------------------------------------------------------------------

ytmx = fill(NaN,n)            #model implied ytm
for i = 1:n
  local ti,d,Qx
  ti      = mod(tm[i],1):tm[i]
  d       = BondNSxPs(ti,NSXbR...)[3]    #... expands into NSXbR[1],NSXbR[2],...
  Qx      = sum(d.*c[i]) + d[end]     #model implied bond price
  #println(ti)
  ytmx[i] = BondYieldToMatPs(Qx,c[i],ti,1,1,0.05,1e-7)[1]
end

println("\nActual and model ytm, %: ")
printTable([y ytmx]*100,["ytm actual","ytm model"],string.(tm),width=15,cell00="maturity")
#----------------------------------------------------------------------------

#calculate model implied rates (spot, forward, yield to maturity) to plot

tmFig      = [1e-8;0.1:0.1:16]           #maturities to plot
(shx,fhx,) = BondNSxPs(tmFig,NSXbR...)
shx        = exp.(shx) .- 1     #effective interest rate
fhx        = exp.(fhx) .- 1

println("\nmodel spot and forward rates, %: ")
printTable([shx fhx]*100,["spot","forward"],string.(tmFig),width=15,cell00="maturity")
#----------------------------------------------------------------------------

#Comment out this if you do not have PyPlot installed.

using PyPlot
#PyPlot.svg(true)           #for ipynb notebooks
close("all")
                                       #plotting
figure()
  plot(tm,c,"+",tm,y,"s",tmFig,shx,"b--",tmFig,fhx,"r-",tm,ytmx,".")
  title("Swedish Interest Rates 29 Dec 1993")
  xlabel("Years to Maturity")
  legend(["Coupon rate","Yield to maturity","Estimated spot rate",
         "Estimated forward rate","Estimated yield to maturity"],loc=1)
  ylim(0.04,0.14)
  #display(gcf())            #uncomment in Atom/Juno
#----------------------------------------------------------------------------

max_stars_repo_name: PaulSoderlind/PaulSoderlindCode

max_issues_repo_name: PaulSoderlind/PaulSoderlindCode

max_forks_repo_name: PaulSoderlind/PaulSoderlindCode


text:
__precompile__()
module ConicNonlinearBridge

using MathProgBase
using JuMP

using Compat.LinearAlgebra
using Compat.SparseArrays

if VERSION > v"0.7.0-"
    # this is required because findall return type changed in v0.7
    function SparseArrays.findnz(A::AbstractMatrix)
        I = findall(!iszero, A)
        return (getindex.(I, 1), getindex.(I, 2), A[I])
    end
end

include("nonlinear_to_conic.jl")

end # module


text:
#=
# Reduction

The `reduce` function takes in a binary operator `⊕` and a ordered collection, applying the operator to that collection effectively reducing it to one final value.

For example, the operator ⊕ can be `minimum` and the collection can be an array of Integers.
⊕ can also be `addition` or `xor` and the collection may represent any type of object as long as the operator makes sense in the context.
=#

using CUDA

a = rand(10)

#-

reduce(min, a)

#-

reduce(*, a)

#-

@doc reduce(*, a)

#=
Writing `reduce` for a CPU is quite straightforward with a single `for-loop`. We will focus on writing a reduction for a linear array in this tutorial. We will iteratively develop a performant version using everything we have learnt in the previous tutorials.

---

# Reduction 1 : Divide and Conquer

The first step is to write something that works on the GPU. If we were given a dual-core machine and expected to parallelize this we would split the input array into two halves and feed each half into a different CPU. Similarly the best hint is to use the [divide-and-conquer](https://en.wikipedia.org/wiki/Divide-and-conquer_algorithm) approach.
By envisioning the reduction process as a binary tree we get:

![reduction-1](../assets/reduction1.png)

It would be a good exercise to try to write model the above process in pseudocode.

One such approach is:

The only issue with this approach on the GPU is that after each step we need to synchronize which won't be possible with arrays which span over a single thread block (1024 threads is the maximum threads in a block). Hence, we will have to use a recursive approach.

Assume for now we have $1024$ threads per block and process one element per thread. If we have $2048$ threads then we will run our algorithm with two thread blocks, storing the results in an intermediate array. After our first kernel is done we will perform a reduction on the intermediate array. And if we have an array whose length is greater than $1024^2$ we will have another level of recursion. If there are $1024*1024 + 1$ elements then the $1^{st}$ level of reduction will return an intermediate array of size $1025$ which will take another level of recursion to process.
=#

function reduction1(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, threadsPerBlock))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda shmem=shmem threads=threadsPerBlock blocks=blocks reduction1_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction1(op, sums)[1]
    end

    CUDA.@allowscalar return sums[1]
end

#-

function reduction1_kernel(op, a, sums)
    shmem = @cuDynamicSharedMem(eltype(a), (blockDim().x, ))
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index = threadIdx().x
    len = blockDim().x

    ## Adjust length for the last block
    if blockIdx().x == gridDim().x
        len = mod1(length(a), blockDim().x)
    end

    if tid <= length(a)
        shmem[index] = a[tid]
    else
        return
    end
    sync_threads()

    steps = floor(Int, CUDA.log2(convert(Float32, len)))
    for i = 0:steps
        if mod(index - 1, 2^(i + 1)) == 0 && (index + 2^i) <= len
            shmem[index] = op(shmem[index], shmem[index + 2^i])
        end
        sync_threads()
    end

    if index == 1
        sums[blockIdx().x] = shmem[1]
    end
    return
end

#-

a = CUDA.ones(1025);
reduction1(+, a)

#-

a = CUDA.rand(100_000);
println(reduction1(+, a),"  ",reduce(+, a))

#=
The two results above not being exactly equal is expected since IEEE floats are neither associative nor commutative. Since associativity is tough to achieve on a parallel algorithm we can expect some deviation.

---

# Reduction 2 : Strided Index

One problem with the last reduction was divergent branching, for example thread three is active for exactly one computation (`a[3] op a[4]`) and is never used again. Because GPU's operate warpwise we want to use all the resources of a warp instead of a small subset. When threads in a warp do different things it has diverged and it's efficiency drops. In this case with each iteration half the number of threads go inactive.

A simple fix is to change the way threads map to the elements by using a strided index

![reduction-2](../assets/reduction2.png)
=#

function reduction2_kernel(op, a, sums)
    shmem = @cuDynamicSharedMem(eltype(a), (blockDim().x, ))
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index = threadIdx().x
    len = blockDim().x

    ## Adjust length for the last block
    if blockIdx().x == gridDim().x
        len = mod1(length(a), blockDim().x)
    end

    if tid <= length(a)
        shmem[index] = a[tid]
    else
        return
    end
    sync_threads()

    steps = floor(Int, CUDA.log2(convert(Float32, len)))
    for i = 0:steps
        stride = 2^i
        sindex = 2*stride*(index - 1) + 1
        if sindex + stride <= len
            shmem[sindex] = op(shmem[sindex], shmem[sindex + stride])
        end
        sync_threads()
    end

    if index == 1
        sums[blockIdx().x] = shmem[1]
    end
    return
end

#-

function reduction2(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, threadsPerBlock))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda shmem=shmem threads=threadsPerBlock blocks=blocks reduction2_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction2(op, sums)
    end

    CUDA.@allowscalar return sums[1]
end

#-

a = CUDA.ones(100_000);
reduction2(+, a)

#-

@time CUDA.@sync reduction1(+, a);
@time CUDA.@sync reduction2(+, a);

#=
# Reduction 3 : Sequential access

In both the above implementations our memory-access pattern is *strided* which is difficult to coalesce. We discussed *coalesced* memory access in the **Shared Memory** tutorial.

**TL;DR** When consecutive threads access consecutive locations in memory, the GPU combines several transactions into a fewer transactions which is called coalesced memory access. When memory accesses are not consecutive which happens when the locations are non-sequantial, sparse or misaligned the GPU hardware is unable to reduce the number of transactions. Since transactions are serviced sequentially there is a significant performance penalty for non-coalesced access.

To make use of sequantial access instead of `stride` iterating from 1 to `length ÷ 2` we can do it the other way around (`length ÷ 2`:1)

**NOTE**: The algorithm below assumes that the `blockDim` is a power of two. Transforming it to become friendly with non-power of two can be done as an exercise.
=#

function reduction3_kernel(op, a, sums)
    shmem = @cuDynamicSharedMem(eltype(a), (blockDim().x, ))
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index = threadIdx().x
    len = blockDim().x

    ## Adjust length for the last block
    if blockIdx().x == gridDim().x
        len = mod1(length(a), blockDim().x)
    end

    if tid <= length(a)
        @inbounds shmem[index] = a[tid]
    else
        return
    end
    sync_threads()

    stride = len ÷ 2
    while stride > 0
        if index <= stride && index + stride <= len
            shmem[index] = op(shmem[index], shmem[index + stride])
        end
        stride = stride ÷ 2
        sync_threads()
    end

    if index == 1
        @inbounds sums[blockIdx().x] = shmem[1]
    end

    return
end

#-

function reduction3(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, threadsPerBlock))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda shmem=shmem threads=threadsPerBlock blocks=blocks reduction3_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction3(op, sums)[1]
    end

    CUDA.@allowscalar return sums[1]
    return sums
end

#-

a = CUDA.ones(1024);
reduction3(+, a)

#-

a = CUDA.ones(Int, 1024 * 1024);
reduction3(+, a)

#-

CUDA.@time reduction2(+, a);
CUDA.@time reduction3(+, a);

#=
# Reduction 4 : Warp Shuffle

A powerful feature in modern GPUs is the ability to communicate within warps with the help of special instructions. Currently we transfer data with the help of shared memory which obviously requires `sync_threads()` and access to shared memory. Warp shuffle functions allow transferring memory within a warp without the use of shared memory also being much faster and not requiring any explicit barrier. The only drawback is that only the following primitive types are supported: `Int32, UInt32, Int64, UInt64, Float32, Float64` and any arbitrary source to destination lane mapping is not permitted.

There are four shuffle methods.
- `shfl_sync`
- `shfl_up_sync`
- `shfl_down_sync`
- `shfl_xor_sync`
=#

@doc shfl_sync

#=
`shfl_sync` acts as a broadcast transferring a lane's value to all other lane.
=#

function broadcast_gpu(lane)
    id = threadIdx().x
    val = id
    mask = typemax(UInt32) # 0xffffffff
    newval = shfl_sync(mask, val, lane)
    @cuprint("id: ", id, "\t value: ", val, "\t new value: ", newval, "\n")
    return
end

@cuda threads=32 blocks=1 broadcast_gpu(19)

#=
`shfl_up_sync` and `shfl_down_sync` copy the value from lane = current_lane ± delta. If lane is out of bounds from the warp then
=#

@doc shfl_down_sync

#-

function shfldown_gpu(delta)
    id = threadIdx().x
    val = id
    mask = typemax(UInt32) # 0xffffffff
    newval = shfl_down_sync(mask, val, delta)
    @cuprint("id: ", id, "\t old value: ", val, "\t new value: ", newval, "\n")
    return
end

@cuda threads=32 blocks=1 shfldown_gpu(2)
synchronize()

#=
We can use `shfl_down_sync` to reduce a warp much faster than shared memory.
=#

@inline function reducewarp(op, val, mask = typemax(UInt32))
    val = op(val, shfl_down_sync(mask, val, 1))
    val = op(val, shfl_down_sync(mask, val, 2))
    val = op(val, shfl_down_sync(mask, val, 4))
    val = op(val, shfl_down_sync(mask, val, 8))
    val = op(val, shfl_down_sync(mask, val, 16))
    return val
end

function reduction4_kernel(op, a, sums)
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    lane_id = tid % 32
    warp_id = cld(tid, 32)

    ## exit
    tid > length(a) && return

    ## set essential
    tid <= length(a) && (val = a[tid])

    val = reducewarp(op, val)
    lane_id == 1 && (sums[warp_id] = val)
    return
end

function reduction4(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, 32))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda threads=threadsPerBlock blocks=blocks reduction4_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction4(op, sums)
    end

    CUDA.@allowscalar return sums[1]
    return sums
end

#-

a = CUDA.ones(Int, 320_000)
reduction4(+, a)

#=
There is one big problem with our implementation, the input length is expected to be a multiple of 32. This problem can be solved either by defining a neutral element for `op` (`zero(eltype(a))` for `+`, `one(eltype(a))` for `*`) however this won't work with `xor`. Another is to force the last warp's computation via shared memory like earlier examples.
We correct this by having only 1 thread work for the last warp; this is a simple solution that is inefficient but when the number of warps is large the performance hit shouldn't be too high.
=#

function reduction5_kernel(op, a, sums)
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    lane_id = tid % 32
    warp_id = cld(tid, 32)

    ## exit non essential
    tid > cld(length(a), 32)*32 && return

    ## set essential
    tid <= length(a) && (val = a[tid])

    ## Manage last warp
    if warp_id*32 > length(a)
        if lane_id == 1
            for i=(tid + 1):length(a)
                val = op(val, a[i])
            end
        end
    else
        val = reducewarp(op, val)
    end

    lane_id == 1 && (sums[warp_id] = val)
    return
end

function reduction5(op, a::CuArray)
    threadsPerBlock = 1024
    len = length(a)

    sums = similar(a, cld(len, 32))

    blocks = cld(len, threadsPerBlock)
    shmem = sizeof(eltype(a))*threadsPerBlock
    @cuda threads=threadsPerBlock blocks=blocks reduction5_kernel(op, a, sums)

    ## Recursively call reduction for larger arrays
    if length(sums) > 1
        return reduction5(op, sums)
    end

    CUDA.@allowscalar return sums[1]
    return sums
end

#-

a = CUDA.ones(Int, 5_000_000)
reduction5(+, a)

#-

@time CUDA.@sync reduction3(+, a);
@time CUDA.@sync reduction5(+, a);


text:
using Test

@timed_testset "net_components/" begin
    include("net_components/core_ops.jl")
    include("net_components/layers.jl")
    include("net_components/nets.jl")
end


text:
function randomordinate(low, high)
    if low > high
      low, high = high, low
    end
    return low + rand() * abs(high - low)
end

"""
    randompoint(lowpt, highpt)

Return a random point somewhere inside the rectangle defined by the two points.
"""
function randompoint(lowpt::Point, highpt::Point)
  Point(randomordinate(lowpt.x, highpt.x), randomordinate(lowpt.y, highpt.y))
end

"""
    randompoint(lowx, lowy, highx, highy)

Return a random point somewhere inside a rectangle defined by the four values.
"""
function randompoint(lowx, lowy, highx, highy)
    Point(randomordinate(lowx, highx), randomordinate(lowy, highy))
end

"""
    randompointarray(lowpt, highpt, n)

Return an array of `n` random points somewhere inside the rectangle defined by two points.
"""
function randompointarray(lowpt::Point, highpt::Point, n)
  array = Point[]
  for i in 1:n
    push!(array, randompoint(lowpt, highpt))
  end
  array
end

"""
    randompointarray(lowx, lowy, highx, highy, n)

Return an array of `n` random points somewhere inside the rectangle defined by the four
coordinates.
"""
function randompointarray(lowx, lowy, highx, highy, n)
    array = Point[]
    for i in 1:n
        push!(array, randompoint(lowx, lowy, highx, highy))
    end
    array
end

# internal function used by Poisson disk sampling below
"""
_empty_neighbourhood(sample, w, h, cellsize, d, points, grid)

Uses entries in `grid` to check whether the `sample`
point is more than `d` units away from any other point
in `points`.

The region we're analyzing lies between the origin and
`Point(w, h)``.
"""
function _empty_neighbourhood(sample::Point, w, h, cellsize, d, points::Array{Point, 1}, grid)
    if sample.x >= 0 &&
            sample.x < w &&
            sample.y >= 0 &&
            sample.y < h
        cellx    = 1 + convert(Int, ceil(sample.x/cellsize))
        celly    = 1 + convert(Int, ceil(sample.y/cellsize))
        _w, _h   = size(grid)
        sstartx  = max(1, cellx - 2)
        sfinishx = min(cellx + 2, _w)
        sstarty  = max(1, celly - 2)
        sfinishy = min(celly + 2, _h)
        for x in sstartx:sfinishx
            for y = sstarty:sfinishy
                ptindex = grid[x, y]
                if ptindex > 0
                    dist = distance(sample, points[ptindex])
                    if dist < d
                        return false
                    end
                end
            end
        end
        return true
    end
    return false
end

"""
    randompointarray(w, h, d; attempts=20)

Return an array of randomly positioned points inside the
rectangle defined by the current origin (0/0) and the
`width` and `height`. `d` determines the minimum
distance between each point. Increase `attempts` if you want
the function to try harder to fill empty spaces; decrease it
if it's taking too long to look for samples that work.

This uses Bridson's Poisson Disk Sampling algorithm:
https://www.cs.ubc.ca/~rbridson/docs/bridson-siggraph07-poissondisk.pdf

## Example

```
for pt in randompointarray(BoundingBox(), 20)
    randomhue()
    circle(pt, 10, :fill)
end
```
"""
function randompointarray(w, h, d;
        attempts=20)
    cellsize = d/sqrt(2)
    grid = zeros(Int,
        convert(Int, ceil(w/cellsize)),
        convert(Int, ceil(h/cellsize)))
    points = Point[]
    activepoints = Point[]
    # start with point in middle
    push!(activepoints, Point(w/2, h/2))
    while !isempty(activepoints)
        random_index = rand(1:length(activepoints))
        sample_center = activepoints[random_index]
        sample_added = false
        for i in 1:attempts
            angle = rand() * 2π
            dir = (sin(angle), cos(angle))
            sample = sample_center + dir .* (d + rand() * d)
            if _empty_neighbourhood(sample, w, h, cellsize, d, points, grid)
                push!(points, sample)
                push!(activepoints, sample)
                grid[convert(Int, ceil(sample.x/cellsize)), convert(Int, ceil(sample.y/cellsize))] = length(points)
                sample_added = true
                break
            end
        end
        if !sample_added
            deleteat!(activepoints, random_index)
        end
    end
    return points
end

"""
    randompointarray(bbox::BoundingBox, d; attempts=20)

Return an array of randomly positioned points inside the
bounding box `d` units apart.
"""
function randompointarray(bbox::BoundingBox, d; kwargs...)
    bw = boxwidth(bbox)
    bh = boxheight(bbox)
    return randompointarray(bw, bh, d; kwargs...) .- (Point(bw/2, bh/2))
end


text:
# https://adventofcode.com/2020/day/4
abstract type Day4 <: AocDay end
day4 = Day4

function parsepassport(entry)
    kv = [split(e, ":") for e in split(entry, " ")]
    pd = Dict(zip([i[1] for i in kv], [i[2] for i in kv]))
    pd
end

function valid_passport(pd)
    expected_keys = ["byr", "iyr", "eyr", "hgt", "hcl", "ecl", "pid"]
    for k in expected_keys
        if !(k in keys(pd))
            return false
        end
    end
    return true
end

function valid_byr(x)
    x = parse(Int, x)
    x >= 1920 && x <= 2002
end

function valid_iyr(x)
    x = parse(Int, x)
    x >= 2010 && x <= 2020
end

function valid_eyr(x)
    x = parse(Int, x)
    x >= 2020 && x <= 2030
end

function valid_hgt(x)
    m = match(r"^(?<number>\d+)(?<metric>(cm|in))$", x)
    n = parse(Int, m["number"])
    if m["metric"] == "cm"
        return n >= 150 && n <= 193
    end
    if m["metric"] == "in"
        return n >= 59 && n <= 76
    end
    return false
end

valid_hcl(x) = typeof(match(r"^#[a-f0-9]{6}$", x)) == RegexMatch
valid_ecl(x) = x in ["amb", "blu", "brn", "gry", "grn", "hzl", "oth"]
valid_pid(x) = typeof(match(r"^[0-9]{9}$", x)) == RegexMatch

function valid_passport_part2(pd)
    try
        valid_passport(pd) || return false
        valid_byr(pd["byr"]) || return false
        valid_iyr(pd["iyr"]) || return false
        valid_eyr(pd["eyr"]) || return false
        valid_hgt(pd["hgt"]) || return false
        valid_hcl(pd["hcl"]) || return false
        valid_ecl(pd["ecl"]) || return false
        valid_pid(pd["pid"]) || return false
    catch e
        return false
    end
    return true
end

function readinput(::Type{Day4})
    contents = read("input/day4.txt", String)
    entries = split(contents, "\r\n\r\n")
    entries = [replace(entry, "\r\n" => " ") for entry in entries]

    passports = map(parsepassport, entries)
    passports
end

function solve(day::Type{Day4})
    input = readinput(day)
    valid_passports = input[map(valid_passport, input)]
    length(valid_passports)
end

function solve_part2(day::Type{Day4})
    input = readinput(day)
    valid_passports = input[map(valid_passport_part2, input)]
    length(valid_passports)
end


text:
export compute_stress_rivlin_saunders, compute_maxwell_viscoelasticity_matrices

@doc raw"""
    compute_stress_rivlin_saunders(strain::Union{PyObject, Array{Float64,2}},c1::Union{PyObject, Float64},c2::Union{PyObject, Float64})

Computes the stress using the plane stress incompressible Rivlin Saunders model.
"""
function compute_stress_rivlin_saunders(strain::Union{PyObject, Array{Float64,2}},c1::Union{PyObject, Float64},c2::Union{PyObject, Float64})
    rivlin_saunders_ = load_op_and_grad("$(@__DIR__)/../../deps/CustomOp/build/libDataLib","rivlin_saunders")
    strain,c1,c2 = convert_to_tensor([strain,c1,c2], [Float64,Float64,Float64])
    rivlin_saunders_(strain,c1,c2)
end

@doc raw"""
    compute_maxwell_viscoelasticity_matrices(mu::Union{PyObject, Array{Float64,1}},eta::Union{PyObject, Array{Float64,1}},lambda::Union{PyObject, Array{Float64,1}},dt::Union{PyObject, Float64})

Computes the viscoelasticity matrix in the Maxwell model. `mu`, `lambda` are Lamé parameters, and `eta` is the viscosity parameter.
These three inputs are all length $N$ vectors, where $N$ is the total number of Gauss points.
`dt` is the  time step.

The time advancing formula is

$$\sigma^{n+1} = H(\epsilon^{n+1}-\epsilon^n) + S\sigma^n$$

This function returns `S`, `H`, and both are $N\times 3\times 3$ tensors.
"""
function compute_maxwell_viscoelasticity_matrices(mu::Union{PyObject, Array{Float64,1}},lambda::Union{PyObject, Array{Float64,1}},eta::Union{PyObject, Array{Float64,1}},dt::Union{PyObject, Float64})
    viscoelasticity_matrices_ = load_op_and_grad("$(@__DIR__)/../../deps/CustomOp/build/libDataLib","viscoelasticity_matrices", multiple=true)
    mu,eta,lambda,dt = convert_to_tensor(Any[mu,eta,lambda,dt], [Float64,Float64,Float64,Float64])
    viscoelasticity_matrices_(mu,eta,lambda,dt)
end


text:
module KalmanFilters

    using
        DocStringExtensions,
        Distributions,
        LinearAlgebra,
        LazyArrays,
        Statistics

    const liblapack = Base.liblapack_name

    import Statistics: mean, cov

    import ..LinearAlgebra.BLAS.@blasfunc

    import ..LinearAlgebra: BlasFloat, BlasInt,
        DimensionMismatch, chkstride1, checksquare, cholesky

    export
        WanMerweWeightingParameters,
        MeanSetWeightingParameters,
        GaussSetWeightingParameters,
        ScaledSetWeightingParameters,
        Augment,
        KFTUIntermediate,
        SRKFTUIntermediate,
        SRUKFTUIntermediate,
        SRAUKFTUIntermediate,
        AUKFTUIntermediate,
        UKFTUIntermediate,
        KFMUIntermediate,
        SRKFMUIntermediate,
        UKFMUIntermediate,
        SRUKFMUIntermediate,
        SRAUKFMUIntermediate,
        AUKFMUIntermediate,
        get_state,
        get_covariance,
        get_innovation,
        get_innovation_covariance,
        get_kalman_gain,
        get_sqrt_covariance,
        get_sqrt_innovation_covariance,
        time_update,
        time_update!,
        measurement_update,
        measurement_update!,
        calc_nis,
        nis_test,
        sigma_bound_test,
        two_sigma_bound_test,
        ConsideredState,
        ConsideredCovariance,
        ConsideredMeasurementModel,
        ConsideredProcessModel

    abstract type AbstractUpdate{X,P} end
    abstract type AbstractTimeUpdate{X,P} <: AbstractUpdate{X,P} end
    abstract type AbstractMeasurementUpdate{X,P} <: AbstractUpdate{X,P} end
    abstract type AbstractWeightingParameters end

    function get_state(kf::AbstractUpdate)
        kf.state
    end
    function get_covariance(kf::AbstractUpdate)
        kf.covariance
    end
    get_innovation(kf::AbstractMeasurementUpdate) = kf.innovation
    get_innovation_covariance(kf::AbstractMeasurementUpdate) = kf.innovation_covariance
    get_kalman_gain(kf::AbstractMeasurementUpdate) = kf.kalman_gain

    function get_covariance(kf::AbstractUpdate{X,P}) where {X,P <: Cholesky}
        kf.covariance.L * kf.covariance.U
    end

    function get_sqrt_covariance(kf::AbstractUpdate{X,P}) where {X,P <: Cholesky}
        kf.covariance
    end

    function get_innovation_covariance(
        kf::AbstractMeasurementUpdate{X,P}
    ) where {X,P <: Cholesky}
        kf.innovation_covariance.L * kf.innovation_covariance.U
    end

    function get_sqrt_innovation_covariance(
        kf::AbstractMeasurementUpdate{X,P}
    ) where {X,P <: Cholesky}
        kf.innovation_covariance
    end

    @static if VERSION < v"1.1"
        eachcol(A::AbstractVecOrMat) = (view(A, :, i) for i in axes(A, 2))
    end

    include("gels.jl")
    include("kf.jl")
    include("srkf.jl")
    include("sigmapoints.jl")
    include("augmented_sigmapoints.jl")
    include("ukf.jl")
    include("srukf.jl")
    include("aukf.jl")
    include("sraukf.jl")
    include("consider.jl")
    include("tests.jl")

end


text:
using DynamicGrids, Dates, Test, Colors, ColorSchemes, FileIO
using FreeTypeAbstraction
using DynamicGrids: grid_to_image!, imagegen, minval, maxval, normalise, SimData, NoDisplayImageOutput,
    isstored, isasync, initialise!, finalise!, delay, fps, settimestamp!, timestamp, textconfig,
    tspan, setfps!, frames, isshowable, showframe, to_rgb, scale, Extent, extent
using ColorSchemes: leonardo

@testset "to_rgb" begin
    @test to_rgb(0.5) === ARGB32(0.5, 0.5, 0.5, 1.0)
    @test to_rgb((0.5, 0.5, 0.5)) === ARGB32(0.5, 0.5, 0.5, 1.0)
    @test to_rgb((0.5, 0.5, 0.5, 1.0)) === ARGB32(0.5, 0.5, 0.5, 1.0)
    @test to_rgb(RGB(0.5, 0.5, 0.5)) === ARGB32(0.5, 0.5, 0.5, 1.0)
    @test to_rgb(ARGB32(0.5, 0.5, 0.5)) === ARGB32(0.5, 0.5, 0.5, 1.0)
    @test to_rgb(Greyscale(), 0.5) === ARGB32(0.5, 0.5, 0.5, 1.0)
    @test to_rgb(ObjectScheme(), 0.5) === ARGB32(0.5, 0.5, 0.5, 1.0)
end

@testset "normalise" begin
    @test normalise(-.2, 0.0, 1.0) == 0.0
    @test normalise(1.2, 0.0, 1.0) == 1.0
    @test normalise(-.2, 0.0, nothing) == 0.0
    @test normalise(1.2, nothing, 1.0) == 1.0
    @test normalise(1.2, nothing, nothing) == 1.2
end

@testset "scale" begin
    @test scale(0.0, 5.0, 10.0) == 5.0
    @test scale(0.5, 5.0, 10.0) == 7.5
    @test scale(1.0, nothing, 10.0) == 10.0
    @test scale(0.0, -2.0, nothing) == -2.0
    @test scale(1.2, nothing, nothing) == 1.2
end

l0 = ARGB32(get(leonardo, 0))
l05 = ARGB32(get(leonardo, 0.5))
l08 = ARGB32(get(leonardo, 0.8))
l1 = ARGB32(get(leonardo, 1))

images = []

DynamicGrids.showimage(image, o::NoDisplayImageOutput) = begin
    push!(images, image)
    image
end

@testset "basic ImageOutput" begin
    init_ = [8.0 10.0;
             NaN  5.0]
    output = NoDisplayImageOutput(init_; tspan=1:1, maxval=40.0, text=nothing)

    @test all(parent(output)[1] .=== init_)
    @test minval(output) === nothing
    @test maxval(output) === 40.0
    @test textconfig(output) === nothing
    @test imagegen(output).scheme == ObjectScheme()
    @test isasync(output) == false
    @test isstored(output) == false
    @test delay(output, 1.0) === nothing
    @test timestamp(output) === 0.0
    pre = time()
    settimestamp!(output, 1)
    @test timestamp(output) > pre
    @test length(output) == 1
    push!(output, 2init_)
    @test length(output) == 2
    @test all(output[2] .=== 2init_)
    @test tspan(output) == 1:1
    @test fps(output) === 25.0
    @test setfps!(output, 1000.0) === 1000.0
    @test fps(output) === 1000.0
    output[1] = 5init_
    @test all(frames(output)[1] .=== 5init_)
    @test isshowable(output, 1)

    imgen = Image()
    output = NoDisplayImageOutput(init_;
        tspan=1:10, maxval=40.0, imagegen=imgen, text=nothing
    )
    simdata = SimData(output, Ruleset(Life()))
    z0 = DynamicGrids.ZEROCOL
    ref = [ARGB32(0.2, 0.2, 0.2) ARGB32(0.25, 0.25, 0.25)
           z0                    ARGB32(0.125, 0.125, 0.125)]
    @test showframe(output, simdata) == ref
    savegif("test.gif", output)
    gif = load("test.gif")
    @test gif == RGB.(ref)
    rm("test.gif")

    arrayoutput = ArrayOutput(init_; tspan=1:1)
    @test minval(arrayoutput) == nothing
    @test maxval(arrayoutput) == nothing
    @test imagegen(arrayoutput) isa Image
    @test fps(arrayoutput) === nothing
    savegif("test2.gif", arrayoutput; maxval=40.0, imagegen=imgen, text=nothing)
    gif = load("test2.gif")
    @test gif == RGB.(ref)
    rm("test2.gif")
end

@testset "Image generator" begin
    init_ = [8.0 10.0;
             0.0  5.0]
    mask_ = Bool[0 1;
                 1 1]
    imgen = Image(zerocolor=(1.0, 0.0, 0.0), maskcolor=(0.1, 0.1, 0.1))
    ic = DynamicGrids.ImageConfig(init_; imagegen=imgen, textconfig=nothing)
    @test ic.imagegen === imgen
    output = NoDisplayImageOutput((a=init_,);
        tspan=DateTime(2001):Year(1):DateTime(2010), mask=mask_,
        imagegen=imgen, text=nothing, minval=0.0, maxval=10.0, store=true
    )
    @test imagegen(output) === output.imageconfig.imagegen === imgen
    @test minval(output) === 0.0
    @test maxval(output) === 10.0
    @test imagegen(output).zerocolor == Image(zerocolor=(1.0, 0.0, 0.0)).zerocolor
    @test isstored(output) == true
    simdata = SimData(output, Ruleset(Life()))

    # Test level normalisation
    normed = normalise.(output[1][:a], minval(output), maxval(output))
    @test normed == [0.8 1.0
                     0.0 0.5]

    # Test greyscale Image conversion
    img = grid_to_image!(output, simdata)
    @test img == [ARGB32(0.1, 0.1, 0.1, 1.0) ARGB32(1.0, 1.0, 1.0, 1.0)
                  ARGB32(1.0, 0.0, 0.0, 1.0) ARGB32(0.5, 0.5, 0.5, 1.0)]

    output = NoDisplayImageOutput((a=init_,);
        tspan=DateTime(2001):Year(1):DateTime(2010), mask=mask_,
        imagegen=Image(; scheme=leonardo), text=nothing,
        minval=0.0, maxval=10.0, store=true
    )
    img = grid_to_image!(output, simdata)
    @test img == [DynamicGrids.MASKCOL l1
                  DynamicGrids.ZEROCOL l05]
    z0 = ARGB32(1, 0, 0)
    output = NoDisplayImageOutput((a=init_,);
        tspan=DateTime(2001):Year(1):DateTime(2010), mask=mask_,
        imagegen = Image(scheme=leonardo, zerocolor=z0), text=nothing,
        minval=0.0, maxval=10.0, store=true
    )
    img = grid_to_image!(output, simdata)
    @test img == [DynamicGrids.MASKCOL l1
                  z0 l05]

    @testset "text captions" begin
        pixelsize = 20
        timepos = 2pixelsize, pixelsize
        textinit = zeros(200, 200)
        font = "arial"
        face = findfont(font)
        # Swap fonts on linux
        if face === nothing
            font = "sans-serif"
            face = findfont(font)
        end
        if face !== nothing
            refimg = ARGB32.(map(x -> ARGB32(1.0, 0.0, 0.0, 1.0), textinit))
            renderstring!(refimg, string(DateTime(2001)), face, pixelsize, timepos...;
                          fcolor=ARGB32(1.0, 1.0, 1.0, 1.0), bcolor=ARGB32(0.0, 0.0, 0.0, 1.0))
            textconfig=TextConfig(; font=font, timepixels=pixelsize, namepixels=pixelsize, bcolor=ARGB32(0))
            output = NoDisplayImageOutput(textinit;
                tspan=DateTime(2001):Year(1):DateTime(2001),
                imagegen=Image(zerocolor=ARGB32(1.0, 0.0, 0.0, 1.0)),
                text=textconfig, store=true,
            )
            simdata = SimData(output, Ruleset())
            img = grid_to_image!(output, simdata);
            @test img == refimg
        end

    end

end

@testset "SparseOptInspector" begin
    init =  Bool[
             0 0 0 0 0 0 0
             0 0 0 0 1 1 1
             0 0 0 0 0 0 1
             0 0 0 0 0 1 0
             0 0 0 0 0 0 0
             0 0 0 0 0 0 0
            ]
    ruleset = Ruleset(Life();
        timestep=Day(1),
        boundary=Wrap(),
        opt=SparseOpt(),
    )
    imgen = SparseOptInspector()
    output = NoDisplayImageOutput(init;
        tspan=Date(2001, 1, 1):Day(1):Date(2001, 1, 5),
        imagegen=imgen, minval=0.0, maxval=1.0, store=true
    )

    @test minval(output) === 0.0
    @test maxval(output) === 1.0
    @test imagegen(output) == SparseOptInspector()
    @test isstored(output) == true

    global images = []
    sim!(output, ruleset)
    w, y, c = ARGB32(1), ARGB32(.0, .0, .5), ARGB32(.5, .5, .5)
    @test_broken images[1] == [
             y y y y y y y
             y y y c w w w
             y y y c c c w
             y y y y y w c
             y y y y y c c
             y y y y y y y
            ]

end

@testset "Layout ImageGenerator" begin
    init = [8.0 10.0;
            0.0  5.0]
    z0 = DynamicGrids.ZEROCOL
    grey = Greyscale()
    multiinit = (a=init, b=2init)
    imgen = Layout([:a, nothing, :b], (grey, leonardo))
    output = NoDisplayImageOutput(multiinit;
        tspan=DateTime(2001):Year(1):DateTime(2002),
        imagegen=imgen, text=nothing,
        minval=(0, 0), maxval=(10, 20), store=true
    )
    @test minval(output) === (0, 0)
    @test maxval(output) === (10, 20)
    @test imagegen(output) === imgen
    @test isstored(output) == true
    simdata = SimData(output, Ruleset(Life()))

    # Test image is joined from :a, nothing, :b
    @test grid_to_image!(output, simdata) ==
        [ARGB32(0.8, 0.8, 0.8, 1.0) ARGB32(1.0, 1.0, 1.0, 1.0)
         z0                         ARGB32(0.5, 0.5, 0.5, 1.0)
         ARGB32(0.0, 0.0, 0.0, 1.0) ARGB32(0.0, 0.0, 0.0, 1.0)
         ARGB32(0.0, 0.0, 0.0, 1.0) ARGB32(0.0, 0.0, 0.0, 1.0)
         l08                   l1
         z0                    l05                  ]

    @testset "text captions" begin
        timepixels = 20
        timepos = 2timepixels, timepixels
        textinit = (a=zeros(200, 200), b=zeros(200, 200))
        font = "arial"
        face = findfont(font)
        if face === nothing
            font = "cantarell"
            face = findfont(font)
        end
        if face !== nothing
            refimg = cat(fill(z0, 200, 200),
                         fill(ARGB32(0), 200, 200),
                         fill(z0, 200, 200); dims=1)
            renderstring!(refimg, string(DateTime(2001)), face, timepixels, timepos...;
                          fcolor=ARGB32(RGB(1.0), 1.0), bcolor=ARGB32(RGB(0.0), 1.0))
            namepixels = 15
            nameposa = 3timepixels + namepixels, timepixels
            renderstring!(refimg, "a", face, namepixels, nameposa...;
                          fcolor=ARGB32(RGB(1.0), 1.0), bcolor=ARGB32(RGB(0.0), 1.0))
            nameposb = 3timepixels + namepixels + 400, timepixels
            renderstring!(refimg, "b", face, namepixels, nameposb...;
                          fcolor=ARGB32(RGB(1.0), 1.0), bcolor=ARGB32(RGB(0.0), 1.0))
            textconfig = TextConfig(; font=font, timepixels=timepixels, namepixels=namepixels, bcolor=ARGB32(0))
            imgen = Layout([:a, nothing, :b], (grey, leonardo))
            output = NoDisplayImageOutput(textinit;
                 tspan=DateTime(2001):Year(1):DateTime(2001), imagegen=imgen, text=textconfig,
                 store=true, minval=(0, 0), maxval=(1, 1)
            )
            simdata = SimData(output, Ruleset())
            img = grid_to_image!(output, simdata);
            @test img == refimg
        end
    end
    @testset "errors" begin
        output = NoDisplayImageOutput(multiinit; tspan=1:10, imagegen=imgen,
            minval=(0, 0, 0),
            maxval=(10, 20),
        )
        simdata = SimData(output, Ruleset(Life()))
        @test_throws ArgumentError grid_to_image!(output, simdata)
        broken_imgen = Layout([:d, :c], (grey, leonardo))
        output = NoDisplayImageOutput(multiinit;
            tspan=1:10, imagegen=broken_imgen, text=nothing, minval=(0, 0), maxval=(10, 20),
        )
        simdata = SimData(output, Ruleset(Life()))
        @test_throws ArgumentError grid_to_image!(output, simdata)
        @test_throws ArgumentError TextConfig(; font="not_a_font")
        @test_throws ArgumentError TextConfig(; font=:not_a_string)
    end
    @testset "Layout is the default for NamedTuple of grids" begin
        output = NoDisplayImageOutput(multiinit; tspan=1:10)
        @test imagegen(output) isa Layout
        @test DynamicGrids.imagesize(imagegen(output), multiinit) == (2, 4)
    end
end


text:
using LambertW

"""
    lambertw_gpu(N)

Returns `real(LambertW.lambertw(Float64(N - 1) / MathConstants.e))`,
valid for `N = 2` and `N = 3`.

TODO: add `LambertW.lambertw` support to KernelAbstractions.
"""
function lambertw_gpu(N)
    if !(N == 2 || N == 3)
        error("Bad N in lambertw_gpu")
    end
    return (0.2784645427610738, 0.46305551336554884)[N - 1]
end

"""
    lamb_smooth_minimum(
        l::AbstractArray{FT};
        frac_upper_bound::FT,
        reg_min::FT,
    ) where {FT}

Returns the smooth minimum of the elements of an array
following the formulation of
Lopez-Gomez et al. (JAMES, 2020), Appendix A, given:
 - `l`, an array of candidate elements
 - `frac_upper_bound`, defines the upper bound of the
        smooth minimum as `smin(x) = min(x)*(1+frac_upper_bound)`
 - `reg_min`, defines the minimum value of the regularizer Λ
"""
function lamb_smooth_minimum(
    l::AbstractArray{FT},
    frac_upper_bound::FT,
    reg_min::FT,
) where {FT}

    xmin = minimum(l)

    # Get regularizer for exponential weights
    N_l = length(l)
    denom = FT(lambertw_gpu(N_l))
    Λ = max(FT(xmin) * frac_upper_bound / denom, reg_min)

    num = sum(i -> l[i] * exp(-(l[i] - xmin) / Λ), 1:N_l)
    den = sum(i -> exp(-(l[i] - xmin) / Λ), 1:N_l)
    smin = num / den

    return smin
end


text:
function myname()
  MY_NAME
end

max_stars_repo_name: PacktPublishing/Julia-Projects

max_issues_repo_name: PacktPublishing/Julia-Projects

max_forks_repo_name: PacktPublishing/Julia-Projects


text:
using Documenter, Tutorials

makedocs(;
    modules=[Tutorials],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/bates/Tutorials.jl/blob/{commit}{path}#L{line}",
    sitename="Tutorials.jl",
    authors="Douglas Bates",
    assets=String[],
)

deploydocs(;
    repo="github.com/bates/Tutorials.jl",
)


text:
# ------------------------------------------------------------------------------------------
# # `NDSparse` Usage
#
# The purpose of this section is to show you how to use the methods associated with NDSparse
# objects. If you've already worked through the section on Table Usage, this notebook will
# seem very familiar!
#
# Let's get started by loading JuliaDB and editing our default settings so that tables will
# display their entries.
# ------------------------------------------------------------------------------------------

using JuliaDB

IndexedTables.set_show_compact!(false);

# ------------------------------------------------------------------------------------------
# - `loadndsparse` is an analog to `loadtable`:
# ------------------------------------------------------------------------------------------

nd = loadndsparse("stocksample", filenamecol = :Ticker, indexcols = [:Ticker, :Date])

# ------------------------------------------------------------------------------------------
# - Having made `:Ticker` and `:Date` the columns by which we index entries to `nd`, we can
# perform a lookup in `nd` via
# ------------------------------------------------------------------------------------------

nd["aapl.us.txt", Date(1986, 2, 10)]

# ------------------------------------------------------------------------------------------
# # `selectkeys` and `selectvalues`
#
# - Rather than `select`, for `NDSparse` you can use `selectkeys` and `selectvalues`.
#
# - In this example, we generate a NDSparse object that contains two indexable columns
# (`:Ticker` and `:Date`) and one column of selected values, `:Close`:
# ------------------------------------------------------------------------------------------

selectvalues(nd, :Close)

# ------------------------------------------------------------------------------------------
# # `map`
#
# - The syntax for `map` is the same for `Table` and `NDSparse` objects:
#
# `map(function, table; select)`
# ------------------------------------------------------------------------------------------

map(r -> r.High - r.Low, nd; select = (:High, :Low))

# ------------------------------------------------------------------------------------------
# # `filter`
#
# - The syntax for filter is the same for `Table` and `NDSparse` objects:
#
# `filter(function, tablename; select)`
#
# - In the following example, we filter `nd` to create a new `NDSparse` object that only
# contains days/stocks that increased in value:
# ------------------------------------------------------------------------------------------

filter(r -> r.Open < r.Close, nd)

# ------------------------------------------------------------------------------------------
# # `reduce` and `groupreduce`
#
# - Syntax is the same as for `Table`:
#
# `reduce(reducer, table; select)`
#
# `groupreduce(reducer, table, by; select)`
# ------------------------------------------------------------------------------------------

groupreduce(+, nd, :Ticker; select = :Volume)

# ------------------------------------------------------------------------------------------
# # `groupby`
#
# - Syntax is the same as for `Table`:
#
# `groupby(function, table, by; select)`
# ------------------------------------------------------------------------------------------

groupby(extrema, nd, :Ticker; select = :Close)

# ------------------------------------------------------------------------------------------
# # `summarize`
#
# - Syntax is the same as for `Table`:
#
# `summarize(function[s], table[, by]; select)`
# ------------------------------------------------------------------------------------------

summarize((mean, std), nd, :Ticker; select = (:Open, :Close))

# ------------------------------------------------------------------------------------------
# # `columns` and `rows`
#
# - Syntax is the same as for `Table`
# ------------------------------------------------------------------------------------------

columns(nd)

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
@views maximize_noise_distributions(edge_lengths::Vector{Float64}, assignment_probs::Matrix{Float64}; updating_ids::T2 where T2 <: AbstractArray{Int} = 1:length(edge_lengths)) =
    (Normal(wmean_std(edge_lengths, assignment_probs[:,i], non_zero_ids=updating_ids)...) for i in 1:2)

function expect_noise_probabilities!(assignment_probs::Matrix{Float64}, d1::Normal, d2::Normal, edge_lengths::Vector{Float64},
        adjacent_points::Vector{Vector{Int}}, adjacent_weights::Vector{Vector{Float64}}; updating_ids::T where T <: AbstractVector{Int} = 1:length(edge_lengths),
        min_confidence::Union{Vector{Float64}, Nothing}=nothing)
    norm_denses = (pdf.(d1, edge_lengths), pdf.(d2, edge_lengths));
    n1 = sum(assignment_probs[:, 1])
    n2 = size(assignment_probs, 1) - n1

    dists = (d1, d2)
    for i in updating_ids
        cur_weights = adjacent_weights[i]
        cur_points = adjacent_points[i]

        c_d1 = c_d2 = 0.0
        for j in 1:length(cur_points)
            ap = assignment_probs[cur_points[j], 1]
            c_d1 += cur_weights[j] * ap
            c_d2 += cur_weights[j] * (1. - ap)
        end

        d1 = n1 * exp(c_d1) * norm_denses[1][i]
        d2 = n2 * exp(c_d2) * norm_denses[2][i]
        p1 = d1 / fmax(d1 + d2, 1e-20)
        if min_confidence !== nothing
            p1 = min_confidence[i] + p1 * (1. - min_confidence[i])
        end

        assignment_probs[i, 1] = p1
        assignment_probs[i, 2] = 1. - p1
    end
end

function fit_noise_probabilities(edge_lengths::Vector{Float64}, adjacent_points::Vector{Vector{Int}}, adjacent_weights::Vector{Vector{Float64}};
        min_confidence::Union{Vector{Float64}, Nothing}=nothing, max_iters::Int=10000, tol::Float64=0.005, verbose::Bool=false, progress::Union{Progress, Nothing}=nothing)
    # Initialization
    init_means = quantile(edge_lengths, (0.1, 0.9))
    init_std = (init_means[2] - init_means[1]) / 4.0
    d1, d2 = Normal(init_means[1], init_std), Normal(init_means[2], init_std)

    assignment_probs = hcat(pdf.(d1, edge_lengths), pdf.(d2, edge_lengths));

    ## Hard-assign outliers to the noise class to avoid numerical problems with super-small probabilities
    outlier_mask = (edge_lengths .> (init_means[2] + 3 * init_std))
    updating_ids = findall(.!outlier_mask)
    assignment_probs[outlier_mask, 2] .= 1.0
    assignment_probs ./= sum(assignment_probs, dims=2)
    max_diffs = Float64[]

    if verbose && progress === nothing
        progress = ProgressUnknown(0.3)
    end

    # EM iterations
    n_iters = max_iters
    for i in 1:max_iters
        expect_noise_probabilities!(assignment_probs, d1, d2, edge_lengths, adjacent_points, adjacent_weights, updating_ids=updating_ids, min_confidence=min_confidence)
        d1n, d2n = maximize_noise_distributions(edge_lengths, assignment_probs, updating_ids=updating_ids)

        ## Estimate parameter differences as convergence criteria
        param_diff = max(abs(d1n.μ - d1.μ) / d1.μ, abs(d2n.μ - d2.μ) / d2.μ, abs(d1n.σ - d1.σ) / d1.σ, abs(d2n.σ - d2.σ) / d2.σ)
        push!(max_diffs, param_diff)
        d1, d2 = d1n, d2n

        if verbose
            next!(progress)
        end

        if max_diffs[end] < tol
            if verbose
                finish!(progress)
            end
            n_iters = i
            break
        end
    end

    if verbose
        println("Algorithm stopped after $n_iters iterations. Error: $(round(max_diffs[end], sigdigits=3)). Converged: $(max_diffs[end] <= tol).")
    end

    if d1.μ > d2.μ
        d2, d1 = d1, d2
    end

    n1 = sum(assignment_probs[:, 1])
    n2 = size(assignment_probs, 1) - n1
    assignment_probs = hcat(n1 .* pdf.(d1, edge_lengths), n2 .* pdf.(d2, edge_lengths))
    assignment_probs[outlier_mask, 2] .= 1.0
    assignment_probs ./= sum(assignment_probs, dims=2)

    if min_confidence !== nothing
        assignment_probs[:, 1] .= min_confidence .+ assignment_probs[:, 1] .* (1. .- min_confidence)
        assignment_probs[:, 2] .= 1.0 .- assignment_probs[:, 1]
    end

    assignment = Int.(assignment_probs[:, 2] .> 0.5) .+ 1

    return assignment_probs, assignment, (d1, d2), max_diffs
end

function append_confidence!(df_spatial::DataFrame, prior_assignment::Union{Vector{Int}, Nothing}=nothing; nn_id::Int, prior_confidence::Float64=0.5)
    pos_data = position_data(df_spatial);
    mean_dists = getindex.(knn(KDTree(pos_data), pos_data, nn_id + 1, true)[2], nn_id + 1)

    min_confidence = (prior_assignment !== nothing) ? (prior_confidence^2 .* (prior_assignment .> 0)) : nothing

    adjacent_points, adjacent_weights = build_molecule_graph(df_spatial, filter=false)[1:2]; # TODO: can be optimized as we already have kNNs
    probs, (d1, d2) = fit_noise_probabilities(mean_dists, adjacent_points, adjacent_weights, min_confidence=min_confidence)[[1, 3]]
    df_spatial[!,:confidence] = probs[:, 1];

    return mean_dists, df_spatial[!,:confidence], d1, d2
end

max_stars_repo_name: YinuoJin/Baysor

max_issues_repo_name: YinuoJin/Baysor

max_forks_repo_name: YinuoJin/Baysor


text:
using Images
#using ImageMagick
#using FileIO
#FileIO.read(file::readformats(ImageMagick.BACKEND)) = read(file, ImageMagick.BACKEND)

function CreateTexture(file::AbstractString)
	img = Images.load(file)
	sz = size(img)
	#imgFile = open(file,"r")
	#image.Decode(imgFile)
	#img2 = convert(GrayAlpha{UFixed8}, img)
	#set_packing_alignment(data)
#=
	rgba := image.NewRGBA(img.Bounds())
	if rgba.Stride != rgba.Rect.Size().X*4
		return 0, fmt.Errorf("unsupported stride")
	end
	draw.Draw(rgba, rgba.Bounds(), img, image.Point{0, 0}, draw.Src)
=#
	texture = GLuint[0];
	glGenTextures(1, texture); id = texture[1];
	glActiveTexture(GL_TEXTURE0)
	glBindTexture(GL_TEXTURE_2D, id)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
	# glTexImage2D(
		# GL_TEXTURE_2D,
		# 0,
		# GL_RGBA,
		# sz[1], # width
		# sz[2], # height
		# 0,
		# GL_RGBA,
		# GL_UNSIGNED_BYTE,
		# pointer(img)
		# )

	return id
end

max_stars_repo_name: Gilga/JGE

max_issues_repo_name: Gilga/JGE

max_forks_repo_name: Gilga/JGE


text:
struct ZSH <: AbstractShell
    path::AbstractString
    prompt::AbstractString
end

ZSH(path::AbstractString) = ZSH(path, "%F{5}%n@%m:%~ (playground)> %f")
ZSH() = ZSH(strip(readstring(`which zsh`)))

function Base.run(shell::ZSH, env::Environment)
    init(env)
    prompt = getprompt(shell, env)
    ENV["PS1"] = prompt

    debug(logger, "Shell prompt: $prompt")

    usr_rc = join(Path(get(ENV, "ZDOTDIR", home())), ".zshrc")
    pg_rc = join(parent(Path(ENV["JULIA_PKGDIR"])), ".zshrc")

    if !exists(pg_rc)
        debug(logger, "Creating shell rc file $pg_rc...")
        exists(usr_rc) ? cp(usr_rc, pg_rc, follow_symlinks=true) : touch(pg_rc)

        content = string(
            "\nunset PROMPT\n",
            "export PATH=\"", ENV["PATH"], "\"\n",
            "export PS1=\"$prompt\"\n",
            "export JULIA_PKGDIR=\"", ENV["JULIA_PKGDIR"], "\"\n",
        )

        if haskey(ENV, "HISTFILE") && exists(Path(ENV["HISTFILE"]))
            content = string(content, "export HISTFILE=\"", ENV["HISTFILE"], "\"\n")
        end

        debug(logger, "Writing rc contents to file.")
        write(pg_rc, content, "a")
    end

    ENV["ZDOTDIR"] = parent(pg_rc)
    @mock run(`$(shell.path) -i`)
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/Materials.jl/blob/master/LICENSE

"""
    Continuum3D

A simplified 3D continuum model for material tests.
"""
mutable struct Continuum3D <: FieldProblem
    material_model :: Union{Symbol, Expr}
end

Continuum3D() = Continuum3D(:nothing)
FEMBase.get_unknown_field_name(::Continuum3D) = "displacement"

abstract type AbstractLoading end

mutable struct AxialStrainLoading <: AbstractLoading
    times :: Vector{Float64}
    loads :: Vector{Float64}
end
mutable struct ShearStrainLoading <: AbstractLoading
    times :: Vector{Float64}
    loads :: Vector{Float64}
end

function update_bc_elements!(bc_elements::Vector{Element{Poi1}}, loading::AxialStrainLoading)
    bc_element_1, bc_element_2, bc_element_3, bc_element_4, bc_element_5, bc_element_6, bc_element_7, bc_element_8 = bc_elements
    # Fix bottom side
    for element in (bc_element_1, bc_element_2, bc_element_3, bc_element_4)
        update!(element, "fixed displacement 3", 0.0)
    end
    # Set loading bc
    for element in (bc_element_5, bc_element_6, bc_element_7, bc_element_8)
        for (t,f) in zip(loading.times, loading.loads)
            # @info "Setting displacement 3 at $(t) to $(f)!"
            update!(element, "fixed displacement 3", t => f)
        end
    end
    # Set rest of boundary conditions
    update!(bc_element_1, "fixed displacement 1", 0.0)
    update!(bc_element_1, "fixed displacement 2", 0.0)
    update!(bc_element_2, "fixed displacement 2", 0.0)
    update!(bc_element_4, "fixed displacement 1", 0.0)
    update!(bc_element_5, "fixed displacement 1", 0.0)
    update!(bc_element_5, "fixed displacement 2", 0.0)
    update!(bc_element_6, "fixed displacement 2", 0.0)
    update!(bc_element_8, "fixed displacement 1", 0.0)
    return nothing
end

function update_bc_elements!(bc_elements::Vector{Element{Poi1}}, loading::ShearStrainLoading)
    bc_element_1, bc_element_2, bc_element_3, bc_element_4, bc_element_5, bc_element_6, bc_element_7, bc_element_8 = bc_elements
    # Fix bottom side
    for element in (bc_element_1, bc_element_2, bc_element_3, bc_element_4)
        update!(element, "fixed displacement 1", 0.0)
        update!(element, "fixed displacement 2", 0.0)
        update!(element, "fixed displacement 3", 0.0)
    end
    # Set top side bcs
    for element in (bc_element_5, bc_element_6, bc_element_7, bc_element_8)
        for (t,f) in zip(loading.times, loading.loads)
            update!(element, "fixed displacement 1", t => f)
        end
        update!(element, "fixed displacement 2", 0.0)
        update!(element, "fixed displacement 3", 0.0)
    end
    return nothing
end

function FEMBase.assemble_elements!(problem::Problem{Continuum3D},
                                    assembly::Assembly,
                                    elements::Vector{Element{Hex8}},
                                    time::Float64)

    bi = BasisInfo(Hex8)

    dim = 3
    nnodes = 8
    ndofs = dim*nnodes

    BL = zeros(6, ndofs)
    Km = zeros(ndofs, ndofs)
    f_int = zeros(ndofs)
    f_ext = zeros(ndofs)

    for element in elements

        u = element("displacement", time)
        fill!(Km, 0.0)
        fill!(f_int, 0.0)
        fill!(f_ext, 0.0)

        for ip in get_integration_points(element)
            J, detJ, N, dN = element_info!(bi, element, ip, time)
            material = ip("material", time)
            w = ip.weight*detJ

            # Kinematic matrix, linear part

            fill!(BL, 0.0)
            for i=1:nnodes
                BL[1, 3*(i-1)+1] = dN[1,i]
                BL[2, 3*(i-1)+2] = dN[2,i]
                BL[3, 3*(i-1)+3] = dN[3,i]
                BL[6, 3*(i-1)+1] = dN[2,i]
                BL[6, 3*(i-1)+2] = dN[1,i]
                BL[4, 3*(i-1)+2] = dN[3,i]
                BL[4, 3*(i-1)+3] = dN[2,i]
                BL[5, 3*(i-1)+1] = dN[3,i]
                BL[5, 3*(i-1)+3] = dN[1,i]
            end

            # Calculate stress response
            integrate_material!(material)
            #D = material.jacobian
            #S = material.stress + material.dstress
            D = tovoigt(material.variables_new.jacobian)
            S = tovoigt(material.variables_new.stress)

            # Material stiffness matrix
            Km += w*BL'*D*BL

            # Internal force vector
            f_int += w*BL'*S

            # External force vector
            for i=1:dim
                haskey(element, "displacement load $i") || continue
                b = element("displacement load $i", ip, time)
                f_ext[i:dim:end] += w*b*vec(N)
            end

        end

        # add contributions to K, Kg, f
        gdofs = get_gdofs(problem, element)
        add!(assembly.K, gdofs, gdofs, Km)
        add!(assembly.f, gdofs, f_ext - f_int)

    end

    return nothing
end

function FEMBase.assemble_elements!(problem::Problem{Continuum3D},
                                    assembly::Assembly,
                                    elements::Vector{Element{Quad4}},
                                    time::Float64)

    nnodes = 4
    ndofs = 3
    f = zeros(nnodes*ndofs)
    bi = BasisInfo(Quad4)

    for element in elements

        fill!(f, 0.0)

        for ip in get_integration_points(element)

            J, detJ, N, dN = element_info!(bi, element, ip, time)
            w = ip.weight*detJ

            if haskey(element, "surface pressure")
                J = element(ip, time, Val{:Jacobian})'
                n = cross(J[:,1], J[:,2])
                n /= norm(n)
                # sign convention, positive pressure is towards surface
                p = element("surface pressure", ip, time)
                f += w*p*vec(n*N)
            end
        end

        gdofs = get_gdofs(problem, element)
        add!(assembly.f, gdofs, f)

    end

    return nothing

end

function FEMBase.assemble_elements!(problem::Problem{Continuum3D},
                                    assembly::Assembly,
                                    elements::Vector{Element{Poi1}},
                                    time::Float64)

    u = zeros(3)
    ip = (0.0, 0.0, 0.0)

    for element in elements
        fill!(u, 0.0)
        gdofs = get_gdofs(problem, element)
        if haskey(element, "displacement")
            u[:] = element("displacement", ip, time)
        end
        for j in 1:3
            if haskey(element, "fixed displacement $j")
                dof = gdofs[j]
                add!(assembly.C1, dof, dof, 1.0)
                add!(assembly.C2, dof, dof, 1.0)
                du = element("fixed displacement $j", ip, time) - u[j]
                add!(assembly.g, dof, du)
            end
        end
    end

    return nothing

end

""" Mechanical material model analysis."""
mutable struct MecaMatSo <: AbstractAnalysis
    t0 :: Float64
    t1 :: Float64
    dt :: Float64
    u :: Vector{Float64}
    du :: Vector{Float64}
    la :: Vector{Float64}
    dla :: Vector{Float64}
    convergence_tolerance :: Float64
    max_iterations :: Int64
    extrapolate_initial_guess :: Bool
end

function MecaMatSo()
    t0 = 0.0
    t1 = 1.0
    dt = 0.1
    u = Vector{Float64}()
    du = Vector{Float64}()
    la = Vector{Float64}()
    dla = Vector{Float64}()
    convergence_tolerance = 1.0e-5
    max_iterations = 10
    extrapolate_initial_guess = false
    return MecaMatSo(t0, t1, dt, u, du, la, dla,
                     convergence_tolerance, max_iterations,
                     extrapolate_initial_guess)
end

# struct Dirichlet <: BoundaryProblem end

function FEMBase.run!(analysis::Analysis{MecaMatSo})

    props = analysis.properties
    time = props.t0
    dtime = props.dt
    u = props.u
    du = props.du
    la = props.la
    dla = props.dla

    function hasmaterialmodel(problem)
        try
            # getfield(Materials, problem.properties.material_model)
            ip = first(get_integration_points(first(get_elements(problem))))
            ym = ip("material",time).properties.youngs_modulus
            return true
        catch
            return false
        end
    end

    for problem in get_problems(analysis)
        FEMBase.initialize!(problem, time)
    end

    function operate_integration_points!(func!, time)
        func_name = last(split(string(func!),"."))
        for problem in get_problems(analysis)
            for element in get_elements(problem)
                for ip in get_integration_points(element)
                    if func_name == "material_preprocess_analysis!"
                        try
                            material_type = problem.properties.material_model
                            material = eval(:($material_type()))
                        catch
                            material = nothing
                        end
                        ip.fields["material"] = field(material)
                        func!(material, element, ip, time)
                    else
                        material = ip("material", time)
                        func!(material, element, ip, time)
                    end
                end
            end
        end
    end

    function extrapolate_next_solution()
        # extrapolate next solution, use linear approximation
        nnodes = Int(length(u)/3)
        u_dict = Dict(j => [u[3*(j-1)+k] for k in 1:3] for j in 1:nnodes)
        du_dict = Dict(j => [du[3*(j-1)+k] for k in 1:3] for j in 1:nnodes)
        for problem in get_problems(analysis)
            for element in get_elements(problem)
                connectivity = get_connectivity(element)
                # This should take into account the possibility of different time steps
                ue = tuple(collect((u_dict[j]+du_dict[j]) for j in connectivity)...)
                update!(element, "displacement", time => ue)
            end
        end
    end

    operate_integration_points!(material_preprocess_analysis!, time)

    solution_norm = 0.0
    solution_norm_prev = Inf

    while time < props.t1

        if time + dtime > props.t1
            dtime = props.t1 - time
        end
        time = round(time + dtime, digits=12)

        if props.extrapolate_initial_guess
            extrapolate_next_solution()
        else
            fill!(du, 0.0)
            solution_norm_prev = Inf
        end

        operate_integration_points!(material_preprocess_increment!, time)

        n = 0
        while true

            n += 1
            @info("Solving for time $time, iteration # $n")

            operate_integration_points!(material_preprocess_iteration!, time)

            K = SparseMatrixCOO()
            C = SparseMatrixCOO()
            f = SparseMatrixCOO()
            g = SparseMatrixCOO()
            for problem in get_problems(analysis)
                empty!(problem.assembly)
                assemble!(problem, time)
                append!(K, problem.assembly.K)
                append!(f, problem.assembly.f)
                append!(C, problem.assembly.C1)
                append!(g, problem.assembly.g)
            end

            dim = size(K, 1)

            if dim != length(u)
                resize!(u, dim)
                resize!(la, dim)
                resize!(du, dim)
                resize!(dla, dim)
                fill!(u, 0.0)
                fill!(la, 0.0)
                fill!(du, 0.0)
                fill!(dla, 0.0)
            end

            K = sparse(K, dim, dim)
            C = sparse(C, dim, dim)
            f = sparse(f, dim, 1)
            g = sparse(g, dim, 1)
            Kb = 1.0e36*C'*C
            fb = 1.0e36*C'*g

            du += Vector((cholesky(Symmetric(K+Kb)) \ (f+fb))[:])

            extrapolate_next_solution()

            # Check convergence
            solution_norm = norm(du)
            norm_err = norm(solution_norm-solution_norm_prev)
            norm_err_rel = norm_err/(max(solution_norm, solution_norm_prev))
            @info(norm_err)
            if norm_err < props.convergence_tolerance
                @info("Solution converged in $n iterations.")
                break
            else
                solution_norm_prev = solution_norm
            end
            if n > props.max_iterations
                error("Model did not converge at time $time in $n iterations.")
            end

            operate_integration_points!(material_postprocess_iteration!, time)

        end # iterations

        u += du

        operate_integration_points!(material_postprocess_increment!, time)

        JuliaFEM.write_results!(analysis, time)

        if time == props.t1
            @info("Simulation step ready.")
        end

    end # step

    operate_integration_points!(material_postprocess_analysis!, time)

    return nothing

end # analysis

"""
    get_material_analysis(material_model)

Create a standardized material test for one element. Returns a tuple:

    (analysis, problem, element, ip1)

"""
function get_one_element_material_analysis(material_model::Union{Symbol, Expr})

    X = Dict(
        1 => [0.0, 0.0, 0.0],
        2 => [1.0, 0.0, 0.0],
        3 => [1.0, 1.0, 0.0],
        4 => [0.0, 1.0, 0.0],
        5 => [0.0, 0.0, 1.0],
        6 => [1.0, 0.0, 1.0],
        7 => [1.0, 1.0, 1.0],
        8 => [0.0, 1.0, 1.0])

    element = Element(Hex8, (1, 2, 3, 4, 5, 6, 7, 8))
    elements = [element]
    update!(elements, "geometry", X)

    bc_element_1 = Element(Poi1, (1,))
    bc_element_2 = Element(Poi1, (2,))
    bc_element_3 = Element(Poi1, (3,))
    bc_element_4 = Element(Poi1, (4,))
    bc_element_5 = Element(Poi1, (5,))
    bc_element_6 = Element(Poi1, (6,))
    bc_element_7 = Element(Poi1, (7,))
    bc_element_8 = Element(Poi1, (8,))
    bc_elements = [bc_element_1, bc_element_2, bc_element_3, bc_element_4,
                   bc_element_5, bc_element_6, bc_element_7, bc_element_8]
    update!(bc_elements, "geometry", X)

    problem = Problem(Continuum3D, "1 element problem", 3)
    problem.properties.material_model = material_model
    add_elements!(problem, elements, bc_elements)
    analysis = Analysis(MecaMatSo, "solve 1 element problem")
    add_problems!(analysis, problem)

    ip = first(get_integration_points(element))
    return analysis, problem, element, bc_elements, ip
end


text:
# Holly Jackson, 2020

using BenchmarkTools
using LazySets
using LinearAlgebra
using Plots
using SparseArrays

## SET UP EXAMPLE SYSTEMS

example_num = 2

if (example_num == 1)
    vertices_x = [0.5, 1.5, 2.5, 0, 1, 2]
    all_vertices_x = [0.5, 1.5, 2.5, 0, 1, 2, 0.5, 1.5, 2.5]
    vertices_y = [2, 2, 2, 1, 1, 1]
    all_vertices_y = [2, 2, 2, 1, 1, 1, 0, 0, 0]
    edges = [[1,4],[1,5],[1,2],[2,5],[2,6],[2,3],[3,6],
        [4,7],[4,5],[5,7],[5,8],[5,6],[6,8],[6,9]]
elseif (example_num == 2)
    all_vertices_x = []
    all_vertices_y = []
    edges = []
    w = 50
    for i = 1:100
        for j = 1:w
            if (i % 2 == 0)
                push!(all_vertices_x, j-1)
                if (i > 1 && j > 1)
                    push!(edges,[(i-2)*w+j-1,(i-1)*w+j])
                end
            else
                # on odd rows offset by 0.5
                push!(all_vertices_x, j-1+0.5)
                if (i > 1 && j < w)
                    push!(edges,[(i-2)*w+j+1,(i-1)*w+j])
                end
            end
            push!(all_vertices_y, abs(i-100))
            if (j > 1)
                push!(edges,[(i-1)*w+j-1,(i-1)*w+j])
            end
            if (i > 1)
                push!(edges,[(i-2)*w+j,(i-1)*w+j])
            end

        end
    end
    vertices_x = all_vertices_x[1:length(all_vertices_x)-w]
    vertices_y = all_vertices_y[1:length(all_vertices_y)-w]
else
    # 3rd example, smaller lattice with several supernodes
    all_vertices_x = []
    all_vertices_y = []
    edges = []
    w = 20
    for i = 1:w
        for j = 1:w
            if (i % 2 == 0)
                push!(all_vertices_x, j-1)
                if (i > 1 && j > 1)
                    push!(edges,[(i-2)*w+j-1,(i-1)*w+j])
                end
            else
                # on odd rows offset by 0.5
                push!(all_vertices_x, j-1+0.5)
                if (i > 1 && j < w)
                    push!(edges,[(i-2)*w+j+1,(i-1)*w+j])
                end
            end
            push!(all_vertices_y, abs(i-100))
            if (j > 1)
                push!(edges,[(i-1)*w+j-1,(i-1)*w+j])
            end
            if (i > 1)
                push!(edges,[(i-2)*w+j,(i-1)*w+j])
            end

        end
    end
    vertices_x = all_vertices_x[1:length(all_vertices_x)-w]
    vertices_y = all_vertices_y[1:length(all_vertices_y)-w]
    # connect a few interior nodes to many other nodes
    supernodes = [[w/2,w/2],[w/2,w/4],[w/2,3*w/4]]
    # edges already connected to supernodes
    for i = 1:3
        s = Int((supernodes[i][1]-1)*w+supernodes[i][2])
        # connect to a bunch of other vertices
        for j = 1:length(vertices_x)
            if (j != s && rand() < 0.8)
                new_edge = [s,j]
                if (j < s)
                    new_edge = [j,s]
                end
                push!(edges, new_edge)
            end
        end
    end
end

N = length(vertices_x) # number of unfixed joints
S = length(edges)      # number of struts

v = zeros(2*N)         # original vertex positions
for i = 1:N
    v[2*(i-1)+1] = vertices_x[i]
    v[2*(i-1)+2] = vertices_y[i]
end

if (example_num == 1)
    f_load = 0.5
elseif (example_num == 2)
    f_load = 2.5
else
    f_load = 2
end
η = 5   # material stiffness

# -----------------------------
# Incidence matrix construction
# -----------------------------

A = spzeros(2*N, 2*S)
b_fix = spzeros(2*S)
for i = 1:S
    n1 = edges[i][1]#Int(edges[i,1])
    if (n1 <= N) # not fixed
        A[2*(n1-1)+1,2*(i-1)+1] = 1
        A[2*(n1-1)+2,2*(i-1)+2] = 1
    else
        b_fix[2*(i-1)+1] = all_vertices_x[n1]
        b_fix[2*(i-1)+2] = all_vertices_y[n1]
    end
    n2 = edges[i][2]#Int(edges[i,2])
    if (n2 <= N) # not fixed
        A[2*(n2-1)+1,2*(i-1)+1] = -1
        A[2*(n2-1)+2,2*(i-1)+2] = -1
    else
        b_fix[2*(i-1)+1] = -all_vertices_x[n2]
        b_fix[2*(i-1)+2] = -all_vertices_y[n2]
    end
end

# print(A)
A_T = transpose(A)
# spy(A,legend=false)#, markersize = 6, c = :reds)

f_L = spzeros(2*N)
if (example_num == 1)
    f_L[2*3] = -f_load
    f_L[2*3+1] = -f_load
else # (example_num == 2 or 3)
    f_L[2*w] = -f_load
    f_L[2*w+1] = -f_load
end

# --------------------------
# Helper matrix construction
# --------------------------

α = spzeros(2*S, 2*S)
for i = 1:S
    α[2*(i-1)+1,2*(i-1)+1] = 1
    α[2*(i-1)+1,2*(i-1)+2] = α[2*(i-1)+2,2*(i-1)+1] = 1
    α[2*(i-1)+2,2*(i-1)+2] = 1
end
α_T = transpose(α)
# print(α)
# spy(α, markersize = 6, c = :reds)

# CALCULATE NOMINAL LENGTHS
L0 = sqrt.(α*(A_T*v + b_fix).^2) # nominal lengths

# ------------------------
# Final system formulation
# ------------------------

function f(u)
    L = sqrt.(α * (A_T * (u + v) + b_fix).^2)
    A * (η * (A_T * (u + v) + b_fix) ./ L .* (L0 - L)) - f_L
end

A_dense = Array(A)
A_T_dense = Array(A_T)
v_dense = Array(v)
b_fix_dense = Array(b_fix)
α_dense = Array(α)
L0_dense = Array(L0)
f_L_dense = Array(f_L)
function f_dense(u)
    L = sqrt.(α_dense * (A_T_dense * (u + v_dense) + b_fix_dense).^2)
    A * (η * (A_T_dense * (u + v_dense) + b_fix_dense) ./ L .* (L0_dense - L)) - f_L_dense
end

println("A forward pass through f(u) takes...")
@btime f(rand(2*N))

## FORWARD MOE AD

function forward_mode_ad(D0, Σ)
    # w = A_T * (u .+ v) .+ b_fix
    w_D0 = A_T * (D0 .+ v) .+ b_fix
    w_Σ = A_T * Σ
    # L = sqrt.(α * (w .* w))
    D0 = α * (w_D0 .* w_D0)
    L_Σ = (1/2) .* (D0) .^ (-1/2) .* (α * (2 .* (w_D0) .* w_Σ))
    L_D0 = sqrt.(D0)
    # s = η .* w ./ L .* (L0 .- L)
    # f = A * s .- f_L
    D0 = A * (η .* w_D0 ./ L_D0 .* (L0 .- L_D0)) .- f_L
    Σ =  A * (η ./ L_D0 .* (L0 .- L_D0) .* w_Σ + η .* w_D0 .* L0 .* (-1) .* (L_D0) .^(-2) .* L_Σ)
    return [D0, Σ]
end

using FiniteDifferences
using Distributed
function forward_mode_ad!(J, D0, Σ)
    @distributed for i = 1:size(Σ,2)
        J[i,:] = central_fdm(5, 1)(ε -> f(D0 .+ ε .* Σ[:,i]), 0)
    end
end

u0 = rand(2*N)
using ForwardDiff
println("Julia ForwardDiff.jl takes...")
@btime ForwardDiff.jacobian(u -> f(u), u0)
println("My forward mode AD takes...")
Σ = spzeros(2*N,2*N) + I
@btime forward_mode_ad(u0, Σ)

## REVERSE MODE AD

# Pullback function on vector λ
function pullback(λ, u)
    w_bar = η .* A_T * λ
    z = A_T * (u .+ v) + b_fix
    x = z.^2
    y = α * x
    y_bar = w_bar .* (z .* (-1/2) .* y.^(-3/2) .* (L0 - sqrt.(y))
        + z ./ sqrt.(y) .* (-(1/2) * y.^(-1/2)))
    x_bar = α_T * y_bar
    z_bar = x_bar .* 2 .* z + w_bar .* (1 ./ sqrt.(y) .* (L0 - sqrt.(y)))
    return A * z_bar
end

function reverse_mode_ad(f, u)
    return pullback(spzeros(2*N,2*N) + I,u)
end

using ReverseDiff
println("Julia ReversedDiff.jl takes...")
@btime ReverseDiff.jacobian(u -> f_dense(u), u0)
println("My reverse mode AD takes...")
@btime reverse_mode_ad(f, u0)

## ASSIGN MATRIX COLORS

function assign_matrix_colors(S, rows_to_ignore, forward)
    if (!forward)
        S = transpose(S)
    end
    colors = Array{Int64,1}(zeros(2*N))
    color_num = 1
    colors[1] = color_num # column to color_num
    @inbounds for i = 2:2*N
        # check collisions in every row
        colliding_colors = Set()
        @inbounds for j = 1:2*N
            if (j in rows_to_ignore)
                continue
            end
            if (S[j,i] != 0)
                # col = find_nz(S[j,1:(i-1)])
                # col = S[j,1:(i-1)]
                for k = 1:(i-1)
                    if (S[j,k] != 0)
                        push!(colliding_colors,colors[k])
                    end
                end
            end
        end
        # are there any colors that didn't collide
        need_new_col = true
        @inbounds for j = 1:color_num
            if !(j in colliding_colors)
                valid_color = j
                need_new_col = false
                break
            end
        end
        if (need_new_col)
            valid_color = color_num + 1
            color_num = color_num + 1
        end
        colors[i] = valid_color
    end
    return colors, color_num
end

function decompress!(J, J_compress, colors, S, forward)
    @inbounds for i = 1:2*N
        col = S[:,i]
        J[:,i] = J_compress[:,colors[i]] .* col
    end
end

function matrix_colors(SP, rows_to_ignore, forward)
    # figure out most dense rows, let's say rows that are more than 80% full
    c, max_color = assign_matrix_colors(SP, rows_to_ignore, forward)
    if (forward)
        Σ_color = spzeros(2*N,max_color)
    else
        Σ_color = spzeros(max_color,2*N)
    end
    @inbounds for i = 1:2*N
        if (forward)
            Σ_color[i,c[i]] = 1.0
        else
            Σ_color[c[i],i] = 1.0
        end
    end
    return [Σ_color, c]
end

## FORWARD MODE AD WITH MATRIX COLORS

# Get sparsity pattern
Σ =  spzeros(2*N,2*N) + I
SP = Array(forward_mode_ad(rand(2*N), Σ)[2])#ForwardDiff.jacobian(u -> f(u), rand(2*N))ReverseDiff.jacobian(f, input
SP[SP .> 0] .= 1.0; SP[SP .< 0] .= 1.0;
spy(SP,legend=false)
# savefig("bone_sparsity2.png")

MC = matrix_colors(SP,Set(),true)
Σ_color = MC[1]
c = MC[2]

MC_rev = matrix_colors(SP,Set(),false)
Σ_color_rev = MC[1]
c_rev = MC[2]

function forward_mode_AD_mat_color!(J, u, c, Σ_color)
    D = forward_mode_ad(u, Σ_color)
    decompress!(J, D[2], c, SP, true)
end

function reverse_mode_AD_mat_color!(J, u, c, Σ_color)
    D = pullback(Σ_color, u)
    decompress!(J, D, c, SP, false)
end

println("FORWARD SPARSE")
J = spzeros(2*N,2*N)
@btime forward_mode_AD_mat_color!(J, u0, c, Σ_color)
println("REVERSE SPARSE")
J = spzeros(2*N,2*N)
@btime reverse_mode_AD_mat_color!(J, u0, c_rev, Σ_color_rev)

## COMBINED MODE AD

function combined_mode_AD(u, c, Σ_color, λ, dense_rows)
    D = forward_mode_ad(u, Σ_color)
    J = spzeros(2*N,2*N)
    decompress!(J, D[2], c, SP, true)
    # fill in dense rows
    R = pullback(λ,u)
    count = 1
    for r in dense_rows
        J[r,:] = R[:,count]
        count = count + 1
    end
    return J
end

# find rows to ignore that are dense
rows_to_ignore = Set()
for i = 1:2*N
    if (sum(SP[i,:])/(2*N) > 0.5)
        push!(rows_to_ignore,i)
    end
end
println("DENSE ROWS: ",rows_to_ignore)

MC = matrix_colors(SP,rows_to_ignore,true)
Σ_color = MC[1]
c = MC[2]

# Stack vectors so all dense rows can be calculated at once
λ = zeros(length(rows_to_ignore),2*N)
count = 1
for r in rows_to_ignore
    λ[count,:] = LazySets.Arrays.SingleEntryVector(r,2*N,1)
    global count = count + 1
end
λ = transpose(λ)

println("COMBINED MODE AD, FORWARD")
@btime combined_mode_AD(u0, c, Σ_color, λ, rows_to_ignore)

function combined_mode_AD_reverse(u, c, Σ_color, λ, dense_cols)
    J = spzeros(2*N,2*N)
    reverse_mode_AD_mat_color!(J, u, c, Σ_color)
    # fill in dense rows
    C = forward_mode_ad(u,λ)
    C = C[2]
    count = 1
    for c in dense_cols
        J[:,c] = C[:,count]
        count = count + 1
    end
    return J
end

# find rows to ignore that are dense
cols_to_ignore = Set()
for i = 1:2*N
    if (sum(SP[:,i])/(2*N) > 0.5)
        push!(cols_to_ignore,i)
    end
end
println("DENSE COLS: ",cols_to_ignore)

MC_rev = matrix_colors(SP,cols_to_ignore,false)
Σ_color_rev = MC[1]
c_rev = MC[2]

# Stack vectors so all dense rows can be calculated at once
λ = zeros(2*N,length(cols_to_ignore))
count = 1
for c in cols_to_ignore
    λ[:,count] = LazySets.Arrays.SingleEntryVector(c,2*N,1)
    global count = count + 1
end

println("COMBINED MODE AD, REVERSE")
@btime combined_mode_AD_reverse(u0, c_rev, Σ_color_rev, λ, cols_to_ignore)

## SOLVE WITH NETWON'S METHOD

# ----------------
# Newton iteration
# ----------------

function newton_step(f, u0)
    # -----------------------------------------------
    # -----------------------------------------------
    # REMEMBER TO SET THIS ACCORDING TO METHOD YOU WANT TO USE
    J = reverse_mode_ad(f, u0)
    # -----------------------------------------------
    # -----------------------------------------------
    # J = combined_mode_ad(u0, c, Σ_color, rows_to_ignore)
    # create the (sparse) LU (etc.) factorization
    F = factorize(J)
    # solve J d = f(u0) with LU decompsition
    d = zeros(2*N)
    ldiv!(d, F, f(u0))
    return u0 - d
end

function newton(f, u0, max_iter)
    u_last = u0

    for i in 1:max_iter
        u = newton_step(f, u_last)
        @show u
        if LinearAlgebra.norm(f(u),2) < 1e-10 && LinearAlgebra.norm(u-u_last,2) < 1e-10
            print("converged in ",i," iterations")
            return u
        else
            println(LinearAlgebra.norm(f(u),2))
            println(LinearAlgebra.norm(u-u_last,2))
        end
        u_last = u
    end

    return u_last
end

## EVALUATE AND PLOT

# initial guess
u0 = 0.0001*ones(2*N)
u = newton(f, u0, 100)

# Plot results
new_vertices_x = zeros(length(all_vertices_x))
new_vertices_y = zeros(length(all_vertices_y))
for i = 1:N
    new_vertices_x[i] = all_vertices_x[i] - u[2*(i-1)+1]
    new_vertices_y[i] = all_vertices_y[i] - u[2*(i-1)+2]
end
for i = N+1:length(all_vertices_x)
    new_vertices_x[i] = all_vertices_x[i]
    new_vertices_y[i] = all_vertices_y[i]
end

gr()
p = plot(title = string("Deformation of Lattice Under ",f_load, " Units Force"), legend = false, aspect_ratio = 1)
for i = 1:S
    strut1 = edges[i][1]
    strut2 = edges[i][2]
    plot!(p, [all_vertices_x[strut1], all_vertices_x[strut2]], [all_vertices_y[strut1], all_vertices_y[strut2]], color = :steelblue)
    plot!(p, [new_vertices_x[strut1] .+ 100, new_vertices_x[strut2] .+ 100], [new_vertices_y[strut1], new_vertices_y[strut2]], color = :green)
end
# scatter!(p, all_vertices_x[N+1:end], all_vertices_y[N+1:end], markersize = 1, color = :red)
# scatter!(p, vertices_x, vertices_y, markersize = 1, color = :steelblue)
# scatter!(p, new_vertices_x[1:N] .+ 100, new_vertices_y[1:N], markersize = 1, color = :green)
p

# plot
savefig(p, "example.png")

max_stars_repo_name: hollyjackson/AutomaticDifferentiation

max_issues_repo_name: hollyjackson/AutomaticDifferentiation

max_forks_repo_name: hollyjackson/AutomaticDifferentiation