text:
using JSON
import Base: ==

struct LJH3File{FramePeriod}
    io::IOStream
    index::Vector{Int} # points to the start of records, and potential next record
    header::OrderedDict{String,Any}
end
"""    LJH3File(fname::AbstractString)
Open an LJH3 file. Records can be accessed by indexing with integers like an `AbstractVector`,
eg `ljh[1]` or by iteration eg `collect(ljh)` or `for record in ljh dosomething() end`.
Due to the possibility of unequal record lengths, random access will be slow until
the file is fully indexed. Indexing occurs automatically whenever you access a record.
An `LJH3File` presents a dictionary like interface for accessing the header,
eg `ljh["File Format Version"]` or `keys(ljh)`. Use `frameperiod(ljh)` to get the
time between succesive samples in seconds."""
LJH3File(fname::AbstractString) = LJH3File(open(fname,"r"))
function LJH3File(io::IO; shouldseekstart=true)
    shouldseekstart && seekstart(io)
    header = JSON.parse(io, dicttype=OrderedDict)
    frameperiod = header["frameperiod"]
    @assert read(io, Char) == '\n'
    @assert header["File Format"]=="LJH3"
    @assert header["File Format Version"] == "3.0.0"
    LJH3File{frameperiod}(io,Int[position(io)], header)
end
frameperiod(ljh::LJH3File{FramePeriod}) where FramePeriod = FramePeriod
Base.close(ljh::LJH3File) = close(ljh.io)
filename(ljh::LJH3File) = ljh.io.name[7:end-1]
progresssize(ljh::LJH3File) = stat(ljh.io).size
progressposition(ljh::LJH3File) = position(ljh.io)
"    ljhopen(fname::AbstractString; kwargs...)
Open file `fname` as an `LJH3File` or `LJHFile` depending on contents."
function ljhopen(fname::AbstractString; kwargs...)
    try
        return LJH3File(fname; kwargs...)
    catch
        try
            return LJHFile(fname; kwargs...)
        catch
            error("$fname failed to open as LJH or LJH3")
        end
    end
end
"`LJH3Record` are returned when accessing a record in an `LJH3File`. Use functions
`data(r)`, `first_rising_sample(r)`, count(r)`, and `timestamp_usec(r)` to extract
information from a record `r`."
struct LJH3Record{FramePeriod}
    data::Vector{UInt16}
    first_rising_sample::Int32
    frame1index::Int64
    timestamp_usec::Int64
end
data(r::LJH3Record) = r.data
first_rising_sample(r::LJH3Record) = r.first_rising_sample
frame1index(r::LJH3Record) = r.frame1index
timestamp_usec(r::LJH3Record) = r.timestamp_usec
frameperiod(r::LJH3Record{FramePeriod}) where FramePeriod = FramePeriod
Base.length(r::LJH3Record) = length(r.data)
==(a::LJH3Record, b::LJH3Record) = a.data == b.data && a.first_rising_sample == b.first_rising_sample && a.frame1index == b.frame1index && a.timestamp_usec == b.timestamp_usec
"    write(ljh::LJH3File, trace::Vector{UInt16},first_rising_sample, frame1index::Int64, timestamp_usec::Int64)
Write a single record to `ljh`. Assumes `ljh.io` is at the same position it would be if you had
called `seekend(ljh.io)`."
function Base.write(ljh::LJH3File, trace::Vector{UInt16},first_rising_sample, frame1index::Int64, timestamp_usec::Int64)
    write(ljh.io, Int32(length(trace)), Int32(first_rising_sample), frame1index, timestamp_usec, trace)
    push!(ljh.index,position(ljh.io))
end
"""    create3(filename::AbstractString, frameperiod, header_extra = Dict();version="3.0.0")
Return an `LJH3File` ready for writing with `write`. The header will contain "frameperiod",
"File Format", "File Format Version" and any items in `header_extra`. Items in `header_extra`
will overwrite the header items passed as arguments, and you can make an invalid LJH3 file
this way. You are advised to avoid creating invalid LJH3 files.
"""
function create3(filename::AbstractString, frameperiod, header_extra = Dict();version="3.0.0")
    io = open(filename,"w+")
    header = OrderedDict{String,Any}()
    header["File Format"] = "LJH3"
    header["File Format Version"] = version
    header["frameperiod"]=frameperiod
    for (k,v) in header_extra
        header[k]=v
    end
    JSON.print(io, header, 4) # last argument uses pretty printing
    print(io,"\n")
    LJH3File(io)
end
function seekto(ljh::LJH3File, i::Int)
    if length(ljh.index) >= i
        @inbounds seek(ljh.io, ljh.index[i])
    else
        error("LJH3 Bounds Error")
    end
end
function _readrecord(ljh::LJH3File{FramePeriod},i) where FramePeriod
    num_samples = read(ljh.io, Int32)
    first_rising_sample = read(ljh.io, Int32)
    frame1index = read(ljh.io, Int64)
    timestamp_usec = read(ljh.io, Int64)
    data = Array{UInt16}(undef, num_samples)
    read!(ljh.io, data)
    if i==length(ljh.index)
        deltapos = 2*num_samples + 24
        @inbounds push!(ljh.index,ljh.index[i]+deltapos)
        # now, ljh.index[end] is the first unobserved byte offset into ljh.io
        # it may or may not be the start of a pulse
    end
    LJH3Record{FramePeriod}(data, first_rising_sample, frame1index, timestamp_usec)
end
function tryread(ljh::LJH3File{FramePeriod}) where FramePeriod
    d1 = read(ljh.io,4)
    if length(d1)<4
        seek(ljh.io, position(ljh.io)-length(d1)) # go back to the start of the record
        return Nullable{LJH3Record{FramePeriod}}()
    end
    num_samples = reinterpret(Int32,d1)[1]
    rest_of_record_length = 20+2*num_samples
    d2 = read(ljh.io,rest_of_record_length)
    if length(d2) < rest_of_record_length
        seek(ljh.io, position(ljh.io)-length(d1)-length(d2)) # go back to the start of the record
        return Nullable{LJH3ecord{FramePeriod}}()
    end
    first_rising_sample = reinterpret(Int32,d2[1:4])[1]
    frame1index = reinterpret(Int64,d2[5:12])[1]
    timestamp_usec = reinterpret(Int64,d2[13:20])[1]
    data = reinterpret(UInt16,d2[21:end])
    Nullable(LJH3Record{FramePeriod}(data,first_rising_sample, frame1index, timestamp_usec))
end
# Array interface
function Base.getindex(ljh::LJH3File,i::Int)
    seekto(ljh, i)
    _readrecord(ljh,i)
end
Base.size(ljh::LJH3File) = (length(ljh),)
function Base.length(ljh::LJH3File)
    # iterate from the last entry in index to the end of the file
    # to fill ljh.index
    nknown = length(ljh.index)
    state = nknown, stat(ljh.io).size
    seekto(ljh, nknown)
    next = Base.iterate(ljh, state)
    while next !== nothing
        (i, state) = next
        next = Base.iterate(ljh, state)
    end
    return length(ljh.index)-1
end
Base.lastindex(ljh::LJH3File) = length(ljh)
# iterator interface
function Base.iterate(ljh::LJH3File, state=nothing)
    if state == nothing
        seekto(ljh,1)
        state = 1, stat(ljh.io).size
    end
    j,sz = state
    if ljh.index[j] ≥ sz
        return nothing
    end
    _readrecord(ljh,j),(j+1,sz)
end
Base.IteratorSize(ljh::LJH3File) = Base.SizeUnknown()
# getindex with strings
Base.getindex(ljh::LJH3File,key::AbstractString) = ljh.header[key]
Base.keys(ljh::LJH3File) = keys(ljh.header)
Base.values(ljh::LJH3File) = values(ljh.header)
Base.eltype(ljh::LJH3File) = LJH3Record{frameperiod(ljh)}


text:
# open Demo/ViewController/ViewController.xcworkspace
# Run

using Swifter

iphone = initial("http://192.168.0.9:8080")
simul = initial("http://localhost:8080")

@query iphone.view.backgroundColor = UIColor.cyanColor()
sleep(0.5)
@query simul.view.backgroundColor = iphone.view.backgroundColor
sleep(0.5)

for vc in [iphone, simul]
    @query vc.view.backgroundColor = UIColor.orangeColor()
    sleep(0.5)
end

@query simul.view.backgroundColor = UIColor.blueColor()
sleep(0.5)
@query iphone.view.backgroundColor = simul.view.backgroundColor
sleep(0.5)

for vc in [iphone, simul]
    @query vc.view.backgroundColor = UIColor.whiteColor()
    sleep(0.5)
end

@query iphone.label.text = "Julia"
sleep(0.5)
@query simul.label.text = "Swift"
sleep(0.5)

temporal = @query iphone.label.text
sleep(0.5)
@query iphone.label.text = simul.label.text
sleep(0.5)
@query simul.label.text = temporal
sleep(0.5)

for vc in [iphone, simul]
    @query vc.label.text = "Julia & Swift"
    sleep(0.5)
end

map([iphone, simul]) do vc
    @query vc.label.text
end

ext: jl

max_stars_repo_name: wookay/Console

max_issues_repo_name: wookay/Console

max_forks_repo_name: wookay/Console


text:
using Pkg

ts = ["Combinatorics",
      "Distributions",
      "KrylovKit",
      "LaTeXStrings",
      "LinearAlgebra",
      "Printf",
      "Random",
      "StatsBase",
      "TimerOutputs",
      "WignerSymbols"]
for tmp in ts
    Pkg.add(tmp)
end


text:
import CUDA
include("models.jl")
using MAT, OffsetArrays, FFTViews, ArgParse, Images, ImageMagick, Knet
include("./deeper-cut/loss.jl")
include("./deeper-cut/deeper-cut.config.jl")
include("./utils.jl")

function get_params(params, atype)
    len = length(params["value"])
    ws, ms = [], []
    for k = 1:len
        name = params["name"][k]
        value = convert(Array{Float32}, params["value"][k])

        if endswith(name, "moments")
            push!(ms, reshape(value[:, 1], (1, 1, size(value, 1), 1)))
            push!(ms, reshape(value[:, 2], (1, 1, size(value, 1), 1)))
        elseif startswith(name, "bn")
            push!(ws, reshape(value, (1, 1, length(value), 1)))
        elseif startswith(name, "fc") && endswith(name, "filter")
            push!(ws, transpose(reshape(value, (size(value, 3), size(value, 4)))))
        elseif startswith(name, "conv") && endswith(name, "bias")
            push!(ws, reshape(value, (1, 1, length(value), 1)))
        else
            push!(ws, value)
        end
    end
    map(wi -> convert(atype, wi), ws), map(mi -> convert(atype, mi), ms)
end

function get_params_from_deepercut_pretrained(atype)
    groups_path = "/kuacc/users/gsoykan20/comp541_term_project/deeper-cut/deepercut-weights/test.mat"
    file = matopen(groups_path)
    params = read(file, "deepercut_weights")
    close(file)

    len = length(params["value"])
    ws, ms = [], []
    for k = 1:len
        name = params["name"][k]
        value = convert(Array{Float32}, params["value"][k])

        if endswith(name, "moments")
           # value = reverse_all_dims(value; until_dim=2)
            push!(ms, reshape(value[:, 1], (1, 1, size(value, 1), 1)))
            push!(ms, reshape( sqrt.(value[:, 2]), (1, 1, size(value, 1), 1)))
        elseif endswith(name, "gamma:0_mult") || endswith(name, "beta:0_bias")
            # value = reverse_all_dims(value; until_dim=3)
            push!(ws, reshape(value, (1, 1, length(value), 1)))
        elseif endswith(name, "biases:0_bias")
            value = reverse_all_dims(value; until_dim=3)
            push!(ws, reshape(value, (1, 1, length(value), 1)))
        else
           # permuted_value = permutedims( value, [2,1, 3, 4])
            value = reverse_all_dims(value; until_dim=3)
            push!(ws, value)
        end
    end
    w, ms = map(wi -> convert(atype, wi), ws), map(mi -> convert(atype, mi), ms)
end

function get_modular_resnet(
    should_use_resnet50::Bool;
    should_use_resnet101::Bool = false,
    use_deepercut_resnet101_pretrained = false,
)
    if should_use_resnet101
        model_file_path = "/userfiles/gsoykan20/resnet_pretrained/imagenet-resnet-101-dag.mat"
    elseif should_use_resnet50
        model_file_path = "/userfiles/gsoykan20/resnet_pretrained/imagenet-resnet-50-dag.mat"
    else
        model_file_path = "/userfiles/gsoykan20/resnet_pretrained/imagenet-resnet-152-dag.mat"
    end

    if use_deepercut_resnet101_pretrained
         o = Dict(:atype => KnetArray{Float32}, :model => model_file_path, :top => 10)
        model = matread(abspath(o[:model]))
       # _, ms = get_params(model["params"], o[:atype])

        w, ms = get_params_from_deepercut_pretrained(KnetArray{Float32})
    else
        o = Dict(:atype => KnetArray{Float32}, :model => model_file_path, :top => 10)
        model = matread(abspath(o[:model]))
        w, ms = get_params(model["params"], o[:atype])
    end

    if should_use_resnet101
        modular_resnet = generate_headless_resnet101_from_weights(w, ms)
    elseif should_use_resnet50
        modular_resnet = generate_headless_resnet50_from_weights(w, ms)
    else
        modular_resnet = generate_headless_resnet_from_weights(w, ms)
    end

    return modular_resnet
end

function generate_resnet_from_weights(w, ms)
    conv1 = ResLayerX1(w[1:3], ms; padding = 3, stride = 2, is_initial = true)
    r2 = ResLayerX5(w[4:33], ms; strides = [1, 1, 1, 1])
    r3 = ResLayerX5(w[34:108], ms)
    r4 = ResLayerX5(w[109:435], ms)
    r5 = ResLayerX5(w[436:465], ms; is_next_fc = true)
    fc = Dense(w[466], w[467])
    return Chain(conv1, r2, r3, r4, r5, fc)
end

function generate_headless_resnet_from_weights(w, ms)
    conv1 = ResLayerX1(w[1:3], ms; padding = 3, stride = 2, is_initial = true)
    r2 = ResLayerX5(w[4:33], ms; strides = [1, 1, 1, 1])
    r3 = ResLayerX5(w[34:108], ms; tag = 3, is_conv3_for_deepercut = true)
    r4 = ResLayerX5(w[109:435], ms)
    r5 = ResLayerX5(
        w[436:465],
        ms;
        tag = 5,
        strides = [1, 1, 1, 1],
        b_layer_dilations = [1, 2, 1],
        b_layer_pads = [0, 2, 0],
    )
    return Chain(conv1, r2, r3, r4, r5)
end

function generate_resnet50_from_weights(w, ms)
    layer1 = ResLayerX1_50(w[1:4], ms)
    r2 = ResLayerX5(w[5:34], ms; strides = [1, 1, 1, 1])
    r3 = ResLayerX5(w[35:73], ms)
    r4 = ResLayerX5(w[74:130], ms)
    r5 = ResLayerX5(w[131:160], ms; is_next_fc = true)
    fc = Dense(w[161], w[162])
    return Chain(layer1, r2, r3, r4, r5, fc)
end

# The output size is compatable with Upsampled Conv5 bank
function generate_headless_resnet50_from_weights_until_3rd_bank(w, ms)
    layer1 = ResLayerX1_50(w[1:4], ms)
    r2 = ResLayerX5(w[5:34], ms; strides = [1, 1, 1, 1])
    r3 = ResLayerX5(w[35:73], ms)
    return Chain(layer1, r2, r3)
end

function generate_headless_resnet50_from_weights(w, ms)
    layer1 = ResLayerX1_50(w[1:4], ms)
    r2 = ResLayerX5(w[5:34], ms; strides = [1, 1, 1, 1])
    r3 = ResLayerX5(w[35:73], ms; tag = 3, is_conv3_for_deepercut = true)
    r4 = ResLayerX5(w[74:130], ms)
    r5 = ResLayerX5(
        w[131:160],
        ms;
        tag = 5,
        strides = [1, 1, 1, 1],
        b_layer_dilations = [1, 2, 1],
        b_layer_pads = [0, 2, 0],
    )
    return Chain(layer1, r2, r3, r4, r5)
end

function generate_headless_resnet101_from_weights(w, ms)
    # layer1 = ResLayerX1_50(w[1:4], ms)
    layer1 = ResLayerX1(w[1:3], ms; padding = 3, stride = 2, is_initial = true)
    r2 = ResLayerX5(w[4:33], ms; strides = [1, 1, 1, 1])
    r3 = ResLayerX5(w[34:72], ms; tag = 3, is_conv3_for_deepercut = true)
    r4 = ResLayerX5(w[73:282], ms)
    r5 = ResLayerX5(
        w[283:312],
        ms;
        tag = 5,
        strides = [1, 1, 1, 1],
        b_layer_dilations = [1, 2, 1],
        b_layer_pads = [0, 2, 0],
    )
    return Chain(layer1, r2, r3, r4, r5)
end

function generate_deeper_cut(;
    should_use_resnet101 = false,
    should_use_resnet50 = true,
    is_loc_ref_enabled = false,
    connect_res3_to_res5 = true,
    use_deepercut_resnet101_pretrained = false,
)
    modular_resnet = get_modular_resnet(
        should_use_resnet50;
        should_use_resnet101 = should_use_resnet101,
        use_deepercut_resnet101_pretrained = use_deepercut_resnet101_pretrained,
    )

    if use_deepercut_resnet101_pretrained
        head_weights, _ = get_params_from_deepercut_pretrained(KnetArray{Float32})
        p_d_w = head_weights[end-3] #permutedims(head_weights[end-3], [2, 1, 3, 4])
        p_d_b = head_weights[end-2]
        l_r_w = head_weights[end-1] #permutedims(head_weights[end-1], [2, 1, 3, 4])
        l_r_b = head_weights[end]
        deepercut_head = DeeperCutHead(
             p_d_w,
             p_d_b,
             l_r_w,
            l_r_b;
            is_loc_ref_enabled = is_loc_ref_enabled,
        )
    else
        deepercut_head = DeeperCutHead(;
            part_detection_head = Deconv(
                3,
                3,
                global_num_joints,
                2048;
                padding = 1,
                stride = 2,
                tag = "part_detect_deconv",
            ),
            loc_ref_head = Deconv(
                3,
                3,
                global_num_joints * 2,
                2048;
                padding = 1,
                stride = 2,
                tag = "loc_ref_deconv",
            ),
            is_loc_ref_enabled = is_loc_ref_enabled,
        )
    end

    deeper_cut = Chain(
        modular_resnet.layers...,
        deepercut_head;
        loss = deeper_cut_combined_loss,
        deeperCutOption = DeeperCutOption(; connect_res3_to_res5 = connect_res3_to_res5),
        lambda2 = weight_decay
    )
    return deeper_cut
end

#=
# mode, 0=>train, 1=>test
function resnet50(w,x,ms; mode=1)
    # layer 1
    conv1  = conv4(w[1],x; padding=3, stride=2) .+ w[2]
    bn1    = batchnorm(w[3:4],conv1,ms; mode=mode)
    pool1  = pool(bn1; window=3, stride=2)

    # layer 2,3,4,5
    r2 = reslayerx5(w[5:34], pool1, ms; strides=[1,1,1,1], mode=mode)
    r3 = reslayerx5(w[35:73], r2, ms; mode=mode)
    r4 = reslayerx5(w[74:130], r3, ms; mode=mode) # 5
    r5 = reslayerx5(w[131:160], r4, ms; mode=mode)

    # fully connected layer
    pool5  = pool(r5; stride=1, window=7, mode=2)
    fc1000 = w[161] * mat(pool5) .+ w[162]
end

# mode, 0=>train, 1=>test
function resnet101(w,x,ms; mode=1)
    # layer 1
    conv1 = reslayerx1(w[1:3],x,ms; padding=3, stride=2, mode=mode)
    pool1 = pool(conv1; window=3, stride=2)

    # layer 2,3,4,5
    r2 = reslayerx5(w[4:33], pool1, ms; strides=[1,1,1,1], mode=mode)
    r3 = reslayerx5(w[34:72], r2, ms; mode=mode)
    r4 = reslayerx5(w[73:282], r3, ms; mode=mode)
    r5 = reslayerx5(w[283:312], r4, ms; mode=mode)

    # fully connected layer
    pool5  = pool(r5; stride=1, window=7, mode=2)
    fc1000 = w[313] * mat(pool5) .+ w[314]
end

# mode, 0=>train, 1=>test
function resnet152(w,x,ms; mode=1)
    # layer 1
    conv1 = reslayerx1(w[1:3],x,ms; padding=3, stride=2, mode=mode)
    pool1 = pool(conv1; window=3, stride=2)

    # layer 2,3,4,5
    r2 = reslayerx5(w[4:33], pool1, ms; strides=[1,1,1,1], mode=mode)
    r3 = reslayerx5(w[34:108], r2, ms; mode=mode)
    r4 = reslayerx5(w[109:435], r3, ms; mode=mode)
    r5 = reslayerx5(w[436:465], r4, ms; mode=mode)

    # fully connected layer
    pool5  = pool(r5; stride=1, window=7, mode=2)
    fc1000 = w[466] * mat(pool5) .+ w[467]
end

function reslayerx0(w,x,ms; padding=0, stride=1, mode=1)
    b  = conv4(w[1],x; padding=padding, stride=stride)
    bx = batchnorm(w[2:3],b,ms; mode=mode)
end

function reslayerx1(w,x,ms; padding=0, stride=1, mode=1)
    relu.(reslayerx0(w,x,ms; padding=padding, stride=stride, mode=mode))
end
=#


text:
module SphericalHarmonicArrays
using Base: tail, @propagate_inbounds
using TupleTools
using SphericalHarmonicModes
using SphericalHarmonicModes: ModeRange

export SHArray
export SHVector
export SHMatrix

include("errors.jl")

const RangeOrInteger = Union{ModeRange, AbstractUnitRange, Integer}
const OneBasedAxisType = Union{ModeRange, Integer, Base.OneTo}
const DimOrInd = Union{Integer, AbstractUnitRange}
const RangeOrModeRange = Union{AbstractUnitRange, ModeRange}
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}

# Methods from Unrolled.jl
type_length(tup::Type{T}) where {T<:Tuple} = length(tup.parameters)
function _unrolled_filter(f, tup)
    :($([Expr(:(...), :(f(tup[$i]) ? (tup[$i],) : ()))
         for i in 1:type_length(tup)]...),)
end
function _unrolled_filterdims(f, tup)
    :($([Expr(:(...), :(f(tup[$i]) ? ($i,) : ()))
         for i in 1:type_length(tup)]...),)
end
@generated unrolled_filter(f, tup) = _unrolled_filter(f, tup)
@generated unrolled_filterdims(f, tup) = _unrolled_filterdims(f, tup)

"""
    s = SHArray(arr::AbstractArray{T,N}, modes::NTuple{N, Union{AbstractUnitRange, SphericalHarmonicModes.ModeRange}}) where {T,N}
    s = SHArray(arr::AbstractArray{T,N}, modes::Vararg{Union{AbstractUnitRange, SphericalHarmonicModes.ModeRange}, N}) where {T,N}

Create a wrapper around an array such that certain dimensions may be indexed using
`Tuple` of spherical harmonic modes. Often the indices would be spherical harmonic
degrees `(l,m)`. The argument `modes` dictates the map between spherical harmonic modes and
the indices of the parent array. The indices of `modes` that are of a `ModeRange` type
correspond to the dimensions of `s` that enable indexing with a `Tuple` of
spherical harmonic degrees.

Use the iterators provided by `SphericalHarmonicModes` in `modes` to generate the map between
the indices of the array and `Tuple`s of spherical harmonic degrees.
The resulting map would be of the form
`s[collect(modes[1])[i], collect(modes[2])[j], ...] == s[i,j, ...]`.

# Examples

```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> s = SHArray(reshape(1:4, 2, 2), (LM(0:1, 0:0), 1:2))
2×2 SHArray(reshape(::UnitRange{Int64}, 2, 2), (LM(0:1, 0:0), 1:2)):
 1  3
 2  4

