text:
using Gen: get_child, get_child_num, get_parent

@testset "recurse node numbering" begin

    # max_branch = 1
    @test get_child(1, 1, 1) == 2
    @test get_child(2, 1, 1) == 3
    @test get_child(3, 1, 1) == 4
    @test get_child_num(2, 1) == 1
    @test get_child_num(3, 1) == 1
    @test get_child_num(4, 1) == 1
    @test get_parent(2, 1) == 1
    @test get_parent(3, 1) == 2
    @test get_parent(4, 1) == 3

    # max_branch = 2
    @test get_child(1, 1, 2) == 2
    @test get_child(1, 2, 2) == 3
    @test get_child(2, 1, 2) == 4
    @test get_child(2, 2, 2) == 5
    @test get_child(3, 1, 2) == 6
    @test get_child_num(2, 2) == 1
    @test get_child_num(3, 2) == 2
    @test get_child_num(4, 2) == 1
    @test get_child_num(5, 2) == 2
    @test get_child_num(6, 2) == 1
    @test get_parent(2, 2) == 1
    @test get_parent(3, 2) == 1
    @test get_parent(4, 2) == 2
    @test get_parent(5, 2) == 2
    @test get_parent(6, 2) == 3

    # max_branch = 3
    @test get_child(1, 1, 3) == 2
    @test get_child(1, 2, 3) == 3
    @test get_child(1, 3, 3) == 4
    @test get_child(2, 1, 3) == 5
    @test get_child(2, 2, 3) == 6
    @test get_child(2, 3, 3) == 7
    @test get_child(3, 1, 3) == 8
    @test get_child(3, 2, 3) == 9
    @test get_child(3, 3, 3) == 10
    @test get_child(4, 1, 3) == 11
    @test get_child_num(2, 3) == 1
    @test get_child_num(3, 3) == 2
    @test get_child_num(4, 3) == 3
    @test get_child_num(5, 3) == 1
    @test get_child_num(6, 3) == 2
    @test get_child_num(7, 3) == 3
    @test get_child_num(8, 3) == 1
    @test get_child_num(9, 3) == 2
    @test get_child_num(10, 3) == 3
    @test get_child_num(11, 3) == 1
    @test get_parent(2, 3) == 1
    @test get_parent(3, 3) == 1
    @test get_parent(4, 3) == 1
    @test get_parent(5, 3) == 2
    @test get_parent(6, 3) == 2
    @test get_parent(7, 3) == 2
    @test get_parent(8, 3) == 3
    @test get_parent(9, 3) == 3
    @test get_parent(10, 3) == 3
    @test get_parent(11, 3) == 4
end

@testset "simple pcfg" begin

    @gen function pcfg_production(_::Nothing)
        production_rule = @trace(categorical([0.24, 0.26, 0.23, 0.27]), :rule)
        if production_rule == 1
            # aSa; one child
            num_children = 1
        elseif production_rule == 2
            # bSb; one child
            num_children = 1
        elseif production_rule == 3
            # aa; no child
            num_children = 0
        else
            # bb; no child
            num_children = 0
        end

        return Production(production_rule, [nothing for _=1:num_children])
    end

    @gen function pcfg_aggregation(production_rule::Int, child_outputs::Vector{String})
        prefix = @trace(bernoulli(0.4), :prefix) ? "." : "-"
        local str::String
        if production_rule == 1
            @assert length(child_outputs) == 1
            str = "($(prefix)a$(child_outputs[1])a)"
        elseif production_rule == 2
            @assert length(child_outputs) == 1
            str = "($(prefix)b$(child_outputs[1])b)"
        elseif production_rule == 3
            @assert length(child_outputs) == 0
            str = "($(prefix)aa)"
        else
            @assert length(child_outputs) == 0
            str = "($(prefix)bb)"
        end
        return str
    end

    pcfg = Recurse(pcfg_production, pcfg_aggregation, 1, Nothing, Int, String)

    # test that each of the most probable strings are all produced
    function test_strings(strings)
        @test "(.aa)" in strings
        @test "(-aa)" in strings
        @test "(.bb)" in strings
        @test "(-bb)" in strings
        @test "(.a(.aa)a)" in strings
        @test "(-a(.aa)a)" in strings
        @test "(.a(-aa)a)" in strings
        @test "(-a(-aa)a)" in strings
        @test "(.b(.bb)b)" in strings
        @test "(-b(.bb)b)" in strings
        @test "(.b(-bb)b)" in strings
        @test "(-b(-bb)b)" in strings
        @test "(.a(.bb)a)" in strings
        @test "(-a(.bb)a)" in strings
        @test "(.a(-bb)a)" in strings
        @test "(-a(-bb)a)" in strings
        @test "(.b(.aa)b)" in strings
        @test "(-b(.aa)b)" in strings
        @test "(.b(-aa)b)" in strings
        @test "(-b(-aa)b)" in strings
    end

    # test Julia call
    @test isa(pcfg(nothing, 1), String)

    # test generate
    Random.seed!(1)
    strings = Set{String}()
    for i=1:1000
        (trace, _) = generate(pcfg, (nothing, 1), EmptyChoiceMap())
        push!(strings, get_retval(trace))
    end
    test_strings(strings)

    # test simulate
    Random.seed!(1)
    strings = Set{String}()
    for i=1:1000
        trace = simulate(pcfg, (nothing, 1))
        push!(strings, get_retval(trace))
    end
    test_strings(strings)

    # apply generate to a complete choice map that produces "(.b(.a(.b(-bb)b)a)b)"
    # sequence of production rules: 2 -> 1 -> 2 -> 4
    expected_weight = log(0.26) + log(0.24) + log(0.26) + log(0.27) + log(0.4) + log(0.4) + log(0.4) + log(0.6)
    constraints = choicemap()
    constraints[(1, Val(:production)) => :rule] = 2
    constraints[(1, Val(:aggregation)) => :prefix] = true
    constraints[(2, Val(:production)) => :rule] = 1
    constraints[(2, Val(:aggregation)) => :prefix] = true
    constraints[(3, Val(:production)) => :rule] = 2
    constraints[(3, Val(:aggregation)) => :prefix] = true
    constraints[(4, Val(:production)) => :rule] = 4
    constraints[(4, Val(:aggregation)) => :prefix] = false
    (trace, actual_weight) = generate(pcfg, (nothing, 1), constraints)
    @test isapprox(actual_weight, expected_weight)
    @test isapprox(get_score(trace), actual_weight)
    @test get_args(trace) == (nothing, 1)
    @test get_retval(trace) == "(.b(.a(.b(-bb)b)a)b)"
    choices = get_choices(trace)
    @test choices[(1, Val(:production)) => :rule] == 2
    @test choices[(1, Val(:aggregation)) => :prefix] == true
    @test choices[(2, Val(:production)) => :rule] == 1
    @test choices[(2, Val(:aggregation)) => :prefix] == true
    @test choices[(3, Val(:production)) => :rule] == 2
    @test choices[(3, Val(:aggregation)) => :prefix] == true
    @test choices[(4, Val(:production)) => :rule] == 4
    @test choices[(4, Val(:aggregation)) => :prefix] == false

    # update non-structure choice
    new_constraints = choicemap()
    new_constraints[(3, Val(:aggregation)) => :prefix] = false
    (new_trace, weight, retdiff, discard) = update(
        trace, (nothing, 1), (UnknownChange(), UnknownChange()), new_constraints)
    @test isapprox(weight, log(0.6) - log(0.4))
    expected_score = log(0.26) + log(0.24) + log(0.26) + log(0.27) + log(0.4) + log(0.4) + log(0.6) + log(0.6)
    @test isapprox(get_score(new_trace), expected_score)
    @test get_args(new_trace) == (nothing, 1)
    @test get_retval(new_trace) == "(.b(.a(-b(-bb)b)a)b)"
    choices = get_choices(new_trace)
    @test choices[(1, Val(:production)) => :rule] == 2
    @test choices[(1, Val(:aggregation)) => :prefix] == true
    @test choices[(2, Val(:production)) => :rule] == 1
    @test choices[(2, Val(:aggregation)) => :prefix] == true
    @test choices[(3, Val(:production)) => :rule] == 2
    @test choices[(3, Val(:aggregation)) => :prefix] == false
    @test choices[(4, Val(:production)) => :rule] == 4
    @test choices[(4, Val(:aggregation)) => :prefix] == false
    @test discard[(3, Val(:aggregation)) => :prefix] == true
    @test length(collect(get_submaps_shallow(discard))) == 1
    @test length(collect(get_values_shallow(discard))) == 0
    @test length(collect(get_submaps_shallow(get_submap(discard,(3, Val(:aggregation)))))) == 0
    @test length(collect(get_values_shallow(get_submap(discard,(3, Val(:aggregation)))))) == 1
    @test retdiff == UnknownChange()

    # update structure choice, so that string becomes: (.b(.a(.aa)a)b)
    # note: we reuse the prefix choice from node 3 (true)
    new_constraints = choicemap()
    new_constraints[(3, Val(:production)) => :rule] = 3 # change from rule 2 to rule 3
    (new_trace, weight, retdiff, discard) = update(
        trace, (nothing, 1), (UnknownChange(), UnknownChange()), new_constraints)
    @test isapprox(weight, log(0.23) - log(0.26) - log(0.27) - log(0.6))
    @test isapprox(get_score(new_trace), log(0.26) + log(0.24) + log(0.23) + log(0.4) + log(0.4) + log(0.4))
    @test get_args(new_trace) == (nothing, 1)
    @test get_retval(new_trace) == "(.b(.a(.aa)a)b)"
    choices = get_choices(new_trace)
    @test choices[(1, Val(:production)) => :rule] == 2
    @test choices[(1, Val(:aggregation)) => :prefix] == true
    @test choices[(2, Val(:production)) => :rule] == 1
    @test choices[(2, Val(:aggregation)) => :prefix] == true
    @test choices[(3, Val(:production)) => :rule] == 3
    @test choices[(3, Val(:aggregation)) => :prefix] == true
    @test isempty(get_submap(choices, (4, Val(:production)))) # FAIL
    @test isempty(get_submap(choices, (4, Val(:aggregation))))
    @test discard[(3, Val(:production)) => :rule] == 2
    @test !has_value(discard, (3, Val(:aggregation)) => :prefix)
    @test discard[(4, Val(:production)) => :rule] == 4
    @test discard[(4, Val(:aggregation)) => :prefix] == false
    @test retdiff == UnknownChange()

end

max_stars_repo_name: georgematheos/Gen

max_issues_repo_name: georgematheos/Gen

max_forks_repo_name: georgematheos/Gen


text:
# not working.

using LinearAlgebra, FFTW
import BSON, Statistics, Random
import PyPlot

import NMRSpectraSimulator

include("../src/NMRCalibrate.jl")
import .NMRCalibrate

# for loading something with Interpolations.jl
import OffsetArrays
import Interpolations

import PlotlyJS
import Plots
Plots.plotly()

import Destruct

include("./helpers/final_helpers.jl")
include("./helpers/resonance_helpers.jl")
include("./helpers/plot_resonance.jl")

PyPlot.close("all")
fig_num = 1

Random.seed!(25)
PyPlot.matplotlib["rcParams"][:update](["font.size" => 22, "font.family" => "serif"])

# base_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/final/"
# project_set = filter(isdir, readdir(base_dir; join=true)) # all dirs in base_dir

project_set = Vector{String}(undef, 0)
tmp = "/home/roy/MEGAsync/outputs/NMR/calibrate/final/D-(+)-Glucose-NRC-600"
push!(project_set, tmp)

function loopplotresonance(project_set)

    for i = 1:length(project_set)
        projects_dir = project_set[i]

        println("Now on $(projects_dir)")
        plotresonancegroups(projects_dir)
        println()
    end

end

### batch.
loopplotresonance(project_set)
### end batch.

ext: jl


text:
# list of all tags used by all physics modules

# header guard
if !isdefined(:tag_defs)
  tag_defs = false
end

if !tag_defs
  tag_defs = true
  # advection
  global const TAG_COMPLEX = "TAG_COMPLEX"
  global const TAG_BC = "TAG_BC"
  global const TAG_FLUX = "TAG_FLUX"
  global const TAG_VOLUMEINTEGRALS = "TAG_VOLUMEINTEGRAL"
  global const TAG_CONVERGENCE = "TAG_CONVERGENCE"
  global const TAG_MMS = "TAG_MMS"
  global const TAG_FRONTEND = "TAG_FRONTEND"
  global const TAG_ADJOINT = "TAG_ADJOINT"
  global const TAG_VISCOUS = "TAG_VISCOUS"
  global const TAG_CN = "TAG_CN"

  # euler
  global const TAG_ENTROPYVARS = "TAG_ENTROPYVARS"
  global const TAG_NLSOLVERS = "TAG_NLSOLVERS"
  global const TAG_ESS = "TAG_ESS"
  global const TAG_UTILS = "TAG_UTILS"
  global const TAG_INPUT = "TAG_INPUT"
  global const TAG_REVERSEMODE = "TAG_REVERSEMODE"
  global const TAG_MISC = "TAG_MISC"
  global const TAG_CURVILINEAR = "TAG_CURVILINEAR"
  global const TAG_HOMOTOPY = "TAG_HOMOTOPY"
  global const TAG_STAGGERED = "TAG_STAGGERED"
  global const TAG_PARALLEL_DERIVATIVES = "TAG_PARALLEL_DERIVATIVES"
  global const TAG_CHECKPOINT = "TAG_CHECKPOINT"
  global const TAG_JAC = "TAG_JAC"
  global const TAG_FUNCTIONAL = "TAG_FUNCTIONAL"

  global const TAG_PARALLEL = "TAG_PARALLEL"

  # types of tests
  global const TAG_SHORTTEST = "TAG_SHORTTEST"
  global const TAG_LONGTEST = "TAG_LONGTEST"
  global const LengthTags = Set([TAG_SHORTTEST, TAG_LONGTEST])

  global const TAG_TMP = "TAG_TMP"  # a temporary tag
end


text:
module DiscretePredictors

    using Printf

    include( "./Trie.jl" )
    include( "./BasePredictor.jl" )
    include( "./ALZ.jl" )
    include( "./DG.jl" )
    include( "./KOM.jl")
    include( "./LeZiUpdate.jl" )
    include( "./LZ78.jl" )
    include( "./dHedgePPM.jl" )      # Discounted HEDGEd KOM
    include( "./dHedgePPM_1.jl" )
    include( "./AdaptiveMPP.jl" )
    include( "./ewPPM.jl" )

    export
        Trie,
        # Predictors
        AdaptiveMPP,
        BasePredictor,
        ALZ,
        DG,
        dHedgePPM,
        dHedgePPM_1,
        acc_dHPPM,
        ewPPM,
        KOM,
        LeZiUpdate,
        LZ78,
        # functions
        add!,
        predict,
        info_string,
        unique_string,
        get_best_symbol,
        size

"""
A Julia package for online discrete sequence prediction.

API Overview:
- `p = Predictor{SymbolType}(parameters...)` creates a predictor instance `p` with
    prediction algorithm `Predictor` (see choices below) of symbols of `SymbolType` with
    specified the `parameters`
- `add!( p, sym )` adds a symbol to the model
- `predict( p )` computes a probability distribution over the symbols seen so far by the
    model and returns it as dictionary of `Dict{SymbolType,Float64}`
- `info_string( p )` returns a human readable string about the predictor `p`
- `unique_string( p )` returns a unique string for the predictor `p` with paramter configuration
- `get_best_symbol( p )`  returns the symbol of type `SymbolType` with highest probability
    under current context
- `size( p )` returns the number of nodes in current prediction model `p.model`

Available Predictors:
- Adaptive MPP ([`AdaptiveMPP`](@ref))
- Active LeZi ([`ALZ`](@ref))
- Dependency Graph ([`DG`](@ref))
- Discounted HEDGE on KOM ([`dHedgePPM`](@ref))
- Error Weighted PPM ([`ewPPM`](@ref))
- K-th Order Markov Model ([`KOM`](@ref))
- LeZi-Update ([`LeZiUpdate`](@ref))
- LeZi78 ([`LZ78`](@ref))

## Example
```julia
# Initialize a predictor
julia> p = KOM{Char}(4)
DiscretePredictors.KOM{Char}([*] (0)
, Char[], 4)

# Add some symbols
julia> add!( p, 'a' )

julia> add!( p, 'b' )

julia> add!( p, 'c' )

julia> add!( p, 'b' )

# Print out the model
julia> p
DiscretePredictors.KOM{Char}([*] (4)
+---[b] (2)
     +---[c] (1)
          +---[b] (1)
+---[a] (1)
     +---[b] (1)
          +---[c] (1)
               +---[b] (1)
+---[c] (1)
     +---[b] (1)
, ['a', 'b', 'c', 'b'], 4)

# Get prediction
julia> predict( p )
Dict{Char,Float64} with 3 entries:
  'b' => 0.25
  'a' => 0.125
  'c' => 0.625

# Get best symbol
julia> get_best_symbol( p )
'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)

julia> info_string( p )
"KOM(4)"

julia> unique_string( p )
"KOM_04"

julia> size( p )
10

```
"""
DiscretePredictors

end # module


text:
test_str1 = "mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
mem[8] = 11
mem[7] = 101
mem[8] = 0"

test_str2 = "mask = 000000000000000000000000000000X1001X
mem[42] = 100
mask = 00000000000000000000000000000000X0XX
mem[26] = 1"

test_list1 = split(test_str1, '\n')

test_list2 = split(test_str2, '\n')

list = readlines("day14_input.txt")

function mask_value(list)
    lines = map(x -> split(x, " = "), list)
    memory = Dict()
    mask = "X"^36
    for i in 1:length(list)
        if lines[i][1] == "mask"
            mask = lines[i][2]
        else
            value = bitstring(parse(Int, lines[i][2]))[end-35:end]
            masked_value = ""
            for j=1:36
                if mask[j] == 'X'
                    masked_value *= value[j]
                elseif mask[j] == '0'
                    masked_value *= '0'
                else
                    masked_value *= '1'
                end
            end
            if lines[i][1] in keys(memory)
                memory[lines[i][1]] = parse(Int, masked_value, base=2)
            else
                push!(memory, lines[i][1] => parse(Int, masked_value, base=2))
            end
        end
    end
    return sum(values(memory))
end

@show mask_value(test_list1) == 165
@show mask_value(list) == 15919415426101

function mask_address(list)
    lines = map(x -> split(x, " = ") |> y -> y[1] == "mask" ? (y[1], y[2]) :
        (parse(Int,(split(y[1], r"mem\[|\]")[2])), parse(Int, y[2])), list)
    memory = Dict()
    mask = "X"^36
    for i in 1:length(list)
        if lines[i][1] == "mask"
            mask = lines[i][2]
        else
            address = bitstring(lines[i][1])[end-35:end]
            value = lines[i][2]
            masked_address = map(x -> x[1] == '0' ? x[2] : x[1], zip(mask, address))
            num_x = count(==('X'), masked_address)
            pos_x = [i for (i,c) in enumerate(masked_address) if c=='X']
            num_comb = 2^num_x
            comb = [collect(bitstring(j)[end-num_x+1:end]) for j=1:num_comb]
            for j=1:num_comb
                x_address = deepcopy(masked_address)
                x_address[pos_x] = comb[j]
                if x_address in keys(memory)
                    memory[x_address] = value
                else
                    push!(memory, x_address => value)
                end
            end
        end
    end
    return sum(values(memory))
end

@show mask_address(test_list2) == 208
@show mask_address(list) == 3443997590975

nothing


text:
"""
    generate_stop_file(stop_txt::String, params::CityParams, out_file::String)

Convert the stops.txt GTFS file for a transit network to a dict that maps
integer stop ID to lat-long coordinates, and saves it to an output JSON file.
The CityParams is used to filter out trips outside the bounding box.
"""
function generate_stop_file(stop_txt::String, params::CityParams, out_file::String)

    stop_file = CSV.File(stop_txt)

    n_stops = length(stop_file.stop_id)
    stops_dict = Dict{Int64,LatLonCoords}()

    for (i, id_str) in enumerate(stop_file.stop_id)
        id = 0
        if ismissing(id_str)
            @info "Stop missing, continuing!"
        else
            if typeof(id_str) == String
                try
                    id = parse(Int64, id_str)
                catch ArgumentError
                    @info "Found non-integer stop id $(id_str); ignoring!"
                    continue
                end
            else
                id = id_str
            end

            @assert haskey(stops_dict, id) == false
            coords = LatLonCoords((lat = stop_file.stop_lat[i], lon = stop_file.stop_lon[i]))

            if is_in_bounds(params, coords)
                stops_dict[id] = coords
            end # if
        end
    end

    open(out_file, "w") do stop_out_file
        JSON.print(stop_out_file, stops_dict, 2)
    end

    return stops_dict

end

"""
    generate_trip_file(route_txt::String, trip_txt::String,
                       stop_time_txt::String, stop_coords::Dict{Int64,LatLonCoords},
                       out_file::String, rng::RNG) where {RNG <: AbstractRNG}

Parse the GTFS trips.txt file. Takes the routes.txt file, chooses any trip from that
route type, and generates the corresponding sequence of stop IDs from the stop_coords Dict.
"""
function generate_trip_file(route_txt::String, trip_txt::String,
                            stop_time_txt::String, stop_coords::Dict{Int64,LatLonCoords},
                            out_file::String, rng::RNG) where {RNG <: AbstractRNG}

    route_df = CSV.read(route_txt)
    trip_df = CSV.read(trip_txt)
    stop_time_df = CSV.read(stop_time_txt)

    all_trips_dict = Dict()
    out_trip_idx = 1
    trip_points = 0

    max_rel_time = 0.0

    for rid in route_df[:, :route_id]

        @info rid
        # Filter trips for route ID
        # Choose a trip ID at random
        trips_for_route = trip_df[trip_df.route_id .== rid, :trip_id]
        tid = trips_for_route[rand(rng, 1:end)]

        # Now look at the stop_times for that trip ID
        trip_dict = Dict{Int64,Dict}()

        # Get the stop sequence for that trip
        trip_stop_seq = stop_time_df[stop_time_df.trip_id .== tid, :]

        # Don't assume they are necessarily vertical
        n_stops = DataFrames.nrow(trip_stop_seq)
        is_valid_trip = true

        # Get base time
        base_time_str = trip_stop_seq[trip_stop_seq.stop_sequence .== 1, :arrival_time]
        base_time = hhmmss_to_seconds(base_time_str[1])

        for stop_seq = 1:n_stops

            # Get row for that stop sequence
            stop_row = trip_stop_seq[trip_stop_seq.stop_sequence .== stop_seq, :]

            stop_id = 0
            if typeof(stop_row.stop_id[1]) == String
                try
                    stop_id = parse(Int64, stop_row.stop_id[1])
                catch ArgumentError
                    @info "Found non-integer stop id $(stop_row.stop_id[1]), ignoring!"
                    is_valid_trip = false
                    break
                end
            else
                stop_id = stop_row.stop_id[1]
            end

            if ~(haskey(stop_coords, stop_id))
                @info "Trip ",tid," has out-of-bounds stop!"
                is_valid_trip = false
                break
            end

            # Calculate relative time
            stop_time = hhmmss_to_seconds(stop_row.arrival_time[1])
            @assert stop_time >= base_time
            rel_time = stop_time - base_time

            max_rel_time = (rel_time > max_rel_time) ? rel_time : max_rel_time

            trip_dict[stop_seq] = Dict("stop_id"=>stop_id, "arrival_time"=>rel_time)
        end

        if is_valid_trip
            all_trips_dict[out_trip_idx] = trip_dict
            out_trip_idx += 1
            trip_points += length(keys(trip_dict))
        end
    end

    @info "Total number of unique trips - $(out_trip_idx)"
    @info "Transit trip points - $(trip_points)"
    @info "MAx rel time - $(max_rel_time)"

    open(out_file, "w") do f
        JSON.print(f, all_trips_dict, 2)
    end
end


text:
st=Structure()
lcset=LoadCaseSet()
add_uniaxial_metal!(st,"q345",2e11,0.3,7850,1.17e-7)

open(joinpath(@__DIR__,"nodes.mz")) do f
    line=readline(f)
    while line!=""
        parse_block!(st,lcset,[line])
        line=readline(f)
    end
end

open(joinpath(@__DIR__,"trias.mz")) do f
    line=readline(f)
    while line!=""
        parse_block!(st,lcset,[line])
        line=readline(f)
    end
end

for node in values(st.nodes)
    if abs(node.loc[2])<1e-6 || abs(node.loc[2]-7.62)<1e-6
        set_nodal_restraint!(st,node.id,true,true,true,false,false,false)
    end
end

add_static_case!(lcset,"DL",1.0)

@time begin
assembly=assemble!(st,lcset)
end

solve(assembly)

for node in values(st.nodes)
    if abs(node.loc[2]-3.81)<1e-3 && abs(node.loc[3]-3.81)<1e-3
        r=result_nodal_displacement(assembly,"DL",node.id)
        @show r
    end
end


text:
const ComputableType = Union{AbstractFloat,Complex}
const version = "0.1.7"

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
using FastGaussQuadrature
using SpecialFunctions
using ApproxFun
using Interpolations

export TracyWidom

struct TracyWidom <: ContinuousUnivariateDistribution
    beta::Integer
end

function cdf(d::TracyWidom, x::Real)
    if d.beta == 1
        _fredholm_det(_V_airy_kernel, x, 100)
    elseif d.beta == 2
        _fredholm_det(_airy_kernel, x, 100)
    elseif d.beta == 4
        (_fredholm_det(_V_airy_kernel, x * sqrt(2), 100) + _fredholm_det(_V_airy_kernel, x * sqrt(2), 100, -1))/2
    else
        throw(DomainError(d.beta))
    end
end

_TWfs = Dict()