julia> s.modes[1] |> collect
2-element $(Array{Tuple{Int64,Int64},1}):
 (0, 0)
 (1, 0)

julia> s[(1,0), 1]
2
```

See also: [`SHVector`](@ref), [`SHMatrix`](@ref)
"""
struct SHArray{T, N, AA<:AbstractArray{T,N}, TM<:NTuple{N,RangeOrInteger}} <: AbstractArray{T, N}
    parent :: AA
    modes :: TM

    function SHArray{T,N,AA,TM}(arr::AA, modes::TM) where {T, N, AA<:AbstractArray{T,N}, TM<:NTuple{N,RangeOrInteger}}
        foreach(checkaxes, axes(arr), map(moderangeaxes, modes))
        new{T,N,AA,TM}(arr, modes)
    end
end

moderangeaxes(m::ModeRange) = axes(m, 1)
moderangeaxes(m) = to_axis(m) # convert size to axes
moderangelength(m::ModeRange) = length(m)
moderangelength(m) = m
to_axis(n::Integer) = Base.OneTo(n)
to_axis(n) = n
_uncolon(ax, m::Colon) = ax
_uncolon(ax, m) = m

function SHArray(arr::AbstractArray{T,N}, modes::NTuple{N,Union{Colon, RangeOrInteger}}) where {T,N}
    modes_uncolon = map(_uncolon, axes(arr), modes)
    SHArray{T, N, typeof(arr), typeof(modes_uncolon)}(arr, modes_uncolon)
end
SHArray(A::AbstractArray, modes::Vararg{Union{Colon, RangeOrInteger}}) = SHArray(A, modes)

# SHArrays may pop their parent to avoid multiple wrappers
SHArray(S::SHArray{<:Any,N}, modes::NTuple{N,Union{Colon, RangeOrInteger}}) where {N} = SHArray(parent(S), modes)

SHArray{T,N,A,M}(S::SHArray{T,N,A,M}) where {T,N,A<:AbstractArray{T,N}, M<:NTuple{N,RangeOrInteger}} = S
SHArray{T,N,A,M}(S::SHArray) where {T,N,A<:AbstractArray{T,N}, M<:NTuple{N,RangeOrInteger}} = SHArray(convert(A, parent(S)), convert(M, modes(S)))
SHArray{T,N,A,M}(S::SHArray) where {T,N,A<:AbstractArray{T,N}, M<:NTuple{N,Union{AbstractUnitRange{<:Integer},Integer}}} = SHArray(convert(A, parent(S)), convert(M, map(moderangeaxes, modes(S))))
SHArray{T,N,A,M}(S::AbstractArray) where {T,N,A<:AbstractArray{T,N}, M<:NTuple{N,Union{AbstractUnitRange{<:Integer},Integer}}} = SHArray(convert(A, S), convert(M, axes(S)))
Base.convert(::Type{T}, A::AbstractArray) where {T<:SHArray} = A isa T ? A : T(A)

"""
    s = SHArray(arr::AbstractArray{T,N}, dimsmodes::Pairs{Int, <:SphericalHarmonicModes.ModeRange}) where {T,N}

Return an `SHArray` where the dimensions specified as the `keys` of `dimsmodes` may be indexed using `Tuple`s of
spherical harmonic degrees. This constructor is not type-stable in general, therefore this should be avoided within a loop
if performance is a concern.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> s2 = SHArray(reshape(1:4, 2, 2), 1 => LM(0:1, 0:0))
2×2 SHArray(reshape(::UnitRange{Int64}, 2, 2), (LM(0:1, 0:0), Base.OneTo(2))):
 1  3
 2  4
```
"""
function SHArray(arr::AbstractArray{<:Any,N}, modes::Vararg{Pair{Int,<:ModeRange}}) where {N}
    Base.depwarn("constructing an SHArray with pairs of ModeRanges is deprecated. Please use SHArray(arr, modes::Tuple) instead", :SHArray)
    all(x -> 1 <= x <= N, map(first, modes)) || throw(ArgumentError("all dimensions must be positive and <= $N"))
    allmodes = replaceaxeswithmodes(axes(arr), modes)
    SHArray(arr, allmodes)
end

function replaceaxeswithmodes(ax::NTuple{N,AbstractUnitRange}, modes) where {N}
    allmodes = Vector(undef, N)
    allmodes .= ax
    for (d,m) in modes
        allmodes[d] = m
    end
    ntuple(i -> allmodes[i], Val(N))
end
function replaceaxeswithmodes(ax::NTuple{N,AbstractUnitRange}, modes::NTuple{N, Pair{Int,<:ModeRange}}) where {N}
    allmodes = collect(modes)
    sort!(allmodes, by=first)
    ax = map(last, allmodes)
    ntuple(i -> ax[i], Val(N)) :: NTuple{N, ModeRange}
end
function replaceaxeswithmodes(ax::NTuple{N,AbstractUnitRange}, modes::NTuple{N, Pair{Int, T}}) where {N, T<:ModeRange}
    allmodes = TupleTools.sort(modes, by = first)
    map(last, allmodes)
end

# assume no ModeRange axes by default
SHArray(arr::AbstractArray) = SHArray(arr, axes(arr))

"""
    SHArray{T}(init, modes::NTuple{N,Union{Integer, AbstractUnitRange, SphericalHarmonicModes.ModeRange}}) where {T,N}