function _TWf(x, beta)
    if !(beta in keys(_TWfs))
        dom = Chebyshev(-5..5)
        D = Derivative(dom)
        _TWfs[beta] = D * Fun(x -> cdf(TracyWidom(beta), x), dom)
    end
    _TWfs[beta](x)
end

function pdf(d::TracyWidom, x::Real)
    _TWf(x, d.beta)
end

_TWinvs = Dict()

function _TWinv(t, beta)
    (0 < t < 1) || throw(DomainError(t))

    if !(beta in keys(_TWinvs))
        # bounds outside which we use limiting representation
        u = 4
        l = -5

        # TODO: make a more principled grid using tail behaviour
        n = 1000
        y = range(l, u, length=n)
        x = cdf.(TracyWidom(beta), y)

        # Transform limiting ends to maintain differentiability
        righta = pdf(TracyWidom(beta), u)/(beta * sqrt(u) * exp(-beta * u^(3/2)*2/3))
        rightb = x[end] - righta * (1 - exp(-beta * u^(3/2) * 2/3))

        lefta = pdf(TracyWidom(beta), l)/(beta * l^2/8 * exp(beta * l^3/24))
        leftb = x[1] - lefta * exp(-beta * abs(l)^3/24)

        itp = interpolate((x,), y, Gridded(Linear()))

        _TWinvs[beta] = [righta, rightb, x[end], lefta, leftb, x[1], itp]
    end

    # TODO: make the indexing less terrible
    if t > _TWinvs[beta][3]
        s = (t - _TWinvs[beta][2])/_TWinvs[beta][1]
        return (log(1/(1 - s)) * 3/(2 * beta))^(2/3)
    elseif t < _TWinvs[beta][6]
        s = (t - _TWinvs[beta][5])/_TWinvs[beta][4]
        return -(log(1/s) * 24/beta)^(1/3)
    else
        return _TWinvs[beta][end](t)
    end

end

function quantile(d::TracyWidom, q::Real)
    _TWinv(q, d.beta)
end

function _phi(s, xi)
    s + 10 * tan(pi * (xi + 1)/4)
end

function _Dphi(s, xi)
    pi * 5/2 * sec(pi * (xi + 1)/4)^2
end

"""
Transform a kernel on L_2(s, infinity) to one on L_2(-1, 1)
"""
function _transform_kernel(kernel, s)
    (x, y) -> sqrt(_Dphi(s, x) * _Dphi(s, y)) * kernel(_phi(s, x), _phi(s, y))
end

"""
Approximate fredholm determinant of a kernel on (s, infinity)
"""
function _fredholm_det(kernel, s, N, z=1)

    (nodes, weights) = gausslegendre(N)
    trans_kernel = _transform_kernel(kernel, s)

    M = I - z * sqrt.(weights * weights') .* trans_kernel.(nodes', nodes)
    det(M)
end

function _airy_kernel(x, y)
    if x == y
        airyaiprime(x)^2 - x * airyai(x)^2
    else
        (airyai(x) * airyaiprime(y) - airyaiprime(x) * airyai(y))/(x - y)
    end
end

function _V_airy_kernel(x, y)
    airyai((x + y)/2)/2
end


text:
export GeodesicOnTheSphere

"""
    GeodesicOnTheSphere(start, target)

Homotopy is the geodesic from `g=start/|start|` (t=1) to `f=target/|target|` (t=0):

```H(x,t) = (cos(tα) - sin (tα)cos(α)/sin(α)) f + sin(tα) / sin(α) * g```

where ``α = cos |<f,g>|``. The constructor automatically homgenizes `start` and `target`.

`start` and `target` have to match and to be one of the following
* `Vector{<:MP.AbstractPolynomial}` where `MP` is [`MultivariatePolynomials`](https://github.com/blegat/MultivariatePolynomials.jl)
* `MP.AbstractPolynomial`
* `Vector{<:FP.Polynomial}` where `FP` is [`FixedPolynomials`](https://github.com/saschatimme/FixedPolynomials.jl)

    GeodesicOnTheSphere{T}(start, target)

You can also force a specific coefficient type `T`.
"""
mutable struct GeodesicOnTheSphere{T<:Number} <: AbstractPolynomialHomotopy{T}
    start::Vector{FP.Polynomial{T}}
    target::Vector{FP.Polynomial{T}}
    α::Float64

    # only used for conversions
    function GeodesicOnTheSphere{T}(start::Vector{FP.Polynomial{T}}, target::Vector{FP.Polynomial{T}}, α::Float64) where {T<:Number}
        new(start, target, α)
    end

    function GeodesicOnTheSphere{T}(start::Vector{FP.Polynomial{T}}, target::Vector{FP.Polynomial{T}}) where {T<:Number}
        start, target = homogenize(start, target)

        @assert length(start) == length(target) "Expected the same number of polynomials, but got $(length(start)) and $(length(target))"
        s_nvars = maximum(FP.nvariables.(start))
        @assert all(s_nvars .== FP.nvariables.(start)) "Not all polynomials of the start system have $(s_nvars) variables."
        t_nvars = maximum(FP.nvariables.(target))
        @assert all(t_nvars .== FP.nvariables.(target)) "Not all polynomials of the target system have $(t_nvars) variables instead of $(s_nvars)."
        @assert s_nvars == t_nvars "Expected start and target system to have the same number of variables, but got $(s_nvars) and $(t_nvars)."

        FP.scale_coefficients!.(start, inv(FP.weylnorm(start)))
        FP.scale_coefficients!.(target, inv(FP.weylnorm(target)))

        α = acos(convert(Float64, real(FP.weyldot(start,target))))
        if α > π / 2
            α = π / 2 - α
            FP.scale_coefficients!.(start, -one(T))
        end

        new(start, target, α)
    end

    function GeodesicOnTheSphere{T}(start, target) where {T<:Number}
        s, t = construct(T, start, target)
        GeodesicOnTheSphere{T}(s, t)
    end
end

function GeodesicOnTheSphere(start, target)
    T, s, t = construct(start, target)
    GeodesicOnTheSphere{T}(s, t)
end

function Base.deepcopy(H::GeodesicOnTheSphere{T}) where T
    GeodesicOnTheSphere{T}(deepcopy(H.start), deepcopy(H.target), H.α)
end
#
# PROMOTION AND CONVERSION
#
function Base.promote_rule(::Type{GeodesicOnTheSphere{T}}, ::Type{GeodesicOnTheSphere{S}}) where {S<:Number,T<:Number}
    GeodesicOnTheSphere{promote_type(T,S)}
end
function Base.promote_rule(::Type{GeodesicOnTheSphere{T}}, ::Type{S}) where {S<:Number,T<:Number}
    GeodesicOnTheSphere{promote_type(T,S)}
end
function Base.promote_rule(::Type{GeodesicOnTheSphere},::Type{S}) where {S<:Number}
    GeodesicOnTheSphere{S}
end

function Base.convert(::Type{GeodesicOnTheSphere{T}}, H::GeodesicOnTheSphere) where {T}
    GeodesicOnTheSphere{T}(convert.(FP.Polynomial{T}, H.start), convert.(FP.Polynomial{T}, H.target), H.α)
end

# EVALUATION + DIFFERENTATION
#
(H::GeodesicOnTheSphere)(x,t) = evaluate(H,x,t)

function λ(α, t)
    λ_1 = sin(t * α) / sin(α)
    λ_2 = cos(t * α) - λ_1 * cos(α)

    λ_1, λ_2
end

function dλ(α, t)
    λ_1_dot = α * cos(t * α) / sin(α)
    λ_2_dot = - α * sin(t * α) - λ_1_dot * cos(α)

    λ_1_dot, λ_2_dot
end

function evaluate!(u::AbstractVector, H::GeodesicOnTheSphere{T}, x::Vector, t::Number) where T
    λ_1, λ_2 = λ(H.α, t)
    for i = 1:length(H.target)
        u[i] = λ_1 * FP.evaluate(H.start[i], x) + λ_2 * FP.evaluate(H.target[i], x)
    end
    u
end

function evaluate!(u::AbstractVector{T}, H::GeodesicOnTheSphere, x::Vector, t::Number, cfg::PolynomialHomotopyConfig, precomputed=false) where {T<:Number}
    evaluate_start_target!(cfg, H, x, precomputed)
    λ_1, λ_2 = λ(H.α, t)
    u .= λ_1 .* value_start(cfg) .+ λ_2 .* value_target(cfg)
end

function jacobian!(u::AbstractMatrix, H::GeodesicOnTheSphere{T}, x::AbstractVector, t, cfg::PolynomialHomotopyConfig, precomputed=false) where {T<:Number}
    jacobian_start_target!(cfg, H, x, precomputed)
    λ_1, λ_2 = λ(H.α, t)

    u .= λ_1 .* jacobian_start(cfg) .+ λ_2 .* jacobian_target(cfg)
end
function jacobian!(r::JacobianDiffResult, H::GeodesicOnTheSphere{T}, x::AbstractVector, t, cfg::PolynomialHomotopyConfig) where {T<:Number}
    evaluate_and_jacobian_start_target!(cfg, H, x)
    λ_1, λ_2 = λ(H.α, t)

    r.value .= λ_1 .* value_start(cfg) .+ λ_2 .* value_target(cfg)
    r.jacobian .= λ_1 .* jacobian_start(cfg) .+ λ_2 .* jacobian_target(cfg)
    r
end

function dt!(u, H::GeodesicOnTheSphere{T}, x::AbstractVector, t, cfg::PolynomialHomotopyConfig, precomputed=false) where {T<:Number}
    evaluate_start_target!(cfg, H, x, precomputed)
    λ_1_dot, λ_2_dot = dλ(H.α, t)
    u .= λ_1_dot .* value_start(cfg) .+ λ_2_dot .* value_target(cfg)
end
function dt!(r::DtDiffResult, H::GeodesicOnTheSphere{T}, x::AbstractVector, t, cfg::PolynomialHomotopyConfig, precomputed=false) where {T<:Number}
    evaluate_start_target!(cfg, H, x, precomputed)
    λ_1, λ_2 = λ(H.α, t)
    λ_1_dot, λ_2_dot = dλ(H.α, t)

    r.value = λ_1 .* value_start(cfg) .+ λ_2 .* value_target(cfg)
    r.dt = λ_1_dot .* value_start(cfg) .+ λ_2_dot .* value_target(cfg)
    r
end

function weylnorm(H::GeodesicOnTheSphere{T}) where {T<:Number}
    function (t)
         real(one(T))
     end
end

function gammatrick!(H::GeodesicOnTheSphere{T}, γ::Union{AbstractFloat, Complex}) where {T<:Complex}
    FP.scale_coefficients!.(H.start, convert(T, γ))

    FP.scale_coefficients!.(H.start, inv(FP.weylnorm(H.start)))
    α = acos(convert(Float64, real(FP.weyldot(H.start, H.target))))
    if α > π / 2
        α = π / 2 - α
        FP.scale_coefficients!.(H.start, -one(T))
    end
    H.α = α

    H
end


text:
using BlockSparseArrays
using Random
using SparseArrays
using StaticArrays
using Test

include("vector.jl")
include("matrix.jl")


text:
function run_repeats()
    for set in 1:9
        # create .sh file to submit jobs
        filename = "poly.set$set.sh"
        open(filename, "w") do io
            println(io, "#!/bin/bash")
            println(io, "#\$ -cwd")
            println(io, "# error = Merged with joblog")
            println(io, "#\$ -o joblog.\$JOB_ID")
            println(io, "#\$ -j y")
            println(io, "#\$ -pe shared 16")
            println(io, "#\$ -l h_rt=24:00:00,h_data=1G,arch=intel-gold-61*")
            println(io, "# Email address to notify")
            println(io, "## \$ -M \$USER@mal")
            println(io, "# Notify when:")
            println(io, "#\$ -m bea")
            println(io, "")
            println(io, "#save job info on joblog:")
            println(io, "echo \"Job \$JOB_ID started on:   \" `hostname -s`")
            println(io, "echo \"Job \$JOB_ID started on:   \" `date `")
            println(io, "")
            println(io, "# load the job environment:")
            println(io, ". /u/local/Modules/default/init/modules.sh")
            println(io, "module load julia/1.5.4")
            println(io, "module li")
            println(io, "which julia")
            println(io, "")
            println(io, "# run code")
            println(io, "export JULIA_NUM_THREADS=16")
            println(io, "echo 'julia NFBC_chr1_sim.jl (run IHT/GEMMA/mvPLINK on NFBC data (polygenic model))'")
            println(io, "pwd; julia NFBC_chr1_sim.jl $set")
            println(io, "")
            println(io, "#echo job info on joblog:")
            println(io, "echo \"Job \$JOB_ID ended on:   \" `hostname -s`")
            println(io, "echo \"Job \$JOB_ID ended on:   \" `date `")
            println(io, "#echo \" \"")
        end

        # submit job
        run(`qsub $filename`)
        rm(filename, force=true)
        sleep(2)
    end
end
run_repeats()


text:
"""
Lightweight polynomial type
"""
struct Poly{T, A<:AbstractVector{T}}
    c::A
end
Poly(c::Number...) = Poly(c)
Poly(c::Tuple) = Poly([float.(c)...])
Poly(c::AbstractVector{Int}) = Poly(float.(c))

coeffs(p::Poly) = p.c
degree(p::Poly) = length(coeffs(p)) - 1
(p::Poly)(x) = evalpoly(x, coeffs(p))

add!(p::Poly, a::Number) = (p.c[1] += a; return p)
sub!(p::Poly, a::Number) = (p.c[1] -= a; return p)

# Recall: p(x) = Σᵢ cᵢ ⋅ xⁱ⁻¹
derivative(p::Poly{T}) where {T} = Poly(T[(i-1) * coeffs(p)[i] for i in 2:degree(p)+1]) # ∂/∂x p(x) = Σᵢ (i-1) ⋅ cᵢ ⋅ xⁱ⁻²
integral(p::Poly{T}) where {T} = Poly(T[i == 0 ? zero(T) : coeffs(p)[i] / i for i in 0:degree(p)+1]) # ∫₀ˣ p(x) = Σᵢ (cᵢ / i) ⋅ xⁱ
PolynomialRoots.roots(p::Poly) = PolynomialRoots.roots(coeffs(p))
Base.extrema(p::Poly) = PolynomialRoots.roots(derivative(p))

####
#### Spline utils
####

function make_spline(X::AbstractVector, Y::AbstractVector; deg_spline = min(3, length(X)-1))
    @assert length(X) == length(Y) && length(X) > 1
    return Dierckx.Spline1D(X, Y; k = deg_spline, bc = "extrapolate")
end

function build_polynomials(spl::Dierckx.Spline1D)
    k = spl.k
    t = Dierckx.get_knots(spl)[1:end-1]
    coeffs = zeros(k+1, length(t))
    coeffs[1, :] .= spl.(t)
    for m in 1:k
        coeffs[m+1, :] .= Dierckx.derivative(spl, t, m) ./ factorial(m)
    end
    return [Poly(coeffs[:, j]) for j in 1:size(coeffs, 2)]
end

# Global minimization through fitting a spline to data (X, Y)
function spline_opt(spl::Dierckx.Spline1D)
    knots = Dierckx.get_knots(spl)
    polys = build_polynomials(spl)
    x, y = knots[1], polys[1](0) # initial lefthand point
    @inbounds for (i, p) in enumerate(polys)
        x₀, x₁ = knots[i], knots[i+1] # spline section endpoints
        _x, _y = x₁, p(x₁ - x₀) # check right endpoint
        (_y < y) && (x = _x; y = _y)
        for rᵢ in extrema(p) # extrema(p) returns the zeros of derivative(p)
            if imag(rᵢ) ≈ 0 # real roots only
                xᵢ = x₀ + real(rᵢ)
                if x₀ <= xᵢ <= x₁ # filter roots within range
                    _x, _y = xᵢ, p(real(rᵢ))
                    (_y < y) && (x = _x; y = _y)
                end
            end
        end
    end
    return (; x, y)
end
spline_opt(X::AbstractVector, Y::AbstractVector; deg_spline = min(3, length(X)-1)) = spline_opt(make_spline(X, Y; deg_spline))

# Root finding through fitting a spline to data (X, Y)
function spline_root(spl::Dierckx.Spline1D, value::Number = 0)
    knots = Dierckx.get_knots(spl)
    polys = build_polynomials(spl)
    x = eltype(knots)(NaN)
    @inbounds for (i, p) in enumerate(polys)
        x₀, x₁ = knots[i], knots[i+1] # spline section endpoints
        # Solve `p(rᵢ) = value` via `p(rᵢ) - value = 0`
        for rᵢ in PolynomialRoots.roots(sub!(p, value))
            if imag(rᵢ) ≈ 0 # real roots only
                xᵢ = x₀ + real(rᵢ)
                if x₀ <= xᵢ <= x₁ # filter roots within range
                    x = isnan(x) ? xᵢ : min(x, xᵢ)
                end
            end
        end
    end
    return x
end
spline_root(X::AbstractVector, Y::AbstractVector, value::Number = 0; deg_spline = min(3, length(X)-1)) = spline_root(make_spline(X, Y; deg_spline), value)

####
#### Legacy spline utils
####
####    These algorithms were changed compared to the original MATLAB version:
####    instead of brute force searching through the splines, numerical methods
####    are performed which are much more efficient and much more accurate.
####    For direct comparison of results with the MATLAB version, the brute force
####    version is implimented and can be used by setting the `legacy` flag.
####

# MATLAB spline optimization performs global optimization by sampling the spline
# fit to data (X, Y) at points X[1]:0.001:X[end], and uses the minimum value.
# This isn't very efficient; instead we use exact spline optimization and return
# the (x,y) pair such that x ∈ X[1]:0.001:X[end] is nearest to the exact optimum
function spline_opt_legacy(spl::Dierckx.Spline1D)
    xopt, yopt = spline_opt(spl)
    knots = Dierckx.get_knots(spl)
    xs = knots[1]:eltype(knots)(0.001):knots[end] # from MATLAB version
    _, i0 = find_nearest(xs, xopt) # find nearest x in xs to xopt

    # Note that the above finds the x value nearest the true minimizer, but we need the x value corresponding to
    # the y value nearest the true minimum. Since we are near the minimum, search for a local minimum.
    @unpack x, y, i = local_gridsearch(spl, xs, i0)

    return (; x, y)
end
spline_opt_legacy(X::AbstractVector, Y::AbstractVector; deg_spline = min(3, length(X)-1)) = spline_opt_legacy(make_spline(X, Y; deg_spline))

# Similar to above, but removes the extra trick and instead performs
# exactly what the MATLAB implementation does
function spline_opt_legacy_slow(spl::Dierckx.Spline1D)
    knots = Dierckx.get_knots(spl)
    xs = knots[1]:eltype(knots)(0.001):knots[end] # from MATLAB version
    x, y = xs[1], spl(xs[1])
    for (i,xᵢ) in enumerate(xs)
        (i == 1) && continue
        yᵢ = spl(xᵢ)
        (yᵢ < y) && ((x, y) = (xᵢ, yᵢ))
    end
    return (; x, y)
end
spline_opt_legacy_slow(X::AbstractVector, Y::AbstractVector; deg_spline = min(3, length(X)-1)) = spline_opt_legacy_slow(make_spline(X, Y; deg_spline))

# Brute force root finding through fitting a spline to data (X, Y):
# MATLAB implementation of spline root finding performs root finding by sampling the
# spline fit to data (X, Y) at points X[1]:0.001:X[end], and uses the nearest value.
# This isn't very efficient; instead we use exact spline root finding and return
# the nearest x ∈ X[1]:0.001:X[end] such that the y value is nearest zero
function spline_root_legacy(spl::Dierckx.Spline1D, value = 0)
    # Find x value nearest to the root
    knots = Dierckx.get_knots(spl)
    xs = knots[1]:eltype(knots)(0.001):knots[end] # from MATLAB version
    xroot = spline_root(spl, value)
    _, i0 = find_nearest(xs, xroot) # find nearest x in xs to xroot

    # Note that the above finds the x value nearest the true root, but we need the x value corresponding to the
    # y value nearest to `value`. Since we are near the root, search for a local minimum in abs(spl(x)-value).
    dy(x) = abs(spl(x) - value)
    @unpack x, y, i = local_gridsearch(dy, xs, i0)

    return x
end
spline_root_legacy(X::AbstractVector, Y::AbstractVector, value = 0; deg_spline = min(3, length(X)-1)) = spline_root_legacy(make_spline(X, Y; deg_spline), value)

# Similar to above, but removes the extra trick and instead performs
# exactly what the MATLAB implementation does
function spline_root_legacy_slow(spl::Dierckx.Spline1D, value = 0)
    knots = Dierckx.get_knots(spl)
    xs = knots[1]:eltype(knots)(0.001):knots[end] # from MATLAB version
    x, y = xs[1], abs(spl(xs[1]) - value)
    for (i,xᵢ) in enumerate(xs)
        (i == 1) && continue
        yᵢ = abs(spl(xᵢ) - value)
        (yᵢ < y) && ((x, y) = (xᵢ, yᵢ))
    end
    return x
end
spline_root_legacy_slow(X::AbstractVector, Y::AbstractVector, value = 0; deg_spline = min(3, length(X)-1)) = spline_root_legacy_slow(make_spline(X, Y; deg_spline), value)

####
#### Surrogate functions over discrete grids
####

abstract type AbstractSurrogate{D,T} end

gridwidths(surr::AbstractSurrogate) = Tuple(abs.(last(surr.grid) - first(surr.grid)))
gridspacings(surr::AbstractSurrogate) = gridwidths(surr) ./ size(surr.grid)

struct CubicSplineSurrogate{T,F} <: AbstractSurrogate{1,T}
    f::F
    grid::Vector{SVector{1,T}}
    p::Vector{SVector{1,T}}
    u::Vector{T}
    npts::Base.RefValue{Int}
    legacy::Bool
end

function CubicSplineSurrogate(f, grid::Vector{SVector{1,T}}; legacy = false) where {T}
    CubicSplineSurrogate(f, grid, SVector{1,T}[], T[], Ref(0), legacy)
end

function update!(surr::CubicSplineSurrogate, I::CartesianIndex{1})
    p = surr.grid[I]
    pos = surr.npts[] + 1
    @inbounds for i in 1:surr.npts[]
        (p[1] <= surr.p[i][1]) && (pos = i; break)
    end
    u = surr.f(I)
    insertat!(surr.p, pos, p, surr.npts[]+1)
    insertat!(surr.u, pos, u, surr.npts[]+1)
    surr.npts[] += 1
    return surr
end

function insertat!(x::AbstractVector{T}, i, v::T, len = length(x)) where {T}
    if len > length(x)
        append!(x, similar(x, max(length(x), 1)))
    end
    last = v
    @inbounds for j in i:len
        x[j], last = last, x[j]
    end
    return x
end

function Base.empty!(surr::CubicSplineSurrogate)
    surr.npts[] = 0
    return surr
end

function suggest_point(surr::CubicSplineSurrogate{T}) where {T}
    npts = surr.npts[]
    ps = reinterpret(T, view(surr.p, 1:npts))
    us = view(surr.u, 1:npts)
    p, u = surr.legacy ? spline_opt_legacy_slow(ps, us) : spline_opt(ps, us)
    return SVector{1,T}(p), T(u)
end

struct HermiteSplineSurrogate{D,T,F,RK} <: AbstractSurrogate{D,T}
    fg::F
    grid::Array{SVector{D,T},D}
    ugrid::Array{T,D}
    spl::NormalHermiteSplines.ElasticNormalSpline{D,T,RK}
end

function HermiteSplineSurrogate(fg, grid::Array{SVector{D,T},D}, kernel = RK_H1(one(T))) where {D,T}
    spl = NormalHermiteSplines.ElasticNormalSpline(first(grid), last(grid), maximum(size(grid)), kernel)
    HermiteSplineSurrogate(fg, grid, zeros(T, size(grid)), spl)
end

function update!(surr::HermiteSplineSurrogate{D,T}, I::CartesianIndex{D}) where {D,T}
    u, ∇u = surr.fg(I)
    @inbounds p = surr.grid[I]
    insert!(surr.spl, p, u)
    @inbounds for i in 1:D
        eᵢ = basisvector(SVector{D,T}, i)
        insert!(surr.spl, p, eᵢ, ∇u[i])
    end
    return surr
end

function Base.empty!(surr::HermiteSplineSurrogate)
    empty!(surr.spl)
    return surr
end

function suggest_point(surr::HermiteSplineSurrogate{D,T}) where {D,T}
    _, I = findmin(evaluate!(vec(surr.ugrid), surr.spl, vec(surr.grid)))
    @inbounds _, p = nearest_interior_gridpoint(surr.grid, surr.grid[I])
    # @inbounds p = surr.grid[I] #TODO
    p, u = local_search(surr, p)
    return (p, u)
end

####
#### Bounding box for multi-dimensional bisection search
####

struct BoundingBox{D,S,N}
    bounds::NTuple{D,NTuple{2,Int}}
    corners::SArray{S,CartesianIndex{D},D,N}
end
corners(box::BoundingBox) = box.corners
bounds(box::BoundingBox) = box.bounds
widths(box::BoundingBox{D}) where {D} = ntuple(d -> abs(box.bounds[d][2] - box.bounds[d][1]), D)
Base.show(io::IO, ::MIME"text/plain", box::BoundingBox{D}) where {D} = print(io, "$D-D BoundingBox with dimensions: " * join(bounds(box), " × "))

BoundingBox(widths::NTuple{D,Int}) where {D} = BoundingBox(tuple.(1, widths))
BoundingBox(bounds::NTuple{D,NTuple{2,Int}}) where {D} = BoundingBox(bounds, corners(bounds))

@generated function corners(bounds::NTuple{D,NTuple{2,Int}}) where {D}
    corners = Iterators.product([(true, false) for d in 1:D]...)
    S = Tuple{ntuple(d -> 2, D)...}
    vals = [:(CartesianIndex($(ntuple(d -> I[d] ? :(bounds[$d][1]) : :(bounds[$d][2]), D)...))) for I in corners]
    :(Base.@_inline_meta; SArray{$S, CartesianIndex{$D}, $D, $(2^D)}(tuple($(vals...))))