Return an `SHArray` wrapper around a parent array of the appropriate size with `N` dimensions and elements of type `T`,
such that the indices of `modes` that are of a `ModeRange` type may be indexed using
`Tuples` of spherical harmonic degrees.
The elements of the parent array are set according to the initializer `init`.

# Examples

```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> SHArray{Missing}(undef, (LM(0:1, 0:0), 2))
2×2 SHArray(::$(Array{Missing,2}), (LM(0:1, 0:0), Base.OneTo(2))):
 missing  missing
 missing  missing
```

    SHArray{T}(modes::NTuple{N,Union{Integer, AbstractUnitRange, SphericalHarmonicModes.ModeRange}}) where {T,N}

Return an `SHArray` of the appropriate size with `N` dimensions and elements of type `T`.
The elements are set to zero.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> SHArray{Float64}((LM(1:1, 0:1), LM(0:0, 0:0)))
2×1 SHArray(::$(Array{Float64,2}), (LM(1:1, 0:1), LM(0:0, 0:0))):
 0.0
 0.0

julia> SHArray{Float64}((LM(1:1, 0:1), 2))
2×2 SHArray(::$(Array{Float64,2}), (LM(1:1, 0:1), Base.OneTo(2))):
 0.0  0.0
 0.0  0.0
```
"""
function SHArray{T}(init::ArrayInitializer, modes::NTuple{N,RangeOrInteger}) where {T,N}
    SHArray{T,N}(init, modes)
end
function SHArray{T,N}(::UndefInitializer, modes::NTuple{N,RangeOrInteger}) where {T,N}
    arr = similar(Array{T,N}, map(moderangeaxes, modes))
    SHArray(arr, map(to_axis, modes))
end
function SHArray{T,N}(init::ArrayInitializer, modes::NTuple{N,RangeOrInteger}) where {T,N}
    S = SHArray{T,N}(undef, modes)
    fill!(parent(S), init)
    S
end

SHArray{T,N}(init::ArrayInitializer, modes::Vararg{RangeOrInteger,N}) where {T,N} = SHArray{T,N}(init, modes)
SHArray{T}(init::ArrayInitializer, modes::Vararg{RangeOrInteger,N}) where {T,N} = SHArray{T}(init, modes)

@deprecate SHArray{T}(modes::NTuple{N,RangeOrInteger}) where {T,N} fill!(SHArray{T}(undef, modes), zero(T))
@deprecate SHArray{T,N}(modes::NTuple{N,RangeOrInteger}) where {T,N} fill!(SHArray{T,N}(undef, modes), zero(T))

for (DT, f) in ((:zeros, :zero), (:ones, :one))
    @eval function Base.$DT(::Type{T}, modes::NTuple{N,RangeOrInteger}) where {T,N}
        s = SHArray{T,N}(undef, modes)
        fill!(parent(s), $f(eltype(parent(s))))
        s
    end
    @eval Base.$DT(::Type{T}, modes::Vararg{RangeOrInteger}) where {T} = $DT(T, modes)
    @eval Base.$DT(modes::Tuple{Vararg{RangeOrInteger}}) = $DT(Float64, modes)
    @eval Base.$DT(modes::Vararg{RangeOrInteger}) = $DT(Float64, modes)
end

# Convenience constructors
const SHArrayAllModeRange{T, N, AA<:AbstractArray{T,N}, TM<:NTuple{N,ModeRange}} = SHArray{T, N, AA, TM}
const SHVector{T, AA<:AbstractVector{T}, M<:Tuple{ModeRange}} = SHArrayAllModeRange{T, 1, AA, M}

"""
    SHVector(arr::AbstractVector, modes::SphericalHarmonicModes.ModeRange)

Return a wrapper that maps the indices of the parent `Vector` to a `Tuple` of spherical harmonic degrees.
Often the `Tuple` would be a pair of spherical harmonic degrees `(l,m)`.

# Examples

```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> s = SHVector(1:3, LM(1:1))
3-element SHArray(::UnitRange{Int64}, (LM(1:1, -1:1),)):
 1
 2
 3

julia> s.modes[1] |> collect
3-element $(Array{Tuple{Int64,Int64},1}):
 (1, -1)
 (1, 0)
 (1, 1)

julia> s[(1,0)]
2
```

See also: [`SHMatrix`](@ref), [`SHArray`](@ref)
"""
SHVector
@deprecate SHVector(A::AbstractVector, mode::ModeRange) SHArray(A, (mode,))
@deprecate SHVector(A::AbstractVector, modes::Tuple{ModeRange}) SHArray(A, modes)

"""
    SHVector{T}(init, modes::SphericalHarmonicModes.ModeRange) where {T}
    SHVector{T}(init, modes::Tuple{SphericalHarmonicModes.ModeRange}) where {T}

Return a `SHVector` wrapper, with the parent `Vector` having an element type of `T`.
The default value is set by the initializer `init`.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> SHVector{Missing}(undef, LM(1:1, -1:1))
3-element SHArray(::$(Array{Missing,1}), (LM(1:1, -1:1),)):
 missing
 missing
 missing
```

    SHVector{T}(modes::SphericalHarmonicModes.ModeRange) where {T}
    SHVector{T}(modes::Tuple{SphericalHarmonicModes.ModeRange}) where {T}

Return an `SHVector` of the appropriate size and with elements of type `T`.
The elements are set to zero.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> SHVector{ComplexF64}(LM(0:0, 0:0))
1-element SHArray(::$(Array{Complex{Float64},1}), (LM(0:0, 0:0),)):
 0.0 + 0.0im
```
"""
function SHVector{T}(init::ArrayInitializer, modes::Tuple{ModeRange}) where {T}
    Base.depwarn("SHVector{T}(init::$ArrayInitializer, modes::Tuple{ModeRange}) is deprecated, use SHArray{T}(init, modes) instead", :SHVectorTTup)
    SHVector{T}(init, first(modes))
end
function SHVector{T}(init::ArrayInitializer, modes::ModeRange) where {T}
    Base.depwarn("SHVector{T}(init::$ArrayInitializer, modes::ModeRange) is deprecated, use SHArray{T}(init, modes) instead", :SHVectorT)
    arr = Vector{T}(init, length(modes))
    SHVector(arr, (modes,))
end

function SHVector{T}(modes::Tuple{ModeRange}) where {T}
    Base.depwarn("SHVector{T}(modes::ModeRange) is deprecated, use zeros(T, modes) instead", :SHVectorTTupmodes)
    s = SHVector{T}(undef, modes)
    fill!(parent(s), zero(eltype(parent(s))))
    s
end
function SHVector{T}(modes::ModeRange) where {T}
    Base.depwarn("SHVector{T}(modes::ModeRange) is deprecated, use zeros(T, modes) instead", :SHVectorTmodes)
    SHVector{T}((modes,))
end

const SHMatrix{T, AA<:AbstractMatrix{T}, M<:NTuple{2,ModeRange}} = SHArrayAllModeRange{T,2,AA,M}

"""
    SHMatrix(arr::AbstractMatrix, modes::NTuple{2,SphericalHarmonicModes.ModeRange})
    SHMatrix(arr::AbstractMatrix, modes::Vararg{SphericalHarmonicModes.ModeRange,2})

Return a wrapper that maps the indices of the parent matrix along each axis
to a `Tuple` of spherical harmonic degrees.
Often these `Tuple`s would be pairs of spherical harmonic degrees `(l,m)`.

# Examples

```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> s = SHMatrix(reshape(1:4, 2,2), (LM(0:1, 0:0), LM(1:2,0:0)))
2×2 SHArray(reshape(::UnitRange{Int64}, 2, 2), (LM(0:1, 0:0), LM(1:2, 0:0))):
 1  3
 2  4

julia> s[(1,0),(2,0)]
4
```

See also: [`SHVector`](@ref), [`SHArray`](@ref)
"""
SHMatrix
@deprecate SHMatrix(arr::AbstractMatrix, modes::NTuple{2,ModeRange}) SHArray(arr, modes)
@deprecate SHMatrix(arr::AbstractMatrix, modes::Vararg{ModeRange,2}) SHArray(arr, modes)

"""
    SHMatrix{T}(init, modes::NTuple{2,SphericalHarmonicModes.ModeRange}) where {T}
    SHMatrix{T}(init, modes::Vararg{SphericalHarmonicModes.ModeRange, 2}) where {T}

Return a `SHMatrix` wrapper, with the parent `Matrix` having elements of type `T`.
The default value is set by the initializer `init`.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> SHMatrix{Missing}(undef, LM(1:1), LM(0:1, 0:0))
3×2 SHArray(::$(Array{Missing,2}), (LM(1:1, -1:1), LM(0:1, 0:0))):
 missing  missing
 missing  missing
 missing  missing
```

    SHMatrix{T}(modes::Vararg{SphericalHarmonicModes.ModeRange,2}) where {T}
    SHMatrix{T}(modes::NTuple{2,SphericalHarmonicModes.ModeRange}) where {T}

Return an `SHMatrix` of the appropriate size and with elements of type `T`.
The elements are set to zero.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> SHMatrix{ComplexF64}(LM(0:0, 0:0), LM(0:0, 0:0))
1×1 SHArray(::$(Array{Complex{Float64},2}), (LM(0:0, 0:0), LM(0:0, 0:0))):
 0.0 + 0.0im
```
"""
function SHMatrix{T}(init::ArrayInitializer, modes::NTuple{2,ModeRange}) where {T}
    Base.depwarn("SHMatrix{T}(init::$ArrayInitializer, modes::NTuple{2,ModeRange}) is deprecated, use SHArray{T}(init, modes) instead", :SHMatrixTTup)
    arr = Matrix{T}(init, map(length,modes))
    SHMatrix(arr, modes)
end
function SHMatrix{T}(init::ArrayInitializer, modes::Vararg{ModeRange,2}) where {T}
    Base.depwarn("SHMatrix{T}(init::$ArrayInitializer, modes::Vararg{ModeRange,2}) is deprecated, use SHArray{T}(init, modes) instead", :SHMatrixT)
    SHMatrix{T}(init, modes)
end

function SHMatrix{T}(modes::NTuple{2,ModeRange}) where {T}
    Base.depwarn("SHMatrix{T}(modes::NTuple{2,ModeRange}) is deprecated, use zeros(T, modes) instead", :SHMatrixTTupmodes)
    s = SHMatrix{T}(undef, modes)
    fill!(parent(s), zero(eltype(parent(s))))
    s
end
function SHMatrix{T}(modes::Vararg{ModeRange,2}) where {T}
    Base.depwarn("SHMatrix{T}(modes::Vararg{ModeRange,2}) is deprecated, use zeros(T, modes) instead", :SHMatrixTmodes)
    SHMatrix{T}(modes)
end

# Add methods to Base functions

Base.parent(s::SHArray) = s.parent

Base.similar(arr::SHArray) = SHArray(similar(parent(arr)), modes(arr))
Base.similar(arr::SHArray, T::Type) = SHArray(similar(parent(arr), T), modes(arr))
Base.similar(arr::AbstractArray, modes::RangeOrInteger...) = similar(arr, eltype(arr), modes)
Base.similar(arr::AbstractArray, T::Type, modes::RangeOrInteger...) = similar(arr, T, modes)
Base.similar(arr::AbstractArray, T::Type, modes::Tuple{Vararg{RangeOrInteger}}) = SHArray(similar(parent(arr), T, map(moderangelength, modes)), modes)

Base.dataids(A::SHArray) = Base.dataids(parent(A)) # needed for fast broadcasting
function Broadcast.broadcast_unalias(dest::SHArray, src::SHArray)
    parent(dest) === parent(src) ? src : Broadcast.unalias(dest, src)
end

# Accessor methods
modes(s::SHArray) = s.modes

"""
    SphericalHarmonicArrays.shmodes(arr::SHArray)

Returns a `Tuple` containing the elements of `arr.modes` that are of a `SphericalHarmonicModes.ModeRange` type.
This is type-stable even if `arr.modes` contains inhomogeneous types.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> sa = SHArray{ComplexF64}((2,2));

julia> SphericalHarmonicArrays.shmodes(sa)
()

julia> sa = SHArray{ComplexF64}((LM(1:1, 0:0),2));

julia> SphericalHarmonicArrays.shmodes(sa)
(LM(1:1, 0:0),)

julia> sa = SHArray{ComplexF64}((LM(1:1, 0:0), LM(2:2)));

julia> SphericalHarmonicArrays.shmodes(sa)
(LM(1:1, 0:0), LM(2:2, -2:2))
```