end

function opposite_corner(box::BoundingBox{D}, I::CartesianIndex{D}) where {D}
    @inbounds lo, hi = first(corners(box)), last(corners(box))
    return lo + hi - I
end

function bisect(box::BoundingBox{D}) where {D}
    _, i = findmax(widths(box))
    left_bounds = ntuple(D) do d
        i !== d ? box.bounds[d] : (box.bounds[i][1], (box.bounds[i][1] + box.bounds[i][2]) ÷ 2)
    end
    right_bounds = ntuple(D) do d
        i !== d ? box.bounds[d] : ((box.bounds[i][1] + box.bounds[i][2]) ÷ 2, box.bounds[i][2])
    end
    return BoundingBox(left_bounds), BoundingBox(right_bounds)
end

splittable(box::BoundingBox{D}) where {D} = any(widths(box) .> 1)

####
#### Searching on a discrete grid using a surrogate function
####

struct DiscreteSurrogateSearcher{D, T}
    grid::Array{SVector{D,T},D}
    seen::Array{Bool,D}
    numeval::Base.RefValue{Int}
end
function DiscreteSurrogateSearcher(surr::AbstractSurrogate; mineval::Int, maxeval::Int)
    @assert mineval <= maxeval
    state = DiscreteSurrogateSearcher(surr.grid, fill(false, size(surr.grid)), Ref(0))
    initialize!(surr, state; mineval = mineval, maxeval = maxeval)
end

function initialize!(surr::AbstractSurrogate{D}, state::DiscreteSurrogateSearcher{D}; mineval::Int, maxeval::Int) where {D}
    # Evaluate at least `mineval` points by repeatedly bisecting the grid in a breadth-first manner
    box = BoundingBox(size(state.grid))
    for depth in 1:mineval # should never reach `mineval` depth, this is just to ensure the loop terminates in case `mineval` is greater than the number of gridpoints
        initialize!(surr, state, box, depth; mineval = mineval, maxeval = maxeval)
        state.numeval[] >= mineval && break
    end
    return state
end

function initialize!(surr::AbstractSurrogate{D}, state::DiscreteSurrogateSearcher{D}, box::BoundingBox{D}, depth::Int; mineval::Int, maxeval::Int) where {D}
    depth <= 0 && return state
    evaluate_box!(surr, state, box; maxeval = maxeval)
    state.numeval[] ≥ mineval && return state
    left, right = bisect(box)
    initialize!(surr, state, left, depth-1; mineval = mineval, maxeval = maxeval)
    initialize!(surr, state, right, depth-1; mineval = mineval, maxeval = maxeval)
end

function update!(surr::AbstractSurrogate{D}, state::DiscreteSurrogateSearcher{D}, I::CartesianIndex{D}; maxeval::Int) where {D}
    # Update the surrogate function with a new point, returning whether the maximum number of function evaluations has been reached or not
    state.numeval[] >= maxeval && return true # check if already exceeded number of evals
    @inbounds state.seen[I] && return false # point already evaluated
    update!(surr, I) # update surrogate
    @inbounds state.seen[I] = true # mark as now seen
    @inbounds state.numeval[] += 1 # increment function call counter
    return state.numeval[] >= maxeval
end

####
#### Global optimization using multi-dimensional bisection with surrogate functions
####

function bisection_search(
        surr::AbstractSurrogate{D,T},
        state::DiscreteSurrogateSearcher{D,T};
        maxeval::Int,
    ) where {D,T}

    # Algorithm:
    #   0. Get initial optimum suggestion from surrogate
    #   REPEAT:
    #       1. Find smallest bounding box containing the optimum suggestion
    #       2. Evaluate the box corners
    #       3. Get new optimum suggestion from surrogate:
    #           IF: Box is sufficiently small or if the maximum number of evaluations has been reached:
    #               RETURN: Current optimum
    #           ELSE:
    #               GOTO:   1.
    x, u = suggest_point(surr)
    while true
        box = minimal_bounding_box(state, x)
        evaluate_box!(surr, state, box, x; maxeval = maxeval)
        x, u = suggest_point(surr)
        if state.numeval[] ≥ maxeval || converged(state, box)
            return (x, u)
        end
    end
end

# Update observed evaluations, returning true if converged
function minimal_bounding_box(
        state::DiscreteSurrogateSearcher{D,T},
        x::SVector{D,T},
    ) where {D,T}

    box = BoundingBox(size(state.grid))
    while true
        left, right = bisect(box)
        if contains(state, left, x) # left box contains `x`
            if !is_evaluated(state, left) || !splittable(left)
                return left # left box not fully evaluated, or we have reached bottom; return
            else
                box = left # whole left box already evaluated; continue search
            end
        else # contains(state, right, x), i.e. right box contains `x`
            if !is_evaluated(state, right) || !splittable(right)
                return right # right box not fully evaluated, or we have reached bottom; return
            else
                box = right # whole right box already evaluated; continue search
            end
        end
    end
end

function evaluate_box!(
        surr::AbstractSurrogate{D,T},
        state::DiscreteSurrogateSearcher{D,T},
        box::BoundingBox{D},
        x::Union{Nothing, SVector{D,T}} = nothing;
        maxeval::Int,
    ) where {D,T}
    cs = x === nothing ? corners(box) : sorted_corners(state, box, x)
    @inbounds for I in cs
        is_evaluated(state, box) && break # box sufficiently evaluated
        update!(surr, state, I; maxeval = maxeval) && break # update surrogate, breaking if max evals reached
    end
    return state
end

function is_evaluated(state::DiscreteSurrogateSearcher{D}, box::BoundingBox{D}) where {D}
    # Box is considered sufficiently evaluated when all of the corners have been evaluted
    count(I -> @inbounds(state.seen[I]), corners(box)) >= 2^D
end

function converged(::DiscreteSurrogateSearcher{D}, box::BoundingBox{D}) where {D}
    # Convergence is defined as: bounding box has at least one side of length <= 1
    any(widths(box) .<= 1)
end

function centre(state::DiscreteSurrogateSearcher{D,T}, box::BoundingBox{D}) where {D,T}
    @inbounds lo = state.grid[first(corners(box))]
    @inbounds hi = state.grid[last(corners(box))]
    (lo + hi)/2
end

function sorted_corners(state::DiscreteSurrogateSearcher{D,T}, box::BoundingBox{D}, x::SVector{D,T}) where {D,T}
    dist²(I) = @inbounds sum(abs2.(state.grid[I] - x))
    cs = corners(box)
    typeof(cs)(TupleTools.sort(Tuple(cs); by = dist²))
end

function contains(state::DiscreteSurrogateSearcher{D,T}, box::BoundingBox{D}, x::SVector{D,T}) where {D,T}
    @inbounds lo = state.grid[first(corners(box))]
    @inbounds hi = state.grid[last(corners(box))]
    all(lo .<= x .<= hi)
end

function is_inside(grid::AbstractArray{SVector{D,T},D}, x::SVector{D,T}) where {D,T}
    @inbounds lo = first(grid)
    @inbounds hi = last(grid)
    all(lo .< x .< hi)
end
is_inside(state::DiscreteSurrogateSearcher{D,T}, x::SVector{D,T}) where {D,T} = is_inside(state.grid, x)

####
#### Local optimization using surrogate functions
####

function local_search(
        surr::HermiteSplineSurrogate{D,T},
        x₀::SVector{D,T},
        state::Union{Nothing, DiscreteSurrogateSearcher{D,T}} = nothing;
        maxiter::Int = 100,
        maxeval::Int = maxiter,
        xtol_rel = 1e-4,
        xtol_abs = 1e-4,
        initial_step = maximum(gridwidths(surr)) / 100,
        xeval_radius = √sum(abs2, gridspacings(surr)) - sqrt(eps(T)),
    ) where {D,T}

    if state !== nothing
        # Initialize surrogate with domain corners
        box = BoundingBox(size(surr.grid))
        for I in corners(box)
            update!(surr, state, I; maxeval = maxeval)
        end
    end

    x, xlast = x₀, x₀
    xlo, xhi = first(surr.grid), last(surr.grid)
    opt = ADAM{D,T}(initial_step)

    for _ in 1:maxiter
        if state !== nothing
            # Find nearest gridpoint to `x` and update surrogate
            I, xI = nearest_gridpoint(state, x)
            dmin = minimum(NormalHermiteSplines._get_nodes(surr.spl)) do p
                norm(xI - NormalHermiteSplines._unnormalize(surr.spl, p))
            end
            if dmin > xeval_radius
                update!(surr, state, I; maxeval = maxeval)
            end
        end

        # Perform gradient descent step using surrogate function
        ∇u = NormalHermiteSplines.evaluate_gradient(surr.spl, x)
        Δx, opt = update(∇u, opt)
        xlast, x = x, @. clamp(x - Δx, xlo, xhi)

        # Check for convergence
        maximum(abs.(x - xlast)) <= max(T(xtol_abs), T(xtol_rel) * maximum(abs.(x))) && break
    end

    u = NormalHermiteSplines.evaluate(surr.spl, x)

    return (x, u)
end

function nearest_gridpoint(grid::AbstractArray{SVector{D,T},D}, x::SVector{D,T}) where {D,T}
    @inbounds xlo, xhi = first(grid), last(grid)
    @inbounds Ilo, Ihi = first(CartesianIndices(grid)), last(CartesianIndices(grid))
    lo, hi = SVector(Tuple(Ilo)), SVector(Tuple(Ihi))
    i = @. clamp(round(Int, (x - xlo) * (hi - lo) / (xhi - xlo) + lo), lo, hi)
    I = CartesianIndex(Tuple(i))
    xI = @inbounds grid[I]
    return I, xI
end
nearest_gridpoint(state::DiscreteSurrogateSearcher{D,T}, x::SVector{D,T}) where {D,T} = nearest_gridpoint(state.grid, x)

function nearest_interior_gridpoint(grid::AbstractArray{SVector{D,T},D}, x::SVector{D,T}) where {D,T}
    R = CartesianIndices(grid)
    One = CartesianIndex(ntuple(d -> 1, D))
    Ilo, Ihi = first(R) + One, last(R) - One
    return nearest_gridpoint(@views(grid[Ilo:Ihi]), x)
end
nearest_interior_gridpoint(state::DiscreteSurrogateSearcher{D,T}, x::SVector{D,T}) where {D,T} = nearest_interior_gridpoint(state.grid, x)

####
#### Global optimization for NNLS problem
####

struct NNLSDiscreteSurrogateSearch{D, T, TA <: AbstractArray{T}, TdA <: AbstractArray{T}, Tb <: AbstractVector{T}, W}
    As::TA
    ∇As::TdA
    αs::Array{SVector{D,T},D}
    b::Tb
    u::Array{T,D}
    ∂Ax⁺::Vector{T}
    Ax⁺b::Vector{T}
    nnls_work::W
    legacy::Bool
end

function NNLSDiscreteSurrogateSearch(
        As::AbstractArray{T},  # size(As)  = (M, N, P1..., PD)
        ∇As::AbstractArray{T}, # size(∇As) = (M, N, D, P1..., PD)
        αs::NTuple{D},         # size(αs)  = (P1..., PD)
        b::AbstractVector{T};  # size(b)   = (M,)
        legacy::Bool = false,
    ) where {D,T}
    M, N = size(As, 1), size(As, 2)
    @assert ndims(As) == 2 + D && ndims(∇As) == 3 + D # ∇As has extra dimension for parameter gradients
    @assert size(∇As)[1:3] == (M, N, D) # matrix dimensions must match, and gradient dimension must equal number of parameters
    @assert size(As)[3:end] == size(∇As)[4:end] == length.(αs) # dimension size must match parameters lengths
    @assert size(b) == (M,)

    αs   = meshgrid(SVector{D,T}, αs...)
    u    = zeros(T, size(αs))
    ∂Ax⁺ = zeros(T, M)
    Ax⁺b = zeros(T, M)
    nnls_work = lsqnonneg_work(zeros(T, M, N), zeros(T, M))
    NNLSDiscreteSurrogateSearch(As, ∇As, αs, b, u, ∂Ax⁺, Ax⁺b, nnls_work, legacy)
end

load!(prob::NNLSDiscreteSurrogateSearch{D,T}, b::AbstractVector{T}) where {D,T} = copyto!(prob.b, b)

function loss!(prob::NNLSDiscreteSurrogateSearch{D,T}, I::CartesianIndex{D}) where {D,T}
    @unpack As, b, nnls_work = prob
    solve!(nnls_work, uview(As, :, :, I), b)
    ℓ = chi2(nnls_work)
    u = prob.legacy ? ℓ : log(max(ℓ, eps(T))) # loss capped at eps(T) from below to avoid log(0) error
    return u
end