See also: [`shdims`](@ref)
"""
shmodes(b::SHArrayAllModeRange) = modes(b)
shmodes(b::SHArray) = unrolled_filter(x->isa(x, ModeRange), modes(b))

"""
    shdims(arr::SHArray)

Return the dimensions of `arr` that may be indexed using `ModeRange`s.
This is type-stable even if `arr.modes` contains inhomogeneous types.

# Examples
```jldoctest; setup=:(using SphericalHarmonicModes)
julia> using SphericalHarmonicModes: LM

julia> s = SHArray{Float64}((LM(1:1,0:1), LM(0:0)));

julia> SphericalHarmonicArrays.shdims(s)
(1, 2)

julia> s = SHArray{Float64}((LM(1:1,0:1), 2));

julia> SphericalHarmonicArrays.shdims(s)
(1,)
```

See also: [`shmodes`](@ref)
"""
shdims(b::SHArrayAllModeRange{<:Any,N}) where {N} = ntuple(identity, Val(N))
shdims(b::SHArray) = unrolled_filterdims(x->isa(x, ModeRange), modes(b))

Base.size(s::SHArray) = size(parent(s))
Base.axes(s::SHArray) = axes(parent(s))

# Indexing

Base.IndexStyle(::Type{SA}) where {SA<:SHArray} = IndexStyle(parenttype(SA))
parenttype(::Type{<:SHArray{<:Any,<:Any,AA}}) where {AA} = AA
parenttype(A::SHArray) = parenttype(typeof(A))

const ModeRangeIndexType = Union{Tuple{Integer,Integer},ModeRange,
                            Tuple{AbstractUnitRange{<:Integer},
                            AbstractUnitRange{<:Integer}}}

# use modes instead of axes for indexing in general, so that we may convert tuples of modes to indices
@propagate_inbounds Base.to_indices(s::SHArray, I::Tuple) = to_indices(s, modes(s), I)
@propagate_inbounds Base.to_indices(s::SHVector, I::Tuple{ModeRangeIndexType}) = to_indices(s, modes(s), I)

# methods to avoid ambiguity
@inline Base.to_indices(A::SHArray, I::Tuple{Vararg{Int}}) = I
@inline Base.to_indices(A::SHArray, I::Tuple{Vararg{Integer}}) = to_indices(parent(A), I)
@inline Base.to_indices(A::SHArray, ::Tuple{}) = ()
@inline Base.to_indices(s::SHArray, I::Tuple{Any}) = to_indices(parent(s), I)

@inline Base.to_indices(s::SHArray, I::Tuple{Vararg{Union{Integer, CartesianIndex}}}) =
    to_indices(parent(s), I)

@propagate_inbounds function Base.to_indices(s::SHArray, m::Tuple{ModeRange,Vararg{Any}},
    I::Tuple{ModeRangeIndexType,Vararg{Any}})
    (modeindex(first(m), first(I)), to_indices(s, tail(m), tail(I))...)
end

# getindex
@propagate_inbounds function Base.getindex(s::SHArray, I...)
    parent(s)[to_indices(s, I)...]
end

# setindex
@propagate_inbounds function Base.setindex!(s::SHArray, val, I...)
    parent(s)[to_indices(s, I)...] = val
    s
end

# views may pop the parent
Base.view(S::SHArray, I...) = view(parent(S), to_indices(S, I)...)

Base.strides(A::SHArray) = strides(parent(A))
Base.elsize(::Type{<:SHArray{T,N,A}}) where {T,N,A} = Base.elsize(A)
@inline Base.unsafe_convert(::Type{Ptr{T}}, A::SHArray{T}) where {T} = Base.unsafe_convert(Ptr{T}, parent(A))

# optimize deepcopy, but leads to a lot of invalidations
#function Base.deepcopy_internal(S::SHArray, stackdict::IdDict)
#    _P = Base.deepcopy_internal(parent(S), stackdict)::parenttype(S)
#    _modes = Base.deepcopy_internal(modes(S), stackdict)
#    typeof(S)(_P, _modes)
#end

# Broadcasting
Base.BroadcastStyle(::Type{<:SHArray}) = Broadcast.ArrayStyle{SHArray}()

function assert_modes_same(a,b)
    a == b || throw(ModeMismatchError(a,b))
end
@inline function assert_leading_modes_compatible(a::Tuple, b::Tuple)
    assert_modes_same(first(a) ,first(b))
    assert_leading_modes_compatible(tail(a), tail(b))
end
@inline assert_leading_modes_compatible(a::Tuple{},b::Tuple) = nothing
@inline assert_leading_modes_compatible(a::Tuple,b::Tuple{}) = nothing
@inline assert_leading_modes_compatible(a::Tuple{},b::Tuple{}) = nothing

function modes_larger(A::SHArray, B::SHArray)
    modes_A = modes(A)
    modes_out = modes_A
    (A === B) && return modes_out

    modes_B = modes(B)

    assert_leading_modes_compatible(modes_A, modes_B)

    if ndims(A) < ndims(B)
        modes_out = modes_B
    end

    return modes_out
end

function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{SHArray}},
    ::Type{ElType}) where ElType

    A = find_sharray(bc)
    modes_out = modes(A)

    # Check to make sure that all the SHArrays that are
    # being broadcasted over have compatible axes
    for B in Broadcast.flatten(bc).args
        if B isa SHArray
            (A === B) && continue
            modes_out = modes_larger(A,B)
        end
    end
    SHArray(similar(Array{ElType}, axes(bc)), modes_out)
end

find_sharray(bc::Base.Broadcast.Broadcasted) = find_sharray(bc.args)
find_sharray(args::Tuple) = find_sharray(find_sharray(args[1]), tail(args))
find_sharray(x) = x
find_sharray(a::SHArray, rest) = a
find_sharray(::Any, rest) = find_sharray(rest)

# show

function Base.showarg(io::IO, sa::SHArray, toplevel)
    print(io, "SHArray(")
    Base.showarg(io, parent(sa), false)
    print(io, ", ")
    print(io, modes(sa))
    print(io, ")")
end

function Base.replace_in_print_matrix(A::SHArray, i::Integer, j::Integer, s::AbstractString)
    Base.replace_in_print_matrix(parent(A), i, j, s)
end

"""
    permuteflipmodes(v::AbstractVector, p::Union{LM, ML})

Permute `v` with the map obtained by flipping the order of modes in `p`.
"""
function permuteflipmodes(v::AbstractVector, modes::Union{LM,ML})
    modes_new = SphericalHarmonicModes.flip(modes)
    v_new = SHArray(similar(v), modes_new)
    @inbounds for (ind, m) in enumerate(modes)
        v_new[m] = v[ind]
    end
    return v_new
end
permuteflipmodes(v::SHVector) = permuteflipmodes(v, first(modes(v)))

end # module


text:
# Adjoint test for F and J
# Author: Mathias Louboutin, mlouboutin3@gatech.edu
# Date: May 2020
#

using Distributed

# # Set parallel if specified
if nw > 1 && nworkers() < nw
    addprocs(nw-nworkers() + 1; exeflags=["--code-coverage=user", "--inline=no", "--check-bounds=yes"])
end

@everywhere using JUDI, LinearAlgebra, Test, Distributed

### Model
model, model0, dm = setup_model(tti, viscoacoustic, nlayer)
q, srcGeometry, recGeometry, f0 = setup_geom(model; nsrc=nw)
dt = srcGeometry.dt[1]

# testing parameters and utils
tol = 5f-4
(tti && fs) && (tol = 5f-3)
maxtry = 3

#################################################################################################
# adjoint test utility function so that can retry if fails

function run_adjoint(F, q, y, dm; test_F=true, test_J=true)
    adj_F, adj_J = !test_F, !test_J
    if test_F
        # Forward-adjoint
        d_hat = F*q
        q_hat = F'*y

        # Result F
        a = dot(y, d_hat)
        b = dot(q, q_hat)
        @printf(" <F x, y> : %2.5e, <x, F' y> : %2.5e, relative error : %2.5e \n", a, b, (a - b)/(a + b))
        adj_F = isapprox(a/(a+b), b/(a+b), atol=tol, rtol=0)
    end

    if test_J
        # Linearized modeling
        J = judiJacobian(F, q)
        ld_hat = J*dm
        dm_hat = J'*y

        c = dot(ld_hat, y)
        d = dot(dm_hat, dm)
        @printf(" <J x, y> : %2.5e, <x, J' y> : %2.5e, relative error : %2.5e \n", c, d, (c - d)/(c + d))
        adj_J = isapprox(c/(c+d), d/(c+d), atol=tol, rtol=0)
    end
    return adj_F, adj_J
end

test_adjoint(f::Bool, j::Bool, last::Bool) = (test_adjoint(f, last), test_adjoint(j, last))
test_adjoint(adj::Bool, last::Bool) = (adj || last) ? (@test adj) : (@test_skip adj)

###################################################################################################
# Modeling operators
@testset "Adjoint test with $(nlayer) layers and tti $(tti) and viscoacoustic $(viscoacoustic) and freesurface $(fs)" begin
    @timeit TIMEROUTPUT "Adjoint" begin
        opt = Options(sum_padding=true, dt_comp=dt, free_surface=parsed_args["fs"], f0=f0)
        F = judiModeling(model0, srcGeometry, recGeometry; options=opt)

        # Nonlinear modeling
        y = F*q

        # Run test until succeeds in case of bad case
        adj_F, adj_J = false, false
        ntry = 0
        while (!adj_F || !adj_J) && ntry < maxtry
            wave_rand = (.5f0 .+ rand(Float32, size(q.data[1]))).*q.data[1]
            q_rand = judiVector(srcGeometry, wave_rand)

            adj_F, adj_J = run_adjoint(F, q_rand, y, dm; test_F=!adj_F, test_J=!adj_J)
            ntry +=1
            test_adjoint(adj_F, adj_J, ntry==maxtry)
        end
        println("Adjoint test after $(ntry) tries, F: $(success_log[adj_F]), J: $(success_log[adj_J])")
    end
end
###################################################################################################
# Extended source modeling
@testset "Extended source adjoint test with $(nlayer) layers and tti $(tti) and viscoacoustic $(viscoacoustic) and freesurface $(fs)" begin
    @timeit TIMEROUTPUT "Extended source adjoint" begin
        opt = Options(sum_padding=true, dt_comp=dt, free_surface=parsed_args["fs"], f0=f0)
        F = judiModeling(model0, srcGeometry, recGeometry; options=opt)
        Pr = judiProjection(recGeometry)
        Fw = judiModeling(model0; options=opt)
        Pw = judiLRWF(srcGeometry.dt[1], circshift(q.data[1], 51))
        Fw = Pr*Fw*adjoint(Pw)

        # Extended source weights
        w = zeros(Float32, model0.n...)
        w[141:160, 65:84] .= randn(Float32, 20, 20)
        w = judiWeights(w; nsrc=nw)

        # Nonlinear modeling
        y = F*q

        # Run test until succeeds in case of bad case
        adj_F, adj_J = false, false
        ntry = 0
        while (!adj_F || !adj_J) && ntry < maxtry
            wave_rand = (.5f0 .+ rand(Float32, size(q.data[1]))).*q.data[1]
            q_rand = judiVector(srcGeometry, wave_rand)

            adj_F, adj_J = run_adjoint(F, q_rand, y, dm; test_F=!adj_F, test_J=!adj_J)
            ntry +=1
            test_adjoint(adj_F, adj_J, ntry==maxtry)
        end
        println("Adjoint test after $(ntry) tries, F: $(success_log[adj_F]), J: $(success_log[adj_J])")
    end
end


text:
@testset "2d weights" begin
    simplex = SVector{3}(SVector{2, Float64}[[1., -1], [1., 1], [-1.5, 0]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.3, 0.3, 0.4])
    @test isapprox(weights, projection_weights_reference(simplex))

    simplex = SVector{3}(SVector{2, Float64}[[1., -1], [1., 1], [1.5, 0]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.5, 0.5, 0.0])
    @test isapprox(weights, projection_weights_reference(simplex))

    simplex = SVector{3}(SVector{2, Float64}[[1., -1], [1., 1], [0.5, 0]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.0, 0.0, 1.0])
    @test isapprox(weights, projection_weights_reference(simplex))
end

@testset "3d weights" begin
    simplex = SVector{4}(SVector{3, Float64}[[1., -1, 0], [1., 1, 0], [0.4, 0, -0.1], [0.5, 0, 1]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.0, 0.0,0.942623, 0.057377], atol=1e-6)
    @test isapprox(weights, projection_weights_reference(simplex))

    simplex = SVector{4}(SVector{3, Float64}[[1., -1, 0], [1., 1, 0], [1.4, 0, -0.1], [1.5, 0, 1]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.5, 0.5, 0.0, 0.0])
    @test isapprox(weights, projection_weights_reference(simplex))
end

@testset "random simplex" begin
    Random.seed!(1)
    for i in 1:100
        simplex = SVector{4}([rand(SVector{3, Float64}) for i in 1:4])
        weights = projection_weights(simplex)
        @test all(weights .>= 0)
        @test isapprox(sum(weights), 1)
        @test isapprox(weights, projection_weights_reference(simplex))
    end
end

function test_no_penetration(simplex)
    weights = projection_weights(simplex)
    if all(x -> x > 0, weights)
        @show weights
        @test false
    end
end

@testset "numerical issues" begin
    # These simplices are degenerate in that all z-coordinates are the same.
    # Numerical issues due to floating point comparison can cause `projection_weights`
    # to return positive weights.
    let simplex = SVector{4, SVector{3, Float64}}([-3.33728783796428, 0.3321305518800686, 0.11004228580261734], [-3.33728783796428, -0.6678694481199314, 0.11004228580261734], [1.6627121620357201, -0.6678694481199314, 0.11004228580261734], [1.6627121620357201, 0.3321305518800686, 0.11004228580261734])
        test_no_penetration(simplex)
    end
    let simplex = SVector{4, SVector{3, Float64}}([-2.362029787803838, -0.4737843508044094, 0.11167470708499053], [2.637970212196162, 0.5262156491955906, 0.11167470708499053], [2.637970212196162, -0.4737843508044094, 0.11167470708499053], [-2.362029787803838, 0.5262156491955906, 0.11167470708499053])
        test_no_penetration(simplex)
    end
end


text:
# Date: July 1, 2018
# Author: Bernhard König

############################################################
# Title: Decision Tree Model for the French MTPL Data
############################################################

# If you do not yet have a suitable editor for Julia, you may want to consider VS Code.
# Consider this link
# https://github.com/JuliaEditorSupport/julia-vscode
# We note that there are various alternative editors for Julia.

##############################
# Load packages
##############################
@assert VERSION >= v"1.0" "Your Julia version does not satisfy the requirements for this package. You need to use Julia v1.x"
t0 = time_ns()
# cd(string(ENV["HOMEPATH"],"\\Documents\\ASync\\home\\Code\\Julia\\DecisionTrees.jl"))

using Distributed
import Pkg
# Distributed.@everywhere using Revise
Distributed.@everywhere import CSV
Distributed.@everywhere import DataFrames
Distributed.@everywhere import DataFrames: DataFrame

Distributed.@everywhere using DecisionTrees  # this may take some time

tela = (time_ns() - t0) / 1e9
@show tela # precompilation (first time usage) may take some time

##############################
# Read the data
##############################
datafile = joinpath("data", "freMTPL2", "freMTPL2.csv")
if !isdir(splitdir(datafile)[1])
    datafile = joinpath(Pkg.dir("DecisionTrees"), datafile)
end
@assert isfile(datafile) "You need to unzip the file freMTPL2.zip in the folder data/freMTPL2"
@time fullData = CSV.read(datafile, DataFrame,strict=true);

# the data is described in the R package CASdatasets
# also the Noll, Salzmann, Wüthrich Paper has some descriptive graphs of the data.

# We will slightly modify the data in the following
# add AreaInteger as new variable, i.e. A:F -> 1:6 (as suggested in the paper of Wüthrich, Noll, Salzmann)
areasSorted = sort(unique(fullData[:Area]))
AreaInteger = map(x->findall((in)([x]), areasSorted)[1], fullData[:Area])
fullData[:AreaInteger] = AreaInteger

# correct for unreasonable observations
for i = 1:size(fullData, 1)
    fullData[:ClaimNb][i] = min(4, fullData[:ClaimNb][i])
    fullData[:Exposure][i] = min(1.0, fullData[:Exposure][i])
end

# set independent variables
selected_explanatory_vars = ["Area","AreaInteger","VehPower","VehAge","DrivAge","BonusMalus","VehBrand","VehGas","Density","Region"]
# note: we refrain from using the exposure as explanatory variable, as this is not meaningful in practical applications
# where one aims to predict the claim frequency of a given policy (without knowing how long that policy will remain within a certain portfolio)

# check type of each column
for x in selected_explanatory_vars
    println(x)
    println(eltype(fullData[Symbol(x)]))
    println("first ten values")
    print(fullData[1:10,Symbol(x)])
    println("")
end

# consider unique elements for each variable
for x in 1:size(fullData, 2)
    println("Variable number ", x, ": ", names(fullData)[x])
    @show size(unique(fullData[:,x]))[1]
end

##############################
# Prepare the data
##############################

# dtmtable::DTMTable is a custom data format for DecisionTrees
# sett are the model settings
# dfprepped is an intermediary dataframe which is generally not needed
dtmtable, sett, dfprepped = prepare_dataframe_for_dtm!(fullData, keycol="IDpol", trnvalcol="trnTest", numcol="ClaimNb", denomcol="Exposure", weightcol="Exposure", independent_vars=selected_explanatory_vars);

#=
#NOTE:if the input data (i.e. the DataFrame fullData in this case) contains a row which is currently 'numerical' but should be treated as a categorical variable,
#then you can use the parameter treat_as_categorical_variable
#consider this code example, which is NOT MEANINGFUL in this context, but showcases the optional parameter
    dtmtableAlt,settAlt,dfpreppedAlt=prepare_dataframe_for_dtm!(fullData,keycol="IDpol",trnvalcol="trnTest",numcol="ClaimNb",denomcol="Exposure",weightcol="Exposure",independent_vars=selected_explanatory_vars,treat_as_categorical_variable=["DrivAge","VehAge"]);
    dtmtable.features[:DrivAge].pool #the data was interpreted as a numerical feature
    dtmtableAlt.features[:DrivAge].pool #the data was interpreted as a categorical feature (i.e. a String) =#

# Consider the following regarding the data structure of dtmtable
fieldnames(typeof(dtmtable))
dtmtable.key # an identifier (String type), ideally it is a unique identifier for each row
dtmtable.numerator # a vector
dtmtable.denominator # a vector
dtmtable.weight # a vector
# each of the above have the same length

dtmtable.trnidx # an index vector for the training data. This is an integer vector
# therefore we generally have
length(dtmtable.trnidx) < length(dtmtable.numerator)
# dtmtable.validx is the corresponding validation data
# we trnidx and validx should generally be a partition of 1:n (where n==length(dtmtable.weights))

dtmtable.features # explanatory variables
# note that the data is generally stored in a compressed format
# for numerical variables, the algorithm considers sett.maxSplittingPoints  (default 250) splitting points which are uniformly spaced
# you can increase this number and re-prepare the data if you want
#=
    sett.maxSplittingPoints=100
    dtmtable=prep_data_from_df(df_prepped,sett,fn) =#
# We realize that it may be sutiable to define the splitting in a different manner (than uniformly spaced).
# This feature might be added at a later time. You can consider the function add_coded_numdata!(...) to see how splitting points are chosen

# We note that 'pooling' the data into X (up to sett.maxSplittingPoints) buckets has a few advantages:
# one can do sorting in linear time (countsort)
# the data takes up less space in RAM and caches which generally increases performance

# keep a copy of original train and test split
originalTrnValIndex = deepcopy(fullData[:trnTest])

# Redefine trn and val data sets
# if you prefer train on X% of the data, you can use this function to adjust the training set. By default it is sampled randomly
# resample_trnvalidx!(dtmtable,.85)

##############################
# Define model SETTINGS
##############################

# the minimum weight of each individual tree shall be 3% of the exposure per leaf
# You can set minWeight to any positive value, e.g. sett.minWeight=20000
# if minWeight is set to a value in [-1,0] its absolute value is interpreted as a percentage
# thus with sett.minWeight=-0.03 we define the min weight of each leaf as 3% of the training data

# model_type: currently only "build_tree" and "boosted_tree" are supported
# we might add the implementation of bagging at a later stage

# calculatePoissonError, is a boolean which we need to enable here in order for output to show the poisson error of the estimates
updateSettingsMod!(sett,minWeight=-0.03,model_type="build_tree",calculatePoissonError=true)

##############################
# Run single tree model
##############################

# resM is the resulting Model
resultingFiles, resM = dtm(dtmtable, sett)
# consider the resulting Excel file for a description of the model
# it should be located here:
@show resultingFiles[1]

############################################################
# Investigate the predictions (fitted values)
############################################################

# For a single decision tree, the mean observed value for each leaf is one choice for the predicted value.

fitted, leafnumbers = predict(resM, dtmtable.features)

# by definition there are only N distinct fitted values, where N = 'number of leaves'
unique(fitted)
unique(leafnumbers)

# todo: add the poisson error calc here in julia
# notably, we are using the difference as splitting criterion, which has (to our knowledge) no explicit relationship
# to the poisson distribution which is often assumed for count (and frequency) data.

# the leaves are numbered in such a manner that the lowest fitted value corresponds to leaf number 1
# conversly, the highest fitted value corresponds to the highest leaf number.

############################################################
# Additional options
############################################################

# Changing the impurity function
# this is detailed further below

# Tree representation
# If you have Graphviz installed, you can generate a plot (on a PDF) of the tree
updateSettingsMod!(sett,write_dot_graph="true",graphvizexecutable="c:\\Program Files (x86)\\Graphviz2.38\\bin\\dot.exe")
resultingFiles, resM = dtm(dtmtable, sett)
# You will also get a *.dot.txt file with a text representation of the graph

# If you do not have graphviz installed, you can do the following:
    # you can create a tree representation like this
dot_graph = graph(resM) # this is the content of the .dot.txt file

    # copy tree graph to Clipboard
clipboard(dot_graph)
    # go to 'a graphviz online too':
    # https://dreampuf.github.io/GraphvizOnline/
    # paste the tree

# output location
# you can provide a file and folder name in the following manner:
resultingFiles, resM = dtm(dtmtable, sett, file="C:\\temp\\somename.txt")
# the extension is not relevant (but we suggest to provide it)
# you should ensure that the folder exists, otherwise the model will not run.

# you can set to disable much of the text in the julia window (i.e. the log in stdout which gets printed by the model)
sett.print_details = false
# notably thenames of the output files will still be shown.

# let us keep it at true for the following
sett.print_details = true

# WE NOTE THAT THERE ARE A NUMBER OF UNDOCUMENTED OPTIONS. MOST OF THESE ARE EXPERIMENTAL OR UNUSED.

# if you want to use the resutling tree in a different programming language you may want to
# consider these options which give you the core structure of the model in other languages
# updateSettingsMod!(sett,writeCsharpCode="true",writeVbaCode="true",writeSasCode="true")
# some of these options may only be applicable to boosting OR single tree models

############################################################
# Prepare a grid search over the parameter space
############################################################

# set some default settings
updateSettingsMod!(sett,
model_type="build_tree",
writeStatistics="true",
writeSasCode="false",
writeIterationMatrix="false",
writeResult="false",
writeCsharpCode="false",
writeVbaCode="false",
saveJLDFile="false",
saveResultAsJLDFile="false",

showProgressBar_time="true",
prroduceEstAndLeafMatrices="false",
write_dot_graph="false",

calculatePoissonError="true",
performanceMeasure="Average Poisson Error Val"
)

# performanceMeasure will determine which metric is considered for the output summary
# performanceMeasure must be in DecisionTrees.global_statsperiter_header
# you may want to consider the Excel file of any boosting to get an idea of the available options for performanceMeasure
# consider the first row of the sheet ModelStatistics of the Excel file which lists all measures
# We note that some missing (or constant) measures are either not fully implemented yet, or were not calculated (e.g. because a boolean in the settings was set to false)
@show sett.performanceMeasure
@assert in(sett.performanceMeasure, DecisionTrees.global_statsperiter_header)

# create a vector of ModelSettings
# note, we can provide an arbitrary number of vectors over which we want to loop
# the function createGridSearchSettings will copy the 'input' settings X times into a vector of settings
# X is the size of the cartesian product of all parameter vectors over which we want to loop.

# Let us consider different stopping criteria (minimum weight per leaf)
minweight_list = -collect(range(0.0001, stop=0.1, length=25))

settV = createGridSearchSettings(sett,
    minWeight=minweight_list
    );

# consider the length(settV) which is the number of models that will be run
length(settV)

# notably all these models will be constructed independently of each other#
# if we are only varying the minimum weight this could of course be improved (by pruning a granular tree)
# however, the functionality is much more flexible (see the boosting example). Therefore make no assmptions on the
# relationship between settV[1] and settV[2], etc.

# depending on the size of settV (and the size of the data set!), you may want to restart Julia with additional workers for parallelization, i.e.
# shell> julia -p 8
# to start julia with 8 workers. Depending on your CPU you may want to choose a different number

Sys.CPU_CORES # might be give you an indication of the number of workers() you could use

############################################################
# Run grid search
############################################################

tt0 = time_ns()
gridResult = dtm(dtmtable, settV, file="R:\\temp\\1\\MTPLsingleTree.CSV")
@show ela = (-tt0 + time_ns()) / 1e9
@info ".....done"

# The excelfile 'MTPLsingleTree_multistats.xlsx' will give you a summary of the statistics of each model run
# The validation error may not be evaluated for the most granular tree
# This is because the tree has found leaves with a fitted frequency of zero (on the training data)

# Surprisingly the validation error is smallest for the tree with minWeight=-0.0042625 (which is 0.42% or 1374 policies)

# We definitely want to challenge this on a different validation data set
# Let us therefore consdier a cross validated model

# this defines a CV sampler (which will be an additional parameter for the model run)
# with -10 as first argument we specify that there will be 10 disjoint samples
# the second argument must be 0.0 if we want to have a 10 fold partion of the data.
# the third argument is NOT relevant in this case
cvsampler = CVOptions(-10, 0.0, true)

# Let us consider an alternative CV sampler
# anotherCVsampler (below) will randomly draw 10 samples (without replacement)
# each sample will have have a size of 65% of the total size of dtmtable (train and validation!)
# thus, the individual samples will have common data (as 10*0.65>1.0)
anotherCVsampler = CVOptions(10, 0.65, true)

# the sampler below (yetAnotherCVsampler) is similar to anotherCVsampler
# by setting the last argument to false, none of the 10 samples will use the rows defined by validx
# as such you can still keep validx as a true hold out sample whic is not used by the CV models
yetAnotherCVsampler = CVOptions(10, 0.65, false)

# for the following models we will consdier cvsampler
# Note: for cross validation we can only provide a single 'setting' (and not the settings vector settV)
updateSettingsMod!(sett,minWeight=1700);
CVresult = dtm_single_threaded(dtmtable, sett, cvsampler)

# CV result is dataframe of the cross valiation.
# you can also consider the file dtmresult_multistats.xlsx
# it will be located in a temporary directory (see the julia log above the dataframe to find its location)

# from the result we can see that the model generated a degenerate poisson estimate of 0 a number of times
# let us try to increase the size of the validation data
# 7 fold cross valiation
CVresult = dtm_single_threaded(dtmtable, sett, CVOptions(-7, 0.0, true))

# the averaged error across the validation sets is around 0.314 which does not look unreasonable

# let us creat additional cross valiation models
# we create 20 models (which randomly sampled data) where the
CVresult = dtm_single_threaded(dtmtable, sett, CVOptions(20, 0.7, true))
# the error measure remains relatively stable across several CV samples
# for some samples it is even less than 0.31 on the validation data

# try another model
updateSettingsMod!(sett,minWeight=500);
CVresult = dtm_single_threaded(dtmtable, sett, cvsampler)
sett.minWeight = 500
originalTrnValIndex

############################################################
# Changing the impurity function
############################################################

# Until now we have used the 'difference' measure to select the best split
# Computationally the difference has the favorable property that we only need 'one pass' over the data
# That is to say, we can go once over the data and sum up the numerator and denominator for each bucket (in linear time)
# then, to evaluate each of the possible M splits, we only need to do M-1 comparisons for numerical variables (for categorical variables the case is more complex)

# For now the only possible impurity functions are
# poissondeviance and difference

# Let us consider the poisson error
# to evaluate the poisson error for a given split, we need the mean estimated frequency in the left and right node
# to obtain these means, we need a pass over the data
# this increase the computation time considerably compared to the difference measure

# We can set the poisson error as follows
updateSettingsMod!(sett,crit="poissondeviance")
resultingFiles, resM = dtm(dtmtable, sett)
# For comparison reasons
# on my machine the above model took 15.8 seconds (only the build tree operation)
# with the difference measure, it takes about 1.46 seconds

# Note: We have only marginally performace optimized this package.
# It is highly likely that a skilled programmer could improve various bits of the code in this package.

# Upon inspecting the resulting errors, we can see a comparable performace as with the difference measure.

# Let us consider a grid search for the Poisson Error impurity function.
minweight_list = -collect(range(0.0001, stop=0.1, length=25))

settV = createGridSearchSettings(sett,
    minWeight=minweight_list
    );

gridResult = dtm(dtmtable, settV, file="R:\\temp\\13\\MTPLsingleTree.CSV")

############################################################
# Partial Dependence Plots
############################################################
# this is currently in development and in need of review
# NOTE: the function is very slow. This will take a few minutes
# pds=partialDependence(resM,dtmtable,listOfFeatures=[:Density,:VehGas])
# pds["Density"]


text:
const _GROUPALIASES = Dict{QuOpName,Vector{QuOpName}}()

add_groupaliases(groupname,names) = (_GROUPALIASES[QuOpName(groupname)] = QuOpName.(collect(names)))

function unalias(A::BaseOperator)::BaseOperator
    if A.name ∈ keys(_GROUPALIASES)
        names = _GROUPALIASES[A.name]
        ind = first(A.inds)
        @assert isintindex(ind)
        BaseOperator(A.t,names[ind.num],Base.tail(A.inds))
    else
        A
    end
end

unalias(A::T) where T<:Union{ExpVal,Corr} = T(BaseOpProduct(unalias.(A.ops.v)))


text:
using Distributions, StatsBase
using MixtureModels

function test_emstep(::Type{Beta}, N=10000)
    α = rand()
    β = rand()
    samples = rand(Beta(α, β), N)
    weights = uweights(N)

    pred = emstep(Beta(), samples, weights, ones(N))

    isapprox(pred.α, α, atol=0.05) && isapprox(pred.β, β, atol=0.05)
end

function test_emstep(::Type{Normal}, N=10000)
    μ = randn()
    σ = exp(randn())
    samples = rand(Normal(μ, σ), N)
    weights = uweights(N)

    pred = emstep(Normal(), samples, weights, ones(N))

    isapprox(pred.μ, μ, atol=0.05) && isapprox(pred.σ, σ, rtol=0.05)
end

function test_emstep(::Type{Exponential}, N=1000)
    θ = exp(randn())
    samples = rand(Exponential(θ), N)
    weights = uweights(N)

    pred = emstep(Exponential(), samples, weights, ones(N))

    isapprox(pred.θ, θ, rtol=0.05)
end

function test_emstep(::Type{Gamma}, N=10000)
    α = exp(randn())
    θ = exp(randn())
    samples = rand(Gamma(α, θ), N)
    weights = uweights(N)

    pred = emstep(Gamma(), samples, weights, ones(N))

    isapprox(pred.α, α, atol=0.05) && isapprox(pred.θ, θ, rtol=0.05)
end

function test_emstep(::Type{Erlang}, α, N=1000)
    θ = exp(randn())
    samples = rand(Erlang(α, θ), N)
    weights = uweights(N)

    pred = emstep(Erlang(α), samples, weights, ones(N))

    pred.α == α && isapprox(pred.θ, θ, rtol=0.05)
end

N_rep = 100
p_success = 0.8

for T in [ Beta, Normal, Exponential, Gamma ]
    @test count(test_emstep(T) for i in 1:N_rep) >= p_success * N_rep
end

@test count(test_emstep(Erlang, 5) for i in 1:N_rep) >= p_success * N_rep
@test count(test_emstep(Erlang, 20) for i in 1:N_rep) >= p_success * N_rep


text:
sorted_keys = sort ∘ collect ∘ keys

function Base.redirect_stdout(f::Function, stream::IOBuffer)
    backup = stdout
    rd, wr = redirect_stdout()
    try
        f()
    finally
        redirect_stdout(backup)
        close(wr)
        write(stream, read(rd))
        close(rd)
    end
end

function Base.redirect_stderr(f::Function, stream::IOBuffer)
    backup = stderr
    rd, wr = redirect_stderr()
    try
        f()
    finally
        redirect_stderr(backup)
        close(wr)
        write(stream, read(rd))
        close(rd)
    end
end

function stringify(e::Exception)
    buf = IOBuffer()
    showerror(buf, e)
    String(take!(buf))
end


text:
@testset "Physical Constants" begin
    x = rand(keys(Fermi.PhysicalConstants.atom_num))
    @test Fermi.PhysicalConstants.atomic_number(x) == Fermi.PhysicalConstants.atom_num[x]
    @test_throws KeyError Fermi.PhysicalConstants.atomic_number("Invalid")

    @test Fermi.PhysicalConstants.num_core_electrons("Ca") == 18
end

max_stars_repo_name: FermiQC/Fermi

max_issues_repo_name: FermiQC/Fermi

max_forks_repo_name: FermiQC/Fermi


text:
function AESECB(plaintext, key::AbstractAESKey, cache::AbstractAESCache)
	len = length(plaintext)
	pad = 16 - (len % 16)
	ciphertext = similar(Array{UInt8, 1}, pad + len)
	AESECB!(ciphertext, plaintext, key, cache)
end

function AESECB!(ciphertext, plaintext, key::AbstractAESKey, cache::AbstractAESCache)
	len = length(plaintext)
	pad = 16 - (len % 16)
	for i in 1:len
		ciphertext[i] = UInt8(plaintext[i])
	end
	for i in 1:pad
		ciphertext[len+i] = pad
	end
	iters = Int((len + pad) / 16)
	ciphertextblock = AESBlock(ciphertext)
	for i in 1:iters
		AESEncryptBlock!(ciphertextblock, ciphertextblock, key, cache)
		increment!(ciphertextblock)
	end
	ciphertext
end

function AESECB_D(ciphertext::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache)
	len = length(ciphertext)
	iters = Int(len / 16)
	plaintext = similar(Array{UInt8, 1}, len)
	AESECB_D!(plaintext, ciphertext, key, cache)
end

function AESECB_D!(plaintext, ciphertext::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache)
	len = length(ciphertext)
	iters = Int(len / 16)
	ciphertextblock = AESBlock(ciphertext)
	plaintextblock = AESBlock(plaintext)
	for i in 1:iters
		AESDecryptBlock!(plaintextblock, ciphertextblock, key, cache)
		increment!(ciphertextblock)
		increment!(plaintextblock)
	end
	pad = plaintext[end]
	@view(plaintext[1:len-pad])
end


text:
#ENV["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
#using Revise
using ADSeismic
using ADCME
using PyPlot
using MAT
using Optim
using Random
using DelimitedFiles
using Dates
matplotlib.use("Agg")
close("all")
if has_gpu()
  gpu = true
else
  gpu = false
end

data_dir = "data/acoustic"
if !ispath(data_dir)
  mkpath(data_dir)
end
figure_dir = "figure/FWI/marmousi/"
if !ispath(figure_dir)
  mkpath(figure_dir)
end
result_dir = "result/FWI/marmousi/"
if !ispath(result_dir)
  mkpath(result_dir)
end
loss_file = joinpath(result_dir, "loss_$(Dates.now()).txt")
reset_default_graph()

################### Inversion using Automatic Differentiation #####################
model_name = "models/marmousi2-model-smooth.mat"
# model_name = "models/BP-model-smooth.mat"

## load model setting
params = load_params(model_name, vp_ref=3e3, PropagatorKernel=1)
src = load_acoustic_source(model_name)
rcv = load_acoustic_receiver(model_name)
vp = Variable(matread(model_name)["vp"])

## assemble acoustic propagator model
model = x->AcousticPropagatorSolver(params, x, vp)

## load data
std_noise = 0
Random.seed!(1234);
Rs = Array{Array{Float64,2}}(undef, length(src))
for i = 1:length(src)
    Rs[i] = readdlm(joinpath(data_dir, "marmousi-r$i.txt"))
    # Rs[i] = readdlm(joinpath(data_dir, "BP-r$i.txt"))
    Rs[i] .+= ( randn(size(Rs[i])) .+ mean(Rs[i]) ) .* std(Rs[i]) .* std_noise
end

## calculate loss
if gpu
  losses, grads = compute_loss_and_grads_GPU(model, src, rcv, Rs, vp)
  grad = sum(grads)
  loss = sum(losses)
else
  [SimulatedObservation!(model(src[i]), rcv[i]) for i = 1:length(src)]
  loss = sum([sum((rcv[i].rcvv-Rs[i])^2) for i = 1:length(rcv)])
end

global_step = tf.Variable(0, trainable=false)
max_iter = 50000
lr_decayed = tf.train.cosine_decay(1.0, global_step, max_iter)
opt = AdamOptimizer(lr_decayed).minimize(loss, global_step=global_step, colocate_gradients_with_ops=true)

sess = Session(); init(sess)
loss0 = run(sess, loss)
if gpu
g = run(sess, grad)[1]
else
  grad = gradients(loss, vp)
  g = run(sess, grad)[1]
end
@info "Initial loss: ", loss0
# error()

## run inversion
fp = open(loss_file, "w")
write(fp, "0,$loss0\n")
function callback(vs, iter, loss)
  if iter%10==0
    # if gpu
      x = Array(reshape(vs[1:(params.NY+2)*(params.NX+2)], params.NY+2, params.NX+2))
    # else
    #   x = vs[1]'
    # end
    clf()
    pcolormesh([0:params.NX+1;]*params.DELTAX/1e3,[0:params.NY+1;]*params.DELTAY/1e3,  x)
    axis("scaled")
    colorbar(shrink=0.4)
    xlabel("x (km)")
    ylabel("z (km)")
    gca().invert_yaxis()
    title("Iteration = $iter")
    savefig(joinpath(figure_dir, "inv_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
    writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), x)
  end
  write(fp, "$iter,$loss\n")
  flush(fp)
end

## BFGS
Optimize!(sess, loss, 10000, vars=[vp], grads=grad,  callback=callback)

## ADAM
# time = 0
# for iter = 0:max_iter
#   global time += @elapsed  _, ls, lr = run(sess, [opt, loss, lr_decayed])
#   callback(run(sess, vp)', iter, ls)
#   println("   $iter\t$ls\t$lr")
#   println(" * time: $time")
# end

close(fp)


text:
"""
This is a Julia translation of the schema expected by typescript.