function ∇loss!(prob::NNLSDiscreteSurrogateSearch{D,T}, I::CartesianIndex{D}) where {D,T}
    @unpack As, ∇As, b, ∂Ax⁺, Ax⁺b, nnls_work = prob
    ℓ = chi2(nnls_work)
    ℓ <= eps(T) && return zero(SVector{D,T}) # loss capped at eps(T) from below; return zero gradient
    x = solution(nnls_work)
    @inbounds @. Ax⁺b = -b
    @inbounds for j in 1:size(As, 2)
        (x[j] > 0) && axpy!(x[j], uview(As, :, j, I), Ax⁺b)
    end
    ∇u = ntuple(D) do d
        @inbounds ∂Ax⁺ .= zero(T)
        @inbounds for j in 1:size(∇As, 2)
            (x[j] > 0) && axpy!(x[j], uview(∇As, :, j, d, I), ∂Ax⁺)
        end
        ∂ℓ = 2 * dot(∂Ax⁺, Ax⁺b)
        ∂u = prob.legacy ? ∂ℓ : ∂ℓ / ℓ
        return ∂u
    end
    return SVector{D,T}(∇u)
end

function CubicSplineSurrogate(prehook!, prob::NNLSDiscreteSurrogateSearch{1,T}; legacy = false) where {T}
    function f(I)
        prehook!(I)
        loss!(prob, I)
    end
    CubicSplineSurrogate(f, prob.αs, SVector{1,T}[], T[], Ref(0), legacy)
end
CubicSplineSurrogate(prob::NNLSDiscreteSurrogateSearch; kwargs...) = CubicSplineSurrogate(I -> nothing, prob; kwargs...)

function HermiteSplineSurrogate(prehook!, prob::NNLSDiscreteSurrogateSearch{D,T}) where {D,T}
    function fg(I)
        prehook!(I)
        u = loss!(prob, I)
        ∇u = ∇loss!(prob, I)
        return u, ∇u
    end
    return HermiteSplineSurrogate(fg, prob.αs, RK_H1(one(T)))
end
HermiteSplineSurrogate(prob::NNLSDiscreteSurrogateSearch) = HermiteSplineSurrogate(I -> nothing, prob)

function surrogate_spline_opt(
        prob::NNLSDiscreteSurrogateSearch{D},
        surr::AbstractSurrogate{D};
        mineval::Int = min(2^D, length(prob.αs)),
        maxeval::Int = length(prob.αs),
    ) where {D}
    state = DiscreteSurrogateSearcher(surr; mineval = mineval, maxeval = maxeval)
    bisection_search(surr, state; maxeval = maxeval)
end

function spline_opt(
        spl::NormalSpline{D,T},
        prob::NNLSDiscreteSurrogateSearch{D,T};
        # alg = :LN_COBYLA,        # local, gradient-free, linear approximation of objective
        # alg = :LN_BOBYQA,        # local, gradient-free, quadratic approximation of objective
        # alg = :GN_ORIG_DIRECT_L, # global, gradient-free, systematically divides search space into smaller hyper-rectangles via a branch-and-bound technique, systematic division of the search domain into smaller and smaller hyperrectangles, "more biased towards local search"
        # alg = :GN_AGS,           # global, gradient-free, employs the Hilbert curve to reduce the source problem to the univariate one.
        # alg = :GD_STOGO,         # global, with-gradient, systematically divides search space into smaller hyper-rectangles via a branch-and-bound technique, and searching them by a gradient-based local-search algorithm (a BFGS variant)
        alg = :LD_SLSQP,         # local, with-gradient, "Sequential Least-Squares Quadratic Programming"; uses dense-matrix methods (ordinary BFGS, not low-storage BFGS)
    ) where {D,T}

    evaluate!(prob.u, spl, prob.αs)
    _, i = findmin(prob.u)
    α₀ = prob.αs[i]

    opt = NLopt.Opt(alg, D)
    opt.lower_bounds = Float64[prob.αs[begin]...]
    opt.upper_bounds = Float64[prob.αs[end]...]
    opt.xtol_rel = 0.001
    opt.min_objective = function (x, g)
        if length(g) > 0
            @inbounds g[1] = Float64(evaluate_derivative(spl, x[1]))
        end
        @inbounds Float64(evaluate(spl, x[1]))
    end
    minf, minx, ret = NLopt.optimize(opt, Vector{Float64}(α₀))
    return (SVector{D,T}(minx), T(minf))
end

function mock_surrogate_search_problem(
        b::AbstractVector,
        opts::T2mapOptions,
        ::Val{D},
        ::Val{ETL};
        alphas = range(50, 180, length = opts.nRefAngles),
        betas = range(50, 180, length = opts.nRefAngles),
    ) where {D,ETL}

    # Mock CPMG image
    @assert opts.nTE == ETL
    opt_vars = D == 1 ? (:α,) : (:α, :β)
    opt_ranges = D == 1 ? (alphas,) : (alphas, betas)
    As = zeros(ETL, opts.nT2, length.(opt_ranges)...)
    ∇As = zeros(ETL, opts.nT2, D, length.(opt_ranges)...)
    T2s = logrange(opts.T2Range..., opts.nT2)
    θ = EPGOptions((; α = opts.SetFlipAngle, TE = opts.TE, T2 = 0.0, T1 = opts.T1, β = opts.SetRefConAngle), Val(ETL), Float64)
    j! = EPGJacobianFunctor(θ, opt_vars)

    _, Rαs = SplitCartesianIndices(As, Val(2))
    for Iαs in Rαs
        @inbounds for j in 1:opts.nT2
            θαs = D == 1 ?
                restructure(θ, (T2 = T2s[j], α = alphas[Iαs[1]],)) :
                restructure(θ, (T2 = T2s[j], α = alphas[Iαs[1]], β = alphas[Iαs[2]]))
            j!(uview(∇As, :, j, :, Iαs), uview(As, :, j, Iαs), θαs)
        end
    end

    return NNLSDiscreteSurrogateSearch(As, ∇As, opt_ranges, b)
end
function mock_surrogate_search_problem(::Val{D}, ::Val{ETL}, opts = mock_t2map_opts(; MatrixSize = (1,1,1), nTE = ETL); kwargs...) where {D,ETL}
    b = vec(mock_image(opts))
    mock_surrogate_search_problem(b, opts, Val(D), Val(ETL); kwargs...)
end
function mock_surrogate_search_problem(b::AbstractVector, opts::T2mapOptions, ::Val{D}; kwargs...) where {D}
    @assert length(b) == opts.nTE
    mock_surrogate_search_problem(b, opts, Val(D), Val(length(b)); kwargs...)
end


text:
include("../required.jl")

outcome = "bart"
outcome = "raw"

files = [
    "fig/$(outcome)/04-3_$(outcome)_lcbd-transf.png"                 "fig/$(outcome)/04_$(outcome)_lcbd.png";
    "fig/$(outcome)/04-4_$(outcome)_relationship2d-transf.png"       "fig/$(outcome)/04_$(outcome)_relationship.png";
    "fig/$(outcome)/04-2_$(outcome)_richness.png"                    "fig/$(outcome)/04_$(outcome)_richness.png";
    "fig/$(outcome)/05-2_$(outcome)_subareas_3scales.png"            "fig/$(outcome)/05_$(outcome)_extents.png";
    "fig/$(outcome)/05-4_$(outcome)_subareas_medians.png"            "fig/$(outcome)/05_$(outcome)_medians.png";
    "fig/$(outcome)/05-3_$(outcome)_subareas.gif"                    "fig/$(outcome)/05_$(outcome)_scaling.gif";
    "fig/$(outcome)/05-1_$(outcome)_subareas_combined.png"           "fig/$(outcome)/05_$(outcome)_subareas.png";
    "fig/$(outcome)/07_$(outcome)_combined-maps.png"                 "fig/$(outcome)/09_$(outcome)_combined.png";
    "fig/$(outcome)/07_$(outcome)_comparison-combined.png"           "fig/$(outcome)/09_$(outcome)_difference.png";
    "fig/$(outcome)/07_$(outcome)_residuals-combined.png"            "fig/$(outcome)/09_$(outcome)_residuals.png";
    "fig/$(outcome)/08_$(outcome)_rare-species_eusrr_thresholds.png" "fig/$(outcome)/06_$(outcome)_eusrr.png";
    "fig/$(outcome)/08_$(outcome)_rare-species_ascending_plots.png"  "fig/$(outcome)/06_$(outcome)_rare-species.png";
    "fig/$(outcome)/08_$(outcome)_rare-species_thresholds.png"       "fig/$(outcome)/06_$(outcome)_thresholds.png";
    "fig/$(outcome)/06-0_$(outcome)_moving-windows_full.png"          "archive/fig/$(outcome)/06-0_$(outcome)_moving-windows_full.png";
    "fig/$(outcome)/06-1_$(outcome)_moving-windows_subareas.png"      "archive/fig/$(outcome)/06-1_$(outcome)_moving-windows_subareas.png";
    "fig/$(outcome)/06-2_$(outcome)_moving-windows_3scales.png"       "archive/fig/$(outcome)/06-2_$(outcome)_moving-windows_3scales.png";
    "fig/$(outcome)/06-3_$(outcome)_moving-windows.gif"               "archive/fig/$(outcome)/06-3_$(outcome)_moving-windows.gif";
    "fig/$(outcome)/08_$(outcome)_rare-species_spatial_subareas.png"  "archive/fig/$(outcome)/08_$(outcome)_rare-species_spatial_subareas.png";
]

oldfiles = files[:, 1]
newfiles = files[:, 2]

isfile.(oldfiles)
filter(!isfile, oldfiles)

for (old, new) in zip(oldfiles, newfiles)
    if isfile(old)
        mv(old, new)
    end
end

isfile.(newfiles)
filter(!isfile, newfiles)

oldfiles = replace.(oldfiles, "fig/$(outcome)/" => "")
newfiles = replace.(newfiles, "archive/fig/$(outcome)/" => "")
newfiles = replace.(newfiles, "fig/$(outcome)/" => "")

file = "src/04_full-extent.jl"
(tmppath, tmpio) = mktemp()
cp(file, tmppath; force=true)
for (old, new) in zip(oldfiles[1:3], newfiles[1:3])
    open(tmppath) do io
        for line in eachline(io, keep=true) # keep so the new line isn't chomped
            line = replace(line, old => new)
            write(tmpio, line)
        end
    end
    close(tmpio)
end
mv(tmppath, file, force=true)

max_stars_repo_name: gabrieldansereau/betadiversity-hotspots

max_issues_repo_name: gabrieldansereau/betadiversity-hotspots

max_forks_repo_name: gabrieldansereau/betadiversity-hotspots


text:
using Knet, StatsBase, LinearAlgebra, Random

include("data.jl")
include("model.jl")

struct SimpleLSTMModel
    embed::Embed
    rnn::RNN
    projection::Linear
    dropout::Real
    vocab::Vocab
end

"""

    SimpleLSTMModel(embsz::Int, hidden::Int, vocab::Vocab; layers=1, dropout=0)

```
embsz: Embeddings size
hidden: Size of LSTM hidden layer
vocab: Vocab set of type ::Vocab
layers=1: number of LSTM layers
dropout=0: dropout value
```

Returns:

    Language model consisting of LSTM with the given parameters
"""
function SimpleLSTMModel(embsz::Int, hidden::Int, vocab::Vocab; layers=1, dropout=0)

    embed = Embed(length(vocab.i2v), embsz)
    rnn = RNN(embsz, hidden; bidirectional=false, numLayers=layers, dropout=dropout, atype=aType)
    projection = Linear(hidden, length(vocab.i2v))
    rnn.h, rnn.c = 0, 0

    SimpleLSTMModel(embed, rnn, projection, dropout, vocab)
end

function (s::SimpleLSTMModel)(src, tgt; average=true)
    s.rnn.h, s.rnn.c = value(s.rnn.h), value(s.rnn.c)
    embed = s.embed(src)
    rnn_out = s.rnn(embed)
    dims = size(rnn_out)
    output = s.projection(dropout(reshape(rnn_out, dims[1], dims[2] * dims[3]), s.dropout))
    scores = reshape(output, size(output, 1), dims[2], dims[3])
    nll(scores, tgt; dims=1, average=average)
end

"""
    Generating words using the LM with sampling
"""
function generate(s::SimpleLSTMModel; start="", del="", maxlength=30)
    s.rnn.h, s.rnn.c = 0, 0
    vocabs = fill(s.vocab.eos, 1)

    starting_index = 1
    for (i, t) in enumerate(start)
        push!(vocabs, s.vocab.v2i[string(UInt8(t))])
        embed = s.embed(vocabs[i:i])
        rnn_out = s.rnn(embed)
        starting_index += 1
    end

    for i in starting_index:maxlength
        embed = s.embed(vocabs[i:i])
        rnn_out = s.rnn(embed)
        output = s.projection(dropout(rnn_out, s.dropout))
        push!(vocabs, s.vocab.v2i[ sample(s.vocab.i2v, Weights(Array(softmax(reshape(output, length(s.vocab.i2v)))))) ] )

        if vocabs[end] == s.vocab.eos
            break
        end
    end
    join([ Char(parse(UInt8, s.vocab.i2v[i])) for i in vocabs[2:end-1] ], del)
end