Calling [`to_json`](@ref) on a SemagramSchema should give a schema that
typescript can understand.

There is also a macro here for creating a SemagramSchema: [`@semagramschema`](@ref),
which makes writing down a schema much more convenient.
"""
module Schema

export AttributeType, Numeric, Stringlike,
  BoxProperties, PortProperties, WireProperties,
  SemagramSchema, BoxDesc, PortDesc, WireDesc, DataDesc,
  @semagramschema

import ..Muesli: to_json, from_json
using ..Muesli
using ..SVG
using Catlab.Present, Catlab.Theories
using MLStyle

@enum AttributeType begin
  Numeric
  Stringlike
end

to_json(x::AttributeType) = @match x begin
  Numeric => "Numeric"
  Stringlike => "Stringlike"
end

function from_json(d::Any, ::Type{AttributeType})
  @match d begin
    "Numeric" => Numeric
    "Stringlike" => Stringlike
  end
end

struct BoxProperties
  weights::Vector{Tuple{AttributeType, Symbol}}
  shape::String
  label::Union{Symbol,Nothing}
end

to_json(x::BoxProperties) = generic_to_json(x)
from_json(d::Dict{String,<:Any}, ::Type{BoxProperties}) = generic_from_json(d, BoxProperties)

struct PortProperties
  weights::Vector{Tuple{AttributeType, Symbol}}
  box::Symbol
  box_map::Symbol
  style::String
end

to_json(x::PortProperties) = generic_to_json(x)
from_json(d::Dict{String,<:Any}, ::Type{PortProperties}) = generic_from_json(d, PortProperties)

struct WireProperties
  weights::Vector{Tuple{AttributeType, Symbol}}
  src::Tuple{String, Symbol}
  src_map::Symbol
  tgt::Tuple{String, Symbol}
  tgt_map::Symbol
  style::String
end

to_json(x::WireProperties) = generic_to_json(x)
from_json(d::Dict{String,<:Any}, ::Type{WireProperties}) = generic_from_json(d, WireProperties)

struct SemagramSchema
  box_types::Dict{Symbol, BoxProperties}
  port_types::Dict{Symbol, PortProperties}
  wire_types::Dict{Symbol, WireProperties}
end

to_json(x::SemagramSchema) = generic_to_json(x)
from_json(d::Dict{String, <:Any}, ::Type{SemagramSchema}) = generic_from_json(d, SemagramSchema)

struct BoxDesc
  name::Symbol
  shape::SVGNode
  label::Union{Symbol,Nothing}
  function BoxDesc(name,shape=Circle,label=nothing)
    new(name,shape, label)
  end
end

struct PortDesc
  name::Symbol
  box_map::Symbol
  style::String
  function PortDesc(name,box_map,style="Circular")
    new(name,box_map,style)
  end
end

struct WireDesc
  name::Symbol
  src_map::Symbol
  tgt_map::Symbol
  style::String
  function WireDesc(name,src_map,tgt_map,style="DefaultWire")
    new(name,src_map,tgt_map,style)
  end
end

struct DataDesc
  name::Symbol
  type::AttributeType
end

const EntityDesc = Union{BoxDesc, PortDesc, WireDesc, DataDesc}

function pres_to_semagramschema(p::Presentation, descs::Array)
  ws = SemagramSchema(Dict(),Dict(),Dict())
  datadescs = filter(desc -> typeof(desc) <: DataDesc, descs)
  datatypes = Dict(map(desc -> desc.name => desc.type, datadescs)...)
  function weights(ob::Symbol)
    attrs = filter(attr -> nameof(dom(attr)) == ob, p.generators[:Attr])
    map(attr -> (get(datatypes, nameof(codom(attr)), Stringlike), nameof(attr)), attrs)
  end
  for desc in descs
    ob = desc.name
    if typeof(desc)<:BoxDesc
      ws.box_types[ob] = BoxProperties(weights(ob),write_svg(desc.shape),desc.label)
    elseif typeof(desc)<:PortDesc
      box = nameof(codom(p[desc.box_map]))
      ws.port_types[ob] = PortProperties(weights(ob),box,desc.box_map,desc.style)
    elseif typeof(desc)<:WireDesc
      src = nameof(codom(p[desc.src_map]))
      src_type = if src ∈ keys(ws.box_types)
        "Box"
      else
        "Port"
      end
      tgt = nameof(codom(p[desc.tgt_map]))
      tgt_type = if src ∈ keys(ws.box_types)
        "Box"
      else
        "Port"
      end
      ws.wire_types[ob] = WireProperties(
        weights(ob),
        (src_type, src), desc.src_map,
        (tgt_type, tgt), desc.tgt_map,
        desc.style
      )
    end
  end
  ws
end

q(x) = Expr(:quote,x)

"""
See the examples/general documentation for how to use this.
"""
macro semagramschema(head,body)
  descs = @match body begin
    Expr(:block,lines...) => begin
      map(lines) do line
        @match line begin
          Expr(:macrocall, mname, _, desc, args...) => begin
            (name,morphs) = @match desc begin
              name::Symbol => (name,[])
              Expr(:call, name, args...) => (name,args)
            end
            constructor = if mname == Symbol("@box")
              BoxDesc
            elseif mname == Symbol("@port")
              PortDesc
            elseif mname == Symbol("@wire")
              WireDesc
            elseif mname == Symbol("@data")
              DataDesc
            end
            Expr(:call, constructor, q(name), q.(morphs)..., args...)
          end
          _ => missing
        end
      end
    end
    _ => error("the body must be a block")
  end
  descs = filter(x -> !ismissing(x), descs)
  name, pres = @match head begin
    Expr(:call, name, pres) => (name, pres)
    _ => error("the head must be a name and presentation")
  end
  :($(esc(name)) = pres_to_semagramschema($(esc(pres)), $(esc(Expr(:vect, descs...)))))
end

end


text:
cd(@__DIR__)
using Pkg
Pkg.activate(@__DIR__); Pkg.instantiate()

using OrdinaryDiffEq, Test

include("prinz_neuron_setup.jl")

byhand_prob = ODEProblem{true}(prinz_neuron!, u0, (0.,2000.), p)
byhand_sol = solve(byhand_prob, Rosenbrock23(), reltol=1e-8, abstol=1e-8, saveat=0.025);
current_mtk_sol = solve(sim, Rosenbrock23(), reltol=1e-8, abstol=1e-8, saveat=0.025);

tsteps = 0.0:0.025:2000.0
byhand_out = Array(byhand_sol(tsteps, idxs=2))
current_mtk_out = current_mtk_sol[neuron.sys.Vₘ]

@test isapprox(byhand_out, current_mtk_out, rtol=0.001)


text:
using BrewingRecipes
using Test

@testset "BrewingRecipes.jl" begin
    # Write your tests here.
end


text:
include("./Functions.jl")
using Functions

function factorial(n)
  total = 1
  for i = 1:n
    total = total * i
  end
  total
end

println(digitSum(factorial(BigInt(100))))

max_stars_repo_name: vmchale/julia-problem-solving

max_issues_repo_name: vmchale/julia-problem-solving

max_forks_repo_name: vmchale/julia-problem-solving


text:
# Exemplo de aquisição de dados utilizando PyCall

using PyCall

# incluindo caminho da pasta para o pyimport
push!(pyimport("sys")["path"], pwd());
# ou
push!(pyimport("sys")["path"], "/home/Documentos/Doutorado/scanivalve-master");

# Abrindo conexão e configurando
scani = pyimport("scanivalve")
s = scani.Scanivalve(ip="191.30.80.130")
s.config(AVG=1, PERIOD=1000, FPS=1000) # deve dar 16 segundos de medida

# Lendo dados
p,f = s.acquire();


text:
# p-value for one-sided permutation test on linear regression slope
function pvalue_linreg(m0, x, y; num::Int = Int(1e6))
    p = 0
    for i in 1:num
        m, c, dm, dc = linreg(x=shuffle(x), y=y)
        if m > m0
            p += 1
        end
    end

    return p/num
end

# p-value for one-sided permutation test on Pearson's r
function pvalue_Pearson(rho0, x, y; num::Int = Int(1e6))
    p = 0
    for i in 1:num
        rho = cor(shuffle(x),y)
        if rho > rho0
            p += 1
        end
    end

    return p/num
end


text:
struct SignedDistanceRefinery{F <: Function} <: AbstractRefinery
    signed_distance_func::F
    atol::Float64
    rtol::Float64
end

function needs_refinement(refinery::SignedDistanceRefinery, cell::Cell)
    minimum(cell.boundary.widths) > refinery.atol || return false

    for c in body_and_face_centers(cell.boundary)
        value_interp = evaluate(cell, c)
        value_true = refinery.signed_distance_func(c)
        if !isapprox(value_interp, value_true, rtol=refinery.rtol, atol=refinery.atol)
            return true
        end
    end
    return false
end

function refine_data(refinery::SignedDistanceRefinery, cell::Cell, indices)
    refine_data(refinery, child_boundary(cell, indices))
end

function refine_data(refinery::SignedDistanceRefinery, boundary::HyperRectangle)
    extrapolate(interpolate!(refinery.signed_distance_func.(vertices(boundary)),
                             BSpline(Linear()),
                             OnGrid()), Linear())
end


text:
module VulkanExperiments

module DrawingATriangle
    include("./DrawingATriangle/0_Setup.jl")
    include("./DrawingATriangle/1_Instance.jl")
    include("./DrawingATriangle/2_ValidationLayers.jl")
end # DrawingATriangle

end # VulkanExperiments


text:
module CUBLAS

using CUDAapi

using CUDAdrv
using CUDAdrv: CUstream

using CUDAnative

using ..CuArrays
using ..CuArrays: unsafe_free!
using LinearAlgebra

using CEnum

const libcublas = Ref("libcublas")

# core library
include("libcublas_common.jl")
include("error.jl")
include("libcublas.jl")

# low-level wrappers
include("util.jl")
include("wrappers.jl")

# high-level integrations
include("linalg.jl")

const created_handles = IdDict{CuContext,cublasHandle_t}()
const created_xt_handles = IdDict{CuContext,cublasXtHandle_t}()
const active_handles = Vector{Union{Nothing,cublasHandle_t}}()
const active_xt_handles = Vector{Union{Nothing,cublasXtHandle_t}}()

function handle()
    tid = Threads.threadid()
    if @inbounds active_handles[tid] === nothing
        ctx = context()
        active_handles[tid] = get!(created_handles, ctx) do
            handle = cublasCreate_v2()
            atexit(()->CUDAdrv.isvalid(ctx) && cublasDestroy_v2(handle))

            # enable tensor math mode if our device supports it, and fast math is enabled
            dev = CUDAdrv.device()
            if Base.JLOptions().fast_math == 1 && CUDAdrv.capability(dev) >= v"7.0" && version() >= v"9"
                cublasSetMathMode(CUBLAS_TENSOR_OP_MATH, handle)
            end

            handle
        end
    end
    @inbounds active_handles[tid]
end

function xt_handle()
    tid = Threads.threadid()
    if @inbounds active_xt_handles[tid] === nothing
        ctx = context()
        active_xt_handles[tid] = get!(created_xt_handles, ctx) do
            handle = cublasXtCreate()
            atexit(()->CUDAdrv.isvalid(ctx) && cublasXtDestroy(handle))

            # select the devices
            # TODO: this is weird, since we typically use a single device per thread/context
            devs = convert.(Cint, CUDAdrv.devices())
            cublasXtDeviceSelect(handle, length(devs), devs)

            handle
        end
    end
    @inbounds active_xt_handles[tid]
end

function __init__()
    resize!(active_handles, Threads.nthreads())
    fill!(active_handles, nothing)

    resize!(active_xt_handles, Threads.nthreads())
    fill!(active_xt_handles, nothing)

    CUDAnative.atcontextswitch() do tid, ctx
        # we don't eagerly initialize handles, but do so lazily when requested
        active_handles[tid] = nothing
        active_xt_handles[tid] = nothing
    end
end

end


text:
"""
    (1) swap_observed(model::AbstractSemSingle; kwargs...)

    (2) swap_observed(model::AbstractSemSingle, observed; kwargs...)

Return a new model with swaped observed part.

# Arguments
- `model::AbstractSemSingle`: optimization algorithm.
- `kwargs`: additional keyword arguments; typically includes `data = ...`
- `observed`: Either an object of subtype of `SemObserved` or a subtype of `SemObserved`

# Examples
See the online documentation on [Swap observed data](@ref).
"""
function swap_observed end

"""
    update_observed(to_update, observed::SemObserved; kwargs...)

Update a `SemImply`, `SemLossFunction` or `SemOptimizer` object to use a `SemObserved` object.

# Examples
See the online documentation on [Swap observed data](@ref).

# Implementation
You can provide a method for this function when defining a new type, for more information
on this see the online developer documentation on [Update observed data](@ref).
"""
function update_observed end

############################################################################################
# change observed (data) without reconstructing the whole model
############################################################################################

# use the same observed type as before
swap_observed(model::AbstractSemSingle; kwargs...) =
    swap_observed(model, typeof(observed(model)).name.wrapper; kwargs...)

# construct a new observed type
swap_observed(model::AbstractSemSingle, observed_type; kwargs...) =
    swap_observed(model, observed_type(;kwargs...); kwargs...)

swap_observed(model::AbstractSemSingle, new_observed::SemObserved; kwargs...) =
    swap_observed(model, observed(model), imply(model), loss(model), optimizer(model), new_observed; kwargs...)

function swap_observed(
        model::AbstractSemSingle,
        old_observed,
        imply,
        loss,
        optimizer,
        new_observed::SemObserved;
        kwargs...)

    kwargs = Dict{Symbol, Any}(kwargs...)

    # get field types
    kwargs[:observed_type] = typeof(new_observed)
    kwargs[:old_observed_type] = typeof(old_observed)
    kwargs[:imply_type] = typeof(imply)
    kwargs[:loss_types] = [typeof(lossfun) for lossfun in loss.functions]
    kwargs[:optimizer_type] = typeof(optimizer)

    # update imply
    imply = update_observed(imply, new_observed; kwargs...)
    kwargs[:imply] = imply
    kwargs[:n_par] = n_par(imply)

    # update loss
    loss = update_observed(loss, new_observed; kwargs...)
    kwargs[:loss] = loss

    # update optimizer
    optimizer = update_observed(optimizer, new_observed; kwargs...)

    #new_imply = update_observed(model.imply, new_observed; kwargs...)

    return Sem(
        new_observed,
        update_observed(model.imply, new_observed; kwargs...),
        update_observed(model.loss, new_observed; kwargs...),
        update_observed(model.optimizer, new_observed; kwargs...)
        )
end

function update_observed(loss::SemLoss, new_observed; kwargs...)
    new_functions = Tuple(update_observed(lossfun, new_observed; kwargs...) for lossfun in loss.functions)
    return SemLoss(
        new_functions,
        loss.weights)
end


text:
module IterativeMethod

using LinearAlgebra

export
    is_diagonally_dominant,
    jacobi_method,
    gauss_seidel_method, gauss_seidel_method!

#===== Jacobi method =====#

#=

A x = b
A = D + R
D x = -R x + b
  x = -D^{-1} R x + D^{-1} b
⟹
x = B x + z, where B = -D^{-1} R, z = D^{-1}

x_{i + 1} = B x_i + z
          = B (x^* + e_i) + z
          = x^* + B e_i
⟹
e_{i + 1} = B e_i

=#

function is_diagonally_dominant(A::AbstractMatrix{T})::Bool where T<:AbstractFloat
    A = abs.(A)
    d2 = diag(A) * 2
    all(sum(A, dims=1)' .<= d2) && all(sum(A, dims=2) .<= d2)
end

const MAX_ITERATIONS = 1000

function jacobi_method(
    A::Matrix{T},
    b::Vector{T},
    x0::Vector{T} = zero(b),
)::Vector{T} where T<:AbstractFloat
    (nr, nc) = size(A)
    @assert nr == nc "Jacobi: A is not square."
    if !is_diagonally_dominant(A)
        println("Jacobi: A is not diagonally dominant")
    end

    D = Diagonal(A)
    invD = inv(D)
    negR = D - A
    x = x0
    for _ in 1:MAX_ITERATIONS
        x = invD * (negR * x0 + b)
        if x ≈ x0
            return x
        end
        x0 = x
    end
    x
end

#===== Gauss–Seidel method =====#

#=

In Jacobi method,

    x_i^{(k + 1)} = -1/a_ii * ((L + U)_{i, :} * x^{(k)}) + b_i / a_ii

where $U$ ($L$) is the strictly upper (lower) triangular and $R = L + U$. If in the iteration we
use the elements of $x^{(k+1)}$ that have already been computed to compute the unknown ones of
$x^{(k+1)}$,

    x_i^{(k + 1)} = -1/a_ii * ((L_*)_{i, :} * x^{(k + 1)} + U_{i, :} * x^{(k)}) + b_i / a_ii

where $L_* = D + L$ is a lower triangular matrix. Or more compcatly,

    A x = b
    (L_* + U) x = b
    L_* x = b - U x

    ⟹

    L_* x^{(k + 1)} = b - U x^{(k)}
    x^{(k + 1)} = L_*^{-1} (b - U x^{(k)})

=#

function gauss_seidel_method(A::Matrix{T}, b::Vector{T})::Vector{T} where T<:AbstractFloat
    x0 = zero(b)
    gauss_seidel_method!(A, b, x0)
end

# The initial guess `x0` will be updated in-place.
function gauss_seidel_method!(
    A::Matrix{T},
    b::Vector{T},
    x0::Vector{T},
)::Vector{T} where T<:AbstractFloat
    (nr, nc) = size(A)
    @assert nr == nc "Gauss–Seidel: A is not square."
    n = nr
    if !is_diagonally_dominant(A)
        println("Gauss–Seidel: A is not diagonally dominant")
    end

    x = x0
    for _ in 1:MAX_ITERATIONS
        updated = false
        for r in 1:n
            old = x[r]
            x[r] = (b[r] - ((@view A[r, :])' * x - A[r, r] * x[r])) / A[r, r]
            if !(x[r] ≈ old)
                updated = true
            end
        end
        if !updated
            return x
        end
    end
    x
end

end # module

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
"""
ULMFiT - Text Classifier