nothing


text:
"""
# Provide implementation of default layers

Provided layers
- DenseLayer
- DenseNoBiasLayer
- VectorFunctionLayer
"""

#using Random, Zygote
#import ..Utils
#import Base.size

# ------------------------------------------------------------------------------
# DenseLayer layer

"""
   DenseLayer

Representation of a layer in the network

# Fields:
* `w`:  Weigths matrix with respect to the input from previous layer or data (n x n pr. layer)
* `wb`: Biases (n)
* `f`:  Activation function
* `df`: Derivative of the activation function
"""
mutable struct DenseLayer <: AbstractLayer
     w::Array{Float64,2}
     wb::Array{Float64,1}
     f::Function
     df::Union{Function,Nothing}
     """
        DenseLayer(f,n,nₗ;w,wb,df)

     Instantiate a new DenseLayer

     # Positional arguments:
     * `nₗ`: Number of nodes of the previous layer
     * `n`:  Number of nodes
     # Keyword arguments:
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (n,nₗ)]
     * `wb`:  Initial weigths with respect to bias [default: Xavier initialisation, dims = (n)]
     * `f`:   Activation function [def: `identity`]
     * `df`:  Derivative of the activation function [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n))`

     """
     function DenseLayer(nₗ,n;rng = Random.GLOBAL_RNG,w=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),wb=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n),f=identity,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         w  = reshape(w,n,nₗ)
         wb = reshape(wb,n)
         return new(w,wb,f,df)
     end
end

function _zComp(layer::DenseLayer,x)
    w  = layer.w
    wb = layer.wb
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        zn   += wb[n]
        z[n]  = zn
    end
    return z
end

function forward(layer::DenseLayer,x)
  z =  _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) #layer.w * x + layer.wb # _zComp(layer,x) #   layer.w * x + layer.wb # testd @avx
  return layer.f.(z)
end

function backward(layer::DenseLayer,x,nextGradient)
   z = _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) # layer.w * x + layer.wb # _zComp(layer,x) # @avx layer.w * x + layer.wb               # tested @avx
   if layer.df != nothing
       dϵ_dz =  layer.df.(z) .* nextGradient # tested @avx
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI =  layer.w' * dϵ_dz # @avx
end

function getParams(layer::DenseLayer)
  return Learnable((layer.w,layer.wb))
end

function getGradient(layer::DenseLayer,x,nextGradient)
   z      =  _zComp(layer,x) #@avx layer.w * x + layer.wb #  _zComp(layer,x) #layer.w * x + layer.wb # @avx
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  = dϵ_dz * x' # @avx
   dϵ_dwb = dϵ_dz
   return Learnable((dϵ_dw,dϵ_dwb))
end

function setParams!(layer::DenseLayer,w)
   layer.w  = w.data[1]
   layer.wb = w.data[2]
end
function size(layer::DenseLayer)
    return size(layer.w')
end

# ------------------------------------------------------------------------------
# DenseNoBiasLayer layer

"""
   DenseNoBiasLayer

Representation of a layer without bias in the network

# Fields:
* `w`:  Weigths matrix with respect to the input from previous layer or data (n x n pr. layer)
* `f`:  Activation function
* `df`: Derivative of the activation function
"""
mutable struct DenseNoBiasLayer <: AbstractLayer
     w::Array{Float64,2}
     f::Function
     df::Union{Function,Nothing}
     """
        DenseNoBiasLayer(f,nₗ,n;w,df)

     Instantiate a new DenseNoBiasLayer

        DenseNoBiasLayer(f,nₗ,n;w,df)

     Instantiate a new DenseNoBiasLayer

     # Positional arguments:
     * `nₗ`:  Number of nodes of the previous layer
     * `n`:   Number of nodes
     # Keyword arguments:
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:   Activation function [def: `identity`]
     * `df`:  Derivative of the activation function [def: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ,n))`
     """
     function DenseNoBiasLayer(nₗ,n;rng = Random.GLOBAL_RNG,w=rand(rng,Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),f=identity,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         w  = reshape(w,n,nₗ)
         return new(w,f,df)
     end
end

function _zComp(layer::DenseNoBiasLayer,x)
    w  = layer.w
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        z[n]  = zn
    end
    return z
end

function forward(layer::DenseNoBiasLayer,x)
  z = _zComp(layer,x)
  return layer.f.(z)
end

function backward(layer::DenseNoBiasLayer,x,nextGradient)
   z = _zComp(layer,x)
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI = layer.w' * dϵ_dz
end

function getParams(layer::DenseNoBiasLayer)
  return Learnable((layer.w,))
end

function getGradient(layer::DenseNoBiasLayer,x,nextGradient)
   z      = _zComp(layer,x)
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  =  dϵ_dz * x'
   return Learnable((dϵ_dw,))
end

function setParams!(layer::DenseNoBiasLayer,w)
   layer.w = w.data[1]
end

function size(layer::DenseNoBiasLayer)
    return size(layer.w')
end

# ------------------------------------------------------------------------------
# VectorFunctionLayer layer

"""
   VectorFunctionLayer

Representation of a VectorFunction layer in the network. Vector
function layer expects a vector activation function, i.e. a function taking the
whole output of the previous layer an input rather than working on a single node
as "normal" activation functions would do.
Useful for example with the SoftMax function in classification or with the
`pool1D` function to implement a "pool" layer in 1 dimensions.
By default it is weightless, i.e. it doesn't apply any transformation to the
output coming from the previous layer except the activation function. However,
by passing the parameter `wsize` (a touple or array - tested only 1D) you can
pass the learnable parameter to the activation function too. It is your
responsability to be sure the activation function accept only X or also this
learnable array (as second argument).
The number of nodes in input must be set to the same as in the previous layer
(and if you are using this for classification, to the number of classes, i.e.
the _previous_ layer must be set equal to the number of classes in the
predictions).

# Fields:
* `w`:   Weigths (parameter) array passes as second argument to the activation
         function (if not empty)
* `nₗ`:  Number of nodes in input (i.e. length of previous layer)
* `n`:   Number of nodes in output (automatically inferred in the constructor)
* `f`:   Activation function (vector)
* `dfx`: Derivative of the (vector) activation function with respect to the
         layer inputs (x)
* `dfw`: Derivative of the (vector) activation function with respect to the
         optional learnable weigths (w)

# Notes:
* The output `size` of this layer is given by the size of the output function,
that not necessarily is the same as the previous layers.
"""
mutable struct VectorFunctionLayer{N} <: AbstractLayer
     w::Array{Float64,N}
     nₗ::Int64
     n::Int64
     f::Function
     dfx::Union{Function,Nothing}
     dfw::Union{Function,Nothing}
     """
        VectorFunctionLayer(n;rng,wsize,w,f,dfx,dfw,dummyDataToTestOutputSize)

     Instantiate a new VectorFunctionLayer

     # Positional arguments:
     * `nₗ`: Number of nodes (must be same as in the previous layer)
     # Keyword arguments:
     * `wsize`: A tuple or array specifying the size (number of elements) of the
       learnable parameter [def: empty array]
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:  Activation function [def: `softMax`]
     * `dfx`: Derivative of the activation function with respect to the data
              [default: `nothing` (i.e. use AD)]
     * `dfw`: Derivative of the activation function with respect to the
              learnable parameter [default: `nothing` (i.e. use AD)]
     * `dummyDataToTestOutputSize`: Dummy data to test the output size [def:
     `ones(nₗ)`]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - If the derivative is provided, it should return the gradient as a (n,n) matrix (i.e. the Jacobian)
     - To avoid recomputing the activation function just to determine its output size,
       we compute the output size once here in the layer constructor by calling the
       activation function with `dummyDataToTestOutputSize`. Feel free to change
       it if it doesn't match with the activation function you are setting
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(sum(wsize...)),sqrt(6)/sqrt(sum(wsize...))))`
     """
     function VectorFunctionLayer(nₗ;rng = Random.GLOBAL_RNG,wsize=Int64[],w=rand(rng,Uniform(-sqrt(6)/sqrt(sum(wsize)),sqrt(6)/sqrt(sum(wsize))),Tuple(wsize)), f=softMax,dfx=nothing,dfw=nothing,dummyDataToTestOutputSize=ones(nₗ))
        nw = length(wsize)
        if nw ==0
          n  = length(f(dummyDataToTestOutputSize))
        else
         n  = length(f(dummyDataToTestOutputSize,w))
        end
        return new{nw}(w,nₗ,n,f,dfx,dfw)
     end
end

function forward(layer::VectorFunctionLayer{N},x) where {N}
  return N == 0 ? layer.f(x) : layer.f(x,layer.w)
end

function backward(layer::VectorFunctionLayer{N},x,nextGradient) where N
   if N == 0
      if layer.dfx != nothing
         dϵ_dI = layer.dfx(x)' * nextGradient
      else  # using AD
         j = autoJacobian(layer.f,x; nY=layer.n)
         dϵ_dI = j' * nextGradient
      end
      return dϵ_dI
   else
      if layer.dfx != nothing
         dϵ_dI = layer.dfx(x,layer.w)' * nextGradient
      else  # using AD
        tempfunction(x) = layer.f(x,layer.w)
        nYl = layer.n
        j = autoJacobian(tempfunction,x; nY=nYl)
        dϵ_dI = j' * nextGradient
      end
     return dϵ_dI
   end
end

function getParams(layer::VectorFunctionLayer{N}) where {N}
   return N == 0 ? Learnable(()) : Learnable((layer.w,))
end

function getGradient(layer::VectorFunctionLayer{N},x,nextGradient) where {N}
   if N == 0
     return Learnable(()) # parameterless layer
   else
      if layer.dfw != nothing
         dϵ_dw = layer.dfw(x,layer.w)' * nextGradient
      else  # using AD
        j = autoJacobian(wt -> layer.f(x,wt),layer.w; nY=layer.n)
        dϵ_dw = j' * nextGradient
      end
     return Learnable((dϵ_dw,))
   end
end

function setParams!(layer::VectorFunctionLayer{N},w) where {N}
   if N > 0
      layer.w = w.data[1]
   end
   return nothing
end
function size(layer::VectorFunctionLayer{N}) where {N}
    # Output size for the VectorFunctionLayer is given by its activation function
    # We test its length with dummy values
    return (layer.nₗ,layer.n)
end

# ------------------------------------------------------------------------------
# ScalarFunctionLayer layer

"""
   ScalarFunctionLayer

Representation of a ScalarFunction layer in the network.
ScalarFunctionLayer applies the activation function directly to the output of
the previous layer (i.e., without passing for a weigth matrix), but using an
optional learnable parameter (an array) used as second argument, similarly to
[`VectorFunctionLayer`(@ref).
Differently from `VectorFunctionLayer`, the function is applied scalarwise to
each node.

The number of nodes in input must be set to the same as in the previous layer

# Fields:
* `w`:   Weigths (parameter) array passes as second argument to the activation
         function (if not empty)
* `n`:   Number of nodes in output (≡ number of nodes in input )
* `f`:   Activation function (vector)
* `dfx`: Derivative of the (vector) activation function with respect to the
         layer inputs (x)
* `dfw`: Derivative of the (vector) activation function with respect to the
         optional learnable weigths (w)

# Notes:
* The output `size` of this layer is the same as those of the previous layers.
"""
mutable struct ScalarFunctionLayer{N} <: AbstractLayer
     w::Array{Float64,N}
     n::Int64
     f::Function
     dfx::Union{Function,Nothing}
     dfw::Union{Function,Nothing}
     """
        ScalarFunctionLayer(n;rng,wsize,w,f,dfx,dfw)

     Instantiate a new ScalarFunctionLayer

     # Positional arguments:
     * `nₗ`: Number of nodes (must be same as in the previous layer)
     # Keyword arguments:
     * `wsize`: A tuple or array specifying the size (number of elements) of the
       learnable parameter [def: empty array]
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:  Activation function [def: `softMax`]
     * `dfx`: Derivative of the activation function with respect to the data
              [default: `nothing` (i.e. use AD)]
     * `dfw`: Derivative of the activation function with respect to the
              learnable parameter [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - If the derivative is provided, it should return the gradient as a (n,n) matrix (i.e. the Jacobian)
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(sum(wsize...)),sqrt(6)/sqrt(sum(wsize...))))`
     """
     function ScalarFunctionLayer(nₗ;rng = Random.GLOBAL_RNG,wsize=Int64[],w=rand(rng,Uniform(-sqrt(6)/sqrt(sum(wsize)),sqrt(6)/sqrt(sum(wsize))),Tuple(wsize)), f=softMax,dfx=nothing,dfw=nothing)
        nw = length(wsize)
        return new{nw}(w,nₗ,f,dfx,dfw)
     end
end

function forward(layer::ScalarFunctionLayer{N},x) where {N}
  return N == 0 ? layer.f.(x) : layer.f.(x,layer.w)
end

function backward(layer::ScalarFunctionLayer{N},x,nextGradient) where N
   if N == 0
      if layer.dfx != nothing
         dϵ_dI = layer.dfx.(x) .* nextGradient
      else  # using AD
         dϵ_dI = layer.f'.(x) .* nextGradient
      end
      return dϵ_dI
   else
      if layer.dfx != nothing
         dϵ_dI = layer.dfx.(x,Ref(layer.w)) .* nextGradient
      else  # using AD
        #tempfunction(x) = layer.f.(x,Ref(layer.w))
        df_dx = [gradient(xt -> layer.f(xt,layer.w),xi)[1] for xi in x]
        dϵ_dI = df_dx .* nextGradient
      end
     return dϵ_dI
   end
end

function getParams(layer::ScalarFunctionLayer{N}) where {N}
   return N == 0 ? Learnable(()) : Learnable((layer.w,))
end

function getGradient(layer::ScalarFunctionLayer{N},x,nextGradient) where {N}
   if N == 0
     return Learnable(()) # parameterless layer
   else
      if layer.dfw != nothing
         dϵ_dw = [layer.dfw(xi,wj) for wj in layer.w, xi in x] * nextGradient
      else  # using AD
        tempfunction(w) = [layer.f(xi,w) for xi in x]
        j = autoJacobian(tempfunction,layer.w; nY=layer.n)
        dϵ_dw = j' * nextGradient
      end
     return Learnable((dϵ_dw,))
   end
end

function setParams!(layer::ScalarFunctionLayer{N},w) where {N}
   if N > 0
      layer.w = w.data[1]
   end
   return nothing
end
function size(layer::ScalarFunctionLayer{N}) where {N}
    # Output size for the ScalarFunctionLayer is given by its activation function
    # We test its length with dummy values
    return (layer.n,layer.n)
end

# ------------------------------------------------------------------------------
# RNN layer

"""
   RNNLayer

Representation of a layer in the network

# Fields:
* `wx`: Weigths matrix with respect to the input from data (n by n_input)
* `ws`: Weigths matrix with respect to the layer state (n x n )
* `wb`: Biases (n)
* `f`:  Activation function
* `df`: Derivative of the activation function
* `s` : State
"""
mutable struct RNNLayer <: RecursiveLayer
     wx::Array{Float64,2}
     ws::Array{Float64,2}
     wb::Array{Float64,1}
     s::Array{Float64,1}
     f::Function
     df::Union{Function,Nothing}
     """
        RNNLayer(nₗ,n;f,wx,ws,wb,df)

     Instantiate a new RNNLayer

     # Positional arguments:
     * `nₗ`: Number of nodes of the input
     * `n`:  Number of nodes of the state (and the output)
     # Keyword arguments:
     * `wx`:  Initial weigths with respect to input [default: Xavier initialisation, dims = (n,nₗ)]
     * `ws`:  Initial weigths with respect to input [default: Xavier initialisation, dims = (n,n)]
     * `wb`:  Initial weigths with respect to bias [default: Xavier initialisation, dims = (n)]
     * `s`:   Initial states [def: zeros(n)]
     * `f`:   Activation function [def: `relu`]
     * `df`:  Derivative of the activation function [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n))`

     """
     function RNNLayer(nₗ,n;rng = Random.GLOBAL_RNG,wx=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),ws=rand(rng, Uniform(-sqrt(6)/sqrt(n+n),sqrt(6)/sqrt(n+n)),n,n),wb=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n),s=zeros(n),f=relu,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         wx  = reshape(w,n,nₗ)
         ws  = reshape(w,n,n)
         wb  = reshape(wb,n)
         s   = reshape(s,n)
         return new(w,wb,s,f,df)
     end
end
#=
function _zComp(layer::RNNLayer,x)
    w  = layer.w
    wb = layer.wb
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        zn   += wb[n]
        z[n]  = zn
    end
    return z
end
=#

function forward(layer::RNNLayer,x)
  #z =  _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) #layer.w * x + layer.wb # _zComp(layer,x) #   layer.w * x + layer.wb # testd @avx
  z =  layer.wb + layer.wx * x + layer.ws * s
  return layer.f.(z)
end

function backward(layer::RNNLayer,x,nextGradient) #TODO
   z = _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) # layer.w * x + layer.wb # _zComp(layer,x) # @avx layer.w * x + layer.wb               # tested @avx
   if layer.df != nothing
       dϵ_dz =  layer.df.(z) .* nextGradient # tested @avx
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI =  layer.w' * dϵ_dz # @avx
end

function getParams(layer::RNNLayer)
  return Learnable((layer.wb,layer.wx,layer.ws))
end

function getGradient(layer::RNNLayer,x,nextGradient) #TODO
   z      =  _zComp(layer,x) #@avx layer.w * x + layer.wb #  _zComp(layer,x) #layer.w * x + layer.wb # @avx
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  = dϵ_dz * x' # @avx
   dϵ_dwb = dϵ_dz
   return Learnable((dϵ_dw,dϵ_dwb))
end

function setParams!(layer::RNNLayer,w)
   layer.wb = w.data[1]
   layer.wx = w.data[2]
   layer.ws = w.data[3]
end
function size(layer::RNNLayer)
    return size(layer.w')
end


text:
using JuMP
using ConditionalJuMP

export Conv2d

"""
$(TYPEDEF)

Represents 2-D convolution operation.

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

## Fields:
$(FIELDS)
"""
@auto_hash_equals struct Conv2d{T<:JuMPReal, U<:JuMPReal, V<:Int64} <: Layer
    filter::Array{T, 4}
    bias::Array{U, 1}
    stride::V

    function Conv2d{T, U, V}(filter::Array{T, 4}, bias::Array{U, 1}, stride::V) where {T<:JuMPReal, U<:JuMPReal, V<:Int64}
        (filter_height, filter_width, filter_in_channels, filter_out_channels) = size(filter)
        bias_out_channels = length(bias)
        @assert(
            filter_out_channels == bias_out_channels,
            "For this convolution layer, number of output channels in filter, $filter_out_channels, does not match number of output channels in bias, $bias_out_channels."
        )
        return new(filter, bias, stride)
    end

end

function Conv2d(filter::Array{T, 4}, bias::Array{U, 1}, stride::V) where {T<:JuMPReal, U<:JuMPReal, V<:Int64}
    Conv2d{T, U, V}(filter, bias, stride)
end

function Conv2d(filter::Array{T, 4}, bias::Array{U, 1}) where {T<:JuMPReal, U<:JuMPReal}
    Conv2d(filter, bias, 1)
end

"""
$(SIGNATURES)

Convenience function to create a [`Conv2d`](@ref) struct with the specified filter
and zero bias.
"""
function Conv2d(filter::Array{T, 4}) where {T<:JuMPReal}
    bias_out_channels::Int = size(filter)[4]
    bias = zeros(bias_out_channels)
    Conv2d(filter, bias)
end

function check_size(params::Conv2d, sizes::NTuple{4, Int})::Void
    check_size(params.filter, sizes)
    check_size(params.bias, (sizes[end], ))
end

function Base.show(io::IO, p::Conv2d)
    (filter_height, filter_width, filter_in_channels, filter_out_channels) = size(p.filter)
    stride = p.stride
    print(io,
        "Conv2d($filter_in_channels, $filter_out_channels, kernel_size=($(filter_height), $(filter_width)), stride=($(stride), $(stride)), padding=same)"
    )
end

function increment!(s::Real, input_val::Real, filter_val::Real)
    return s + input_val*filter_val
end

function increment!(s::JuMP.AffExpr, input_val::JuMP.AffExpr, filter_val::Real)
    append!(s, input_val*filter_val)
    return s
end

function increment!(s::JuMP.AffExpr, input_val::JuMP.Variable, filter_val::Real)
    push!(s, Float64(filter_val), input_val)
end

function increment!(s::JuMP.AffExpr, input_val::Real, filter_val::JuMP.AffExpr)
    append!(s, filter_val*input_val)
    return s
end

function increment!(s::JuMP.AffExpr, input_val::Real, filter_val::JuMP.Variable)
    push!(s, Float64(input_val), filter_val)
end

"""
$(SIGNATURES)

Computes the result of convolving `input` with the `filter` and `bias` stored in `params`.

Mirrors `tf.nn.conv2d` from the `tensorflow` package, with `strides = [1, 1, 1, 1],
padding = 'SAME'`.

# Throws
* AssertionError if `input` and `filter` are not compatible.
"""
function conv2d(
    input::Array{T, 4},
    params::Conv2d{U, V}) where {T<:JuMPReal, U<:JuMPReal, V<:JuMPReal}

    if T<:JuMPLinearType || U<:JuMPLinearType || V<:JuMPLinearType
        info(MIPVerify.LOGGER, "Applying $(params) ... ")
    end
    filter = params.filter
    stride = params.stride

    (batch, in_height, in_width, input_in_channels) = size(input)
    (filter_height, filter_width, filter_in_channels, filter_out_channels) = size(filter)

    @assert(
        input_in_channels == filter_in_channels,
        "Number of channels in input, $input_in_channels, does not match number of channels, $filter_in_channels, that filters operate on."
    )

    out_height = round(Int, in_height/stride, RoundUp)
    out_width = round(Int, in_width/stride, RoundUp)
    output_size = (batch, out_height, out_width, filter_out_channels)

    # Considered using offset arrays here, but could not get it working.

    # Calculating appropriate offsets so that center of kernel is matched with
    # cell at which correlation is being calculated. Note that tensorflow
    # chooses a specific convention for a dimension with even size which we
    # replicate here.
    filter_height_offset = round(Int, filter_height/2, RoundUp)
    filter_width_offset = round(Int, filter_width/2, RoundUp)
    W = Base.promote_op(+, V, Base.promote_op(*, T, U))
    output = Array{W}(output_size)

    @nloops 4 i output begin
        s::W = 0
        @nloops 4 j filter begin
            if i_4 == j_4
                x = (i_2-1)*stride+1 + j_1 - filter_height_offset
                y = (i_3-1)*stride+1 + j_2 - filter_width_offset
                if x > 0 && y > 0 && x<=in_height && y<=in_width
                    # Doing bounds check to make sure that we stay within bounds
                    # for input. This effectively zero-pads the input.
                    # TODO (vtjeng): Use default checkbounds function here instead?
                    s = increment!(s, input[i_1, x, y, j_3], filter[j_1, j_2, j_3, j_4])
                end
            end
        end
        s += params.bias[i_4]
        if T<:JuMPLinearType
            ConditionalJuMP.simplify!(s)
        end
        (@nref 4 output i) = s
    end

    return output
end

(p::Conv2d)(x::Array{<:JuMPReal, 4}) = conv2d(x, p)


text:
module NeuralOperators
    using Flux
    using FFTW
    using Tullio
    using CUDA
    using CUDAKernels
    using KernelAbstractions
    using Zygote
    using ChainRulesCore

    export DeepONet

    include("Transform/Transform.jl")
    include("operator_kernel.jl")
    include("model.jl")
    include("DeepONet.jl")
    include("subnets.jl")
end


text:
struct OCTHaGOnException <: Exception
    var::String
end

Base.showerror(io::IO, e::OCTHaGOnException) = print(io, e.var)


text:
"""
    explore(domains, concept, param = nothing; search_limit = 1000, solutions_limit = 100)

Search (a part of) a search space and returns a pair of vector of configurations: `(solutions, non_solutions)`. If the search space size is over `search_limit`, then both `solutions` and `non_solutions` are limited to `solutions_limit`.

Beware that if the density of the solutions in the search space is low, `solutions_limit` needs to be reduced. This process will be automatic in the future (simple reinforcement learning).

# Arguments:
- `domains`: a collection of domains
- `concept`: the concept of the targeted constraint
- `param`: an optional parameter of the constraint
- `sol_number`: the required number of solutions (half of the number of configurations), default to `100`
"""
function explore(
    domains,
    concept,
    param=nothing;
    search=:flexible,
    complete_search_limit=10^6,
    max_samplings=sum(domain_size, domains),
    solutions_limit=floor(Int, sqrt(max_samplings)),
)
    if search == :flexible
        search = sum(domain_size, domains) < complete_search_limit ? :complete : :partial
    end
    return explore(Val(search), domains, concept, param, solutions_limit, max_samplings)
end

function explore(::Val{:partial}, domains, concept, param, solutions_limit, max_samplings)
    solutions = Set{Vector{Int}}()
    non_sltns = Set{Vector{Int}}()

    f = isnothing(param) ? ((x; param = p) -> concept(x)) : concept

    for _ in 1:max_samplings
        length(solutions) ≥ solutions_limit && length(non_sltns) ≥ solutions_limit && break
        config = map(rand, domains)
        c = f(config; param) ? solutions : non_sltns
        length(c) < solutions_limit && push!(c, config)
    end
    return solutions, non_sltns
end

function explore(::Val{:complete}, domains, concept, param, ::Int, ::Int)
    solutions = Set{Vector{Int}}()
    non_sltns = Set{Vector{Int}}()

    f = isnothing(param) ? ((x; param = p) -> concept(x)) : concept

    configurations = Base.Iterators.product(map(d -> get_domain(d), domains)...)
    foreach(
        c -> (cv = collect(c); push!(f(cv; param) ? solutions : non_sltns, cv)),
        configurations,
    )
    return solutions, non_sltns
end