This is wrapper around the LanguageMode struct. It has three fields:

vocab           : contains the same vocabulary from the LanguageModel
rnn_layers      : contains same DroppedEmebeddings, LSTM (AWD_LSTM) and VarDrop layers of LanguageModel except for last softmax layer
linear_layers   : contains Chain of two Dense layers [PooledDense and Dense] with softmax layer

To train create and instance and give it as first argument to 'train_classifier!' function
"""
mutable struct TextClassifier
    vocab::Vector
    rnn_layers::Flux.Chain
    linear_layers::Flux.Chain
end

function TextClassifier(lm::LanguageModel=LanguageModel(), clsfr_out_sz::Integer=1, clsfr_hidden_sz::Integer=50, clsfr_hidden_drop::Float64=0.4)
    return TextClassifier(
        lm.vocab,
        lm.layers[1:8],
        Chain(
            gpu(PooledDense(length(lm.layers[7].layer.cell.h), clsfr_hidden_sz)),
            gpu(BatchNorm(clsfr_hidden_sz, relu)),
            Dropout(clsfr_hidden_drop),
            gpu(Dense(clsfr_hidden_sz, clsfr_out_sz)),
            gpu(BatchNorm(clsfr_out_sz)),
            softmax
        )
    )
end

Flux.@functor TextClassifier

"""
Cross Validate

This function will be used to cross-validate the classifier

Arguments:

tc              : Instance of TextClassfier
gen             : 'Channel' to get a mini-batch from validation set
num_of_batches  : specifies the number of batches the validation will be done

If num_of_batches is not specified then all the batches which can be given by the
gen will be used for validation
"""
function validate(tc::TextClassifier, gen::Channel, num_of_batches::Union{Colon, Integer})
    n_classes = size(tc.linear_layers[end-2].W, 1)
    classifier = tc
    Flux.testmode!(classifier)
    loss = 0
    iters = take!(gen)
    ((num_of_batches != :) & (num_of_batches < iters)) && (iters = num_of_batches)
    TP, TN = gpu(zeros(Float32, n_classes, 1)), gpu(zeros(Float32, n_classes, 1))
    FP, FN = gpu(zeros(Float32, n_classes, 1)), gpu(zeros(Float32, n_classes, 1))
    for i=1:num_of_batches
        X = take!(gen)
        Y = gpu(take!(gen))
        X = map(x -> indices(x, classifier.vocab, "_unk_"), X)
        H = classifier.rnn_layers.(X)
        H = classifier.linear_layers(H)
        l = crossentropy(H, Y)
        Flux.reset!(classifier.rnn_layers)
        TP .+= sum(H .* Y, dims=2)
        FN .+= sum(((-1 .* H) .+ 1) .* Y, dims=2)
        FP .+= sum(H .* ((-1 .* Y) .+ 1), dims=2)
        TN .+= sum(((-1 .* H) .+ 1) .* ((-1 .* Y) .+ 1), dims=2)
        loss += l
    end
    precisions = TP ./ (TP .+ FP)
    recalls = TP ./ (TP .+ FN)
    F1 = (2 .* (precisions .* recalls)) ./ (precisions .+ recalls)
    accuracy = (TP[1] + TN[1])/(TP[1] + TN[1] + FP[1] + FN[1])
    return (loss, accuracy, precisions, recalls, F1)
end

"""
Forward pass

This funciton does the main computation of a mini-batch.
It computes the output of the all the layers [RNN and DENSE layers] and returns the predicted output for that pass.
It uses Truncated Backprop through time to compute the output.

Arguments:
tc              : Instance of TextClassifier
gen             : data loader, which will give 'X' of the mini-batch in one call
tracked_steps   : This is the number of tracked time-steps for Truncated Backprop thorugh time,
                  these will be last time-steps for which gradients will be calculated.
"""
function forward(tc::TextClassifier, gen::Channel, tracked_steps::Integer=32)
  	# swiching off tracking
    classifier = tc
    X = take!(gen)
    l = length(X)
    # Truncated Backprop through time
    Zygote.ignore() do
	for i=1:ceil(l/tracked_steps)-1   # Tracking is swiched off inside this loop
	    (i == 1 && l%tracked_steps != 0) ? (last_idx = l%tracked_steps) : (last_idx = tracked_steps)
	    H = broadcast(x -> indices(x, classifier.vocab, "_unk_"), X[1:last_idx])
	    H = classifier.rnn_layers.(H)
	    X = X[last_idx+1:end]
	end
    end
    # set the lated hidden states to original model
    for (t_layer, unt_layer) in zip(tc.rnn_layers[2:end], classifier.rnn_layers[2:end])
        if t_layer isa AWD_LSTM
            t_layer.layer.state = unt_layer.layer.state
            continue
        end
        if !unt_layer.reset
            t_layer.mask = unt_layer.mask
            t_layer.reset = false
        end
    end
    # last part of the sequecnes in X - Tracking is swiched on
    H = broadcast(x -> tc.rnn_layers[1](indices(x, classifier.vocab, "_unk_")), X)
    H = tc.rnn_layers[2:end].(H)
    H = tc.linear_layers(H)
    return H
end

"""
    loss(classifier::TextClassifier, gen::Channel, tracked_steps::Integer=32)

LOSS function

It takes the output of the forward funciton and returns crossentropy loss.

Arguments:

classifier    : Instance of TextClassifier
gen           : 'Channel' [data loader], to give a mini-batch
tracked_steps : specifies the number of time-steps for which tracking is on
"""
function loss(classifier::TextClassifier, gen::Channel, tracked_steps::Integer=32)
    H = forward(classifier, gen, tracked_steps)
    Y = gpu(take!(gen))
    l = crossentropy(H, Y)
    Flux.reset!(classifier.rnn_layers)
    return l
end

function discriminative_step!(layers, classifier::TextClassifier, gen::Channel, tracked_steps::Integer, ηL::Float64, opts::Vector)
    @assert length(opts) == length(layers)
    # Gradient calculation
    grads = Zygote.gradient(() -> loss(classifier, gen, tracked_steps = tracked_steps), get_trainable_params(layers))

    # discriminative step
    ηl = ηL/(2.6^(length(layers)-1))
    for (layer, opt) in zip(layers, opts)
        opt.eta = ηl
        for ps in get_trainable_params([layer])
            Flux.Optimise.update!(opt, ps, grads[ps])
        end
        ηl *= 2.6
    end
    return
end

"""
    train_classifier!(classifier::TextClassifier=TextClassifier(), classes::Integer=1,
            data_loader::Channel=imdb_classifier_data, hidden_layer_size::Integer=50;kw...)

It contains main training loops for training a defined classifer for specified classes and data.
Usage is discussed in the docs.
"""
function train_classifier!(classifier::TextClassifier=TextClassifier(), classes::Integer=1,
    data_loader::Channel=imdb_classifier_data, hidden_layer_size::Integer=50;
    stlr_cut_frac::Float64=0.1, stlr_ratio::Number=32, stlr_η_max::Float64=0.01,
    val_loader::Channel=nothing, cross_val_batches::Union{Colon, Integer}=:,
    epochs::Integer=1, checkpoint_itvl=5000, tracked_steps::Integer=32)

    trainable = []
    append!(trainable, [classifier.rnn_layers[[1, 3, 5, 7]]...])
    push!(trainable, [classifier.linear_layers[1:2]...])
    push!(trainable, [classifier.linear_layers[4:5]...])
    opts = [ADAM(0.001, (0.7, 0.99)) for i=1:length(trainable)]
    gpu!.(classifier.rnn_layers)

    for epoch=1:epochs
        println("Epoch: $epoch")
        gen = data_loader()
        num_of_iters = take!(gen)
        cut = num_of_iters * epochs * stlr_cut_frac
        for iter=1:num_of_iters

            # Slanted triangular learning rates
            t = iter + (epoch-1)*num_of_iters
            p_frac = (iter < cut) ? iter/cut : (1 - ((iter-cut)/(cut*(1/stlr_cut_frac-1))))
            ηL = stlr_η_max*((1+p_frac*(stlr_ratio-1))/stlr_ratio)

            # Gradual-unfreezing Step with discriminative fine-tuning
            unfreezed_layers, cur_opts = (epoch < length(trainable)) ? (trainable[end-epoch+1:end], opts[end-epoch+1:end]) : (trainable, opts)
            discriminative_step!(unfreezed_layers, classifier, gen, tracked_steps,ηL, cur_opts)

            reset_masks!.(classifier.rnn_layers)    # reset all dropout masks
        end
        println("Train set accuracy: $trn_accu , Training loss: $trn_loss")
        if val_loader != nothing
            val_loss, val_acc, val_precisions, val_reacalls, val_F1_scores = validate(classifer, val_loader)
        else
            continue
        end
        #!(val_loader isa nothing) ? (val_loss, val_acc, val_precisions, val_reacalls, val_F1_scores = validate(classifer, val_loader)) : continue
        println("Cross validation loss: $val_loss")
        println("Cross validation accuracy:\n $val_acc")
        println("Cross validation class wise Precisions:\n $val_precisions")
        println("Cross validation class wise Recalls:\n $val_recalls")
        println("Cross validation class wise F1 scores:\n $val_F1_scores")
    end
end

"""
    predict(tc::TextClassifier, text_sents::Corpus)

This function can be used to test the model after training.
It returns the predictions done by the model for given `Corpus` of `Documents`
All the preprocessing related to the used vocabulary should be done before using this function.
Use `prepare!` function to do preprocessing
"""
function predict(tc::TextClassifier, text_sents::Corpus)
    classifier = tc
    Flux.testmode!(classifier)
    predictions = []
    expr(x) = indices(x, classifier.vocab, "_unk_")
    for text in text_sents
        tokens_ = tokens(text)
        h = classifier.rnn_layers.(expr.(tokens_))
        probability_dist = classifier.linear_layers(h)
        class = argmax(probaility_dist)
        push!(predictions, class)
    end
    return predictions
end