text:
using OndaEDF: validate_arrow_prefix, prettyprint_diagnostic_info, mock_edf

function test_preprocessor(l, t)
    l = replace(l, ':' => '-')
    l = replace(l, "\xf6"[1] => 'o') # remove umlaut (German)
    l = replace(l, '\u00F3' => 'o') # remove accute accent (Spanish)
    l = replace(l, '\u00D3' => 'O') # remove accute accent (Spanish)
    l = OndaEDF._safe_lowercase(l)
    t = OndaEDF._safe_lowercase(t)

    l = replace(l, r"\.$"i => "")

    ####
    # ODDBALLS

    m = match(r"\[chin-(?<lr>[lr])-chin-[ar]\]"i, l)
    l = isnothing(m) ? l : "emg chin$(m[:lr])"

    l = endswith(l, "clavicle") ? "ecg avl" : l

    l = l == "ekg8" ? "ekg v8" : l
    l = replace(l, r"^e1 14" => "e1")
    l = replace(l, r"^e2 18" => "e2")

    l = endswith(l, "cardiogr") ? "ecg avl" : l

    l = l == "ecg+" ? "ecg" : l
    l = l == "ecg-" ? "ecg" : l

    if t == "ekg_channel"
        l = "ecg"
    end

    if l == "eog" && t == ""
        l = "eog l"  # recording will have 2 left eogs, no way to say which is which
    end

    l = startswith(l, "eog horz") ? "eog l" : l

    # CHIN1I-CHIN[23][]
    l = replace(l, r"^chin1i-chin.*"i => "emg chin1")

    if t == "leg_channel" && l == "25"
        l = "left_anterior_tibialis"
    end

    if t == "leg2_channel" && l == "26"
        l = "right_anterior_tibialis"
    end

    if t == "chin_channel" && l ∈ ["fz-cz", "24", "p3-p4", "fp1-fp2"]
        l = "emg chin1"
    end

    if t == "chin2_channel" && l ∈ ["25", "p4-pz"]
        l = "emg chin2"
    end

    m = match(r"chin(?<lr>[lr])[\(\s]+ment.*"i, l)
    l = isnothing(m) ? l : "emg chin$(m[:lr])"

    l = l == "emg x1-x6" ? "emg chin1" : l

    m = match(r"emg\s+(?<lr>[lr])at\s+.*"i, l)
    l = isnothing(m) ? l : "emg $(m[:lr])at"

    m = match(r"[lr]leg[\+\-]"i, l)
    l = isnothing(m) ? l : "ignorame"

    m = match(r"^(?<lr>[lr])t\. eye"i, l)
    l = isnothing(m) ? l : "eog $(m[:lr])"

    l = l == "e1 (l)-m2" ? "eog l" : l
    l = l == "e2 (r)-m2" ? "eog r" : l

    l = l == "chinz-chin1" ? "emg chinl" : l
    l = l == "chinz-chin2" ? "emg chinr" : l

    l = l == "eog x4-a2" ? "eog r" : l
    l = l == "eog x9-a2" ? "eog l" : l

    #m = match(r"^eeg (?<lr>[lr]oc)(?<rest>.*)$"i, l)
    #l = isnothing(m) ? l : "eog $(m[:lr])$(m[:rest])"

    l = l == "e.left" ? "eog loc" : l
    l = l == "e.right" ? "eog roc" : l

    ###############
    # AMBIGUOUS ONES
    # these will only be kept if recording also has leg EMG in separate channels

    # "EMG[123]" => emg_ambiguous [123]
    m = match(r"^\s*emg_?(?<n>[lr123])\s*$"i, l)
    l = isnothing(m) ? l : "emg_ambiguous $(m[:n])"

    # "EMG\s*Aux[123]" => emg_ambiguous [123]
    m = match(r"^\s*emg\s*aux(?<n>[123])\s*$"i, l)
    l = isnothing(m) ? l : "emg_ambiguous $(m[:n])"

    # "EMG_[LR]" => emg_ambiguous [12]
    m = match(r"^\s*emg(?<n>[123])\s*$"i, l)
    l = isnothing(m) ? l : "emg_ambiguous $(m[:n])"

    # EMG[+-]? => EMG Aux[123]
    l = l == "emg+" ? "emg_ambiguous 1" : l
    l = l == "emg-" ? "emg_ambiguous 2" : l
    l = l == "emg" ? "emg_ambiguous 3" : l    # postprocess: if only "emg_ambiguous 3" + legs present, replace with "EMG chin1"

    # other ambiguous forms
    l = l == "emg1-emg2" ? "emg_ambiguous 1" : l
    l = l == "emg2-emg1" ? "emg_ambiguous 2" : l
    l = l == "emg-a1" ? "emg_ambiguous 1" : l
    l = l == "emg emg" ? "emg_ambiguous 1" : l

    m = match(r"[\s\[,]*emg(?<i>[123]?)[\s\-]+emg([123]?)[\]\s,]*"i, l)
    l = isnothing(m) ? l : "emg_ambiguous $(m[:i])"

    ####
    ####

    # [sub][clavi]?
    m = match(r"l[^r]*sub"i, l)
    l = (isnothing(m) || occursin("subm", l)) ? l : "ecg avl"
    m = match(r"r[^l]*sub"i, l)
    l = (isnothing(m) || occursin("subm", l)) ? l : "ecg avr"

    # [sub]?[clav]
    m = match(r"l[^r]*clav"i, l)
    l = isnothing(m) ? l : "ecg avl"
    m = match(r"r[^l]*clav"i, l)
    l = isnothing(m) ? l : "ecg avr"

    # recordings with label in the transducer field
    l = l == "spectrum eeg" ? t : l

    # "EOG - L" and "EOG - R" should not be parsed as channel \minus channel
    m = match(r"^\s*EOG[\s\-]+(?<lr>[LR])\s*"i, l)
    if !isnothing(m)
        l = "EOG $(m[:lr])"
    end

    # "L - EOG" and "R- EOG" should not be parsed as channel \minus channel
    m = match(r"^[\[\s,\(]*(?<lr>[LR])[\s\-]+EOG(?<rest>[^\]\),]*)\s*[\]\s,\)]*$"i, l)
    if !isnothing(m)
        l = "EOG $(m[:lr])$(m[:rest])"
    end

    # "C2M1" => "C2-M1" etc
    m = match(r"\s*(?<channel>[fco][1234])\s*(?<ref>[am][12])\s*$"i, l)
    l = isnothing(m) ? l : "$(m[:channel])-$(m[:ref])"

    # "Chin EMG[12]?" => "EMG chin[12]?"
    m = match(r"\s*chin\s+emg(?<n>[12]?)\s*"i, l)
    l = isnothing(m) ? l : "emg chin$(m[:n])"

    # Chin[LR123]?.* => chin[lr123]?
    m = match(r"\s*chin(?<side>[lr123]?)\s*(?<rest>[\-chinlr123]*)\s*$"i, l)
    l = isnothing(m) ? l : "emg chin$(m[:side])$(m[:rest])"

    # "Menton (cen.)" => chin3
    m = match(r"\s*menton.*cen.*"i, l)
    l = isnothing(m) ? l : "chin3"

    # "Lower.Left-Upper" => "EMG chin1"
    l = startswith(l, "lower.left-upp") ? "emg chin1" : l
    l = startswith(l, "lower.right-upp") ? "emg chin2" : l
    l = startswith(l, "lower.left-low") ? "emg chin1" : l    # dipole; dubious assignment
    l = startswith(l, "lower.right-low") ? "emg chin2" : l   # dipole; dubious assignment

    # how to denote that sign is inverted???
    # "Upper-Lower.Left" => "EMG chin1"
    l = l == "upper-lower.left" ? "emg chin1" : l
    l = l == "upper-lower.righ" ? "emg chin2" : l

    # deutsch
    l = l == "emg li" ? "emg chin1" : l
    l = l == "emg re" ? "emg chin2" : l
    l = l == "emg mitte" ? "emg chin3" : l
    l = l == "emg1 kinn" ? "emg chin1" : l
    l = l == "emg2 kinn" ? "emg chin2" : l

    # EMG-subm[12] => EMG chin[12]
    m = match(r"\s*emg\-subm(?<n>[12])\s*"i, l)
    l = isnothing(m) ? l : "emg chin$(m[:n])"

    # label = "EMG", transducer_type = "FP1-FP2" => label = "EEG FP1-FP2"
    l = (l == "emg" && t == "fp1-fp2") ? "eeg fp1-fp2" : l
    l = (l == "emg" && t == "fp2-fp1") ? "eeg fp2-fp1" : l

    # label = "EMG", transducer_type = "1A-1R" => label = "EMG chin1"
    l = (l == "emg" && t == "1a-1r") ? "emg chin1" : l

    # EKG

    # EKG1, ECG2, etc are ambiguous
    # pick same default as for "ECG"
    # TODO should I instead pick[B another default,
    #      or make a new ambiguous ECG signal type for these cases?
    m = match(r"[\s\[,]*e[ck]g([123])[\s\-]*e[ck]g([123])[\]\s,]*"i, l)
    l = isnothing(m) ? l : "ecg $(m[1])"

    m = match(r"[\s\[,]*e[ck]g([123])[\]\s,]*"i, l)
    l = isnothing(m) ? l : "ecg $(m[1])"

    # [LR] ECG
    m = match(r"^[\[\s,]*(?<lr>[lr])\s*E[ck]g"i, l)
    l = isnothing(m) ? l : "ecg av$(m[:lr])"

    # [LR]-Leg[12]?
    m = match(r"^[\[\s,]*(?<lr>[lr])\s*\-\s*leg[1-2][\s\],]?"i, l)
    l = isnothing(m) ? l : "ecg av$(m[:lr])"

    # EMG Leg

    # {left,right,l,r}[\s-_]*leg
    m = match(r"^(\s*emg\s*)?(left|l)[_\-\s]*leg\s*(?<rest>.*)\s*$"i, l)
    l = isnothing(m) ? l : "emg left_anterior_tibialis$(m[:rest])"
    m = match(r"^(\s*emg\s*)?(right|r)[_\-\s]*leg\s*(?<rest>.*)\s*$"i, l)
    l = isnothing(m) ? l : "emg right_anterior_tibialis$(m[:rest])"

    # leg[\s-_]*{left,right,l,r}
    m = match(r"^(\s*emg\s*)?(?<lr>left|l)[_\-\s]*leg\s*(?<rest>.*)\s*$"i, l)
    l = isnothing(m) ? l : "emg left_anterior_tibialis$(m[:rest])"
    m = match(r"^(\s*emg\s*)?(?<lr>right|r)[_\-\s]*leg\s*(?<rest>.*)\s*$"i, l)
    l = isnothing(m) ? l : "emg right_anterior_tibialis$(m[:rest])"

    # (Tib|Leg)(-|/)[LR]
    m = match(r"^\s*(leg|tib)[/\-](?<lr>l|left|right|r)\s*$"i, l)
    l = isnothing(m) ? l : "$(startswith(m[:lr], "l") ? "left" : "right")_anterior_tibialis"

    # [LR]Leg1 - [LR]Leg2
    m = match(r"[\s\[,]*(?<lr>[lr])leg([12]?)[\s\-lr]*leg([12]?)[\]\s,]*"i, l)
    l = isnothing(m) ? l : "emg $(m[:lr])leg"

    # [LR] LEG[12]?  # with a minus `-`, this would be interpreted as ECG
    m = match(r"^[\[\s,]*(?<lr>[lr])\s*leg[1234]?"i, l)
    l = isnothing(m) ? l : "emg leg$(m[:lr])"

    return l
end

custom_labels = deepcopy(OndaEDF.STANDARD_LABELS)
#custom_labels[["emg"]] = [p for p in custom_labels[["emg"]] if startswith(first(p), "chin")]

custom_extractors = [edf -> extract_channels_by_label(edf, signal_names, channel_names; preprocess_labels=test_preprocessor)
                     for (signal_names, channel_names) in custom_labels]

function has_leg(h)
    h.kind != "emg" && return false
    return any(c -> startswith(c, "left_anterior_tibialis") || startswith(c, "right_anterior_tibialis"), h.channels)
end

@testset "Import EDF" begin

    @testset "edf_to_samples_info" begin
        results = map(test_edf_to_samples_info) do r
            edf = mock_edf(r)
            return last(OndaEDF.edf_header_to_onda_samples_info(edf; custom_extractors=custom_extractors))
        end
        # print result of mapping `edf_to_onda_samples` over `test_edf_to_samples_info.in`
        # this makes it easy to see effects of any changes made to OndaEDF by looking at
        # `diff test_edf_to_samples_info.{in,out}`
        prettyprint_diagnostic_info("test_edf_to_samples_info", results)
        prettyprint_diagnostic_info("no_eeg_tested", filter(nt -> !any(h -> h.kind == "eeg", map(first, nt.header_map)), results))
        prettyprint_diagnostic_info("no_eog_tested", filter(nt -> !any(h -> h.kind == "eog", map(first, nt.header_map)), results))
        prettyprint_diagnostic_info("no_chin_tested", filter(nt -> !any(h -> startswith(h.kind, "emg") && (h.kind == "emg_ambiguous" || any(c -> startswith(c, "chin"), h.channels)), map(first, nt.header_map)), results))
        prettyprint_diagnostic_info("no_leg_tested", filter(nt -> !any(has_leg, map(first, nt.header_map)), results))
        prettyprint_diagnostic_info("no_ekg_tested", filter(nt -> !any(h -> h.kind ∈ Set(("ecg", "ekg")), map(first, nt.header_map)), results))
        for (i, (r, expected)) in enumerate(zip(results, test_edf_to_samples_info))
            @test (i, setdiff(r.unextracted_edf_headers, expected.unextracted_edf_headers)) == (i, [])
        end
    end

    n_records = 100
    edf, edf_channel_indices = make_test_data(MersenneTwister(42), 256, 512, n_records)

    @testset "edf_to_onda_samples" begin
        returned_samples, nt = OndaEDF.edf_to_onda_samples(edf)
        @test length(returned_samples) == 13

        samples_info = Dict(s.info.kind => s.info for s in returned_samples)
        @test samples_info["tidal_volume"].channels == ["tidal_volume"]
        @test samples_info["tidal_volume"].sample_unit == "milliliter"
        @test samples_info["respiratory_effort"].channels == ["chest", "abdomen"]
        @test samples_info["respiratory_effort"].sample_unit == "microvolt"
        @test samples_info["snore"].channels == ["snore"]
        @test samples_info["snore"].sample_unit == "microvolt"
        @test samples_info["ecg"].channels == ["avl", "avr"]
        @test samples_info["ecg"].sample_unit == "microvolt"
        @test samples_info["positive_airway_pressure"].channels == ["ipap", "epap"]
        @test samples_info["positive_airway_pressure"].sample_unit == "centimeter_of_water"
        @test samples_info["heart_rate"].channels == ["heart_rate"]
        @test samples_info["heart_rate"].sample_unit == "beat_per_minute"
        @test samples_info["emg"].channels == ["intercostal", "left_anterior_tibialis", "right_anterior_tibialis"]
        @test samples_info["emg"].sample_unit == "microvolt"
        @test samples_info["eog"].channels == ["left", "right"]
        @test samples_info["eog"].sample_unit == "microvolt"
        @test samples_info["eeg"].channels == ["fpz", "f3-m2", "f4-m1", "c3-m2",
                                               "c4-m1", "o1-m2", "o2-a1"]
        @test samples_info["eeg"].sample_unit == "microvolt"
        @test samples_info["pap_device_cflow"].channels == ["pap_device_cflow"]
        @test samples_info["pap_device_cflow"].sample_unit == "liter_per_minute"
        @test samples_info["pap_device_leak"].channels == ["pap_device_leak"]
        @test samples_info["pap_device_leak"].sample_unit == "liter_per_minute"
        @test samples_info["sao2"].channels == ["sao2"]
        @test samples_info["sao2"].sample_unit == "percent"
        @test samples_info["ptaf"].channels == ["ptaf"]
        @test samples_info["ptaf"].sample_unit == "volt"

        @test all(Onda.duration(s) == Nanosecond(Second(200)) for s in returned_samples)
    end

    @testset "store_edf_as_onda" begin
        root = mktempdir()
        uuid = uuid4()
        returned_uuid, (returned_signals, annotations) = OndaEDF.store_edf_as_onda(root, edf, uuid)
        signals = Dict(s.kind => s for s in returned_signals)

        @test isfile(joinpath(root, "edf.onda.signals.arrow"))
        @test isfile(joinpath(root, "edf.onda.annotations.arrow"))

        @test returned_uuid == uuid
        @test length(returned_signals) == 13
        @testset "samples info" begin
            @test signals["tidal_volume"].channels == ["tidal_volume"]
            @test signals["tidal_volume"].sample_unit == "milliliter"
            @test signals["respiratory_effort"].channels == ["chest", "abdomen"]
            @test signals["respiratory_effort"].sample_unit == "microvolt"
            @test signals["snore"].channels == ["snore"]
            @test signals["snore"].sample_unit == "microvolt"
            @test signals["ecg"].channels == ["avl", "avr"]
            @test signals["ecg"].sample_unit == "microvolt"
            @test signals["positive_airway_pressure"].channels == ["ipap", "epap"]
            @test signals["positive_airway_pressure"].sample_unit == "centimeter_of_water"
            @test signals["heart_rate"].channels == ["heart_rate"]
            @test signals["heart_rate"].sample_unit == "beat_per_minute"
            @test signals["emg"].channels == ["intercostal", "left_anterior_tibialis", "right_anterior_tibialis"]
            @test signals["emg"].sample_unit == "microvolt"
            @test signals["eog"].channels == ["left", "right"]
            @test signals["eog"].sample_unit == "microvolt"
            @test signals["eeg"].channels == ["fpz", "f3-m2", "f4-m1", "c3-m2",
                                              "c4-m1", "o1-m2", "o2-a1"]
            @test signals["eeg"].sample_unit == "microvolt"
            @test signals["pap_device_cflow"].channels == ["pap_device_cflow"]
            @test signals["pap_device_cflow"].sample_unit == "liter_per_minute"
            @test signals["pap_device_leak"].channels == ["pap_device_leak"]
            @test signals["pap_device_leak"].sample_unit == "liter_per_minute"
            @test signals["sao2"].channels == ["sao2"]
            @test signals["sao2"].sample_unit == "percent"
            @test signals["ptaf"].channels == ["ptaf"]
            @test signals["ptaf"].sample_unit == "volt"
        end

        for signal in values(signals)
            @test signal.span.start == Nanosecond(0)
            @test signal.span.stop == Nanosecond(Second(200))
            @test signal.file_format == "lpcm.zst"
        end

        for (signal_name, edf_indices) in edf_channel_indices
            onda_samples = load(signals[string(signal_name)]).data
            edf_samples = mapreduce(transpose ∘ EDF.decode, vcat, edf.signals[edf_indices])
            @test isapprox(onda_samples, edf_samples; rtol=0.02)
        end

        @testset "Annotations import" begin
            @test length(annotations) == n_records * 4
            # check whether all four types of annotations are preserved on import:
            for i in 1:n_records
                start = Nanosecond(Second(i))
                stop = start + Nanosecond(Second(i + 1))
                # two annotations with same 1s span and different values:
                @test any(a -> a.value == "$i a" && a.span.start == start && a.span.stop == stop, annotations)
                @test any(a -> a.value == "$i b" && a.span.start == start && a.span.stop == stop, annotations)
                # two annotations with instantaneous (1ns) span and different values
                @test any(a -> a.value == "$i c" && a.span.start == start && a.span.stop == start + Nanosecond(1), annotations)
                @test any(a -> a.value == "$i d" && a.span.start == start && a.span.stop == start + Nanosecond(1), annotations)
            end
        end

        @testset "Table prefixes" begin
            prefix = @test_nowarn validate_arrow_prefix("edf")
            @test prefix == "edf"
            prefix = @test_nowarn validate_arrow_prefix("edf.something")
            @test prefix == "edf.something"
            @test_throws ArgumentError validate_arrow_prefix("/edf.something")
            @test_throws ArgumentError validate_arrow_prefix("subdir/something")
            prefix = @test_logs (:warn, r"Extracting prefix \"edf\"") validate_arrow_prefix("edf.onda.signals.arrow")
            @test prefix == "edf"
            prefix = @test_logs (:warn, r"Extracting prefix \"edf\"") validate_arrow_prefix("edf.onda.annotations.arrow")
            @test prefix == "edf"

            mktempdir() do root
                OndaEDF.store_edf_as_onda(root, edf, uuid; signals_prefix="edfff")
                @test isfile(joinpath(root, "edfff.onda.signals.arrow"))
                @test isfile(joinpath(root, "edfff.onda.annotations.arrow"))
            end

            mktempdir() do root
                OndaEDF.store_edf_as_onda(root, edf, uuid; annotations_prefix="edff")
                @test isfile(joinpath(root, "edf.onda.signals.arrow"))
                @test isfile(joinpath(root, "edff.onda.annotations.arrow"))
            end

            mktempdir() do root
                OndaEDF.store_edf_as_onda(root, edf, uuid; annotations_prefix="edff")
                @test isfile(joinpath(root, "edf.onda.signals.arrow"))
                @test isfile(joinpath(root, "edff.onda.annotations.arrow"))
            end

            mktempdir() do root
                OndaEDF.store_edf_as_onda(root, edf, uuid; signals_prefix="edfff", annotations_prefix="edff")
                @test isfile(joinpath(root, "edfff.onda.signals.arrow"))
                @test isfile(joinpath(root, "edff.onda.annotations.arrow"))
            end

            mktempdir() do root
                @test_logs (:warn, r"Extracting prefix") OndaEDF.store_edf_as_onda(root, edf, uuid; signals_prefix="edff.onda.signals.arrow", annotations_prefix="edf")
                @test isfile(joinpath(root, "edff.onda.signals.arrow"))
                @test isfile(joinpath(root, "edf.onda.annotations.arrow"))
            end

            mktempdir() do root
                @test_throws ArgumentError OndaEDF.store_edf_as_onda(root, edf, uuid; signals_prefix="stuff/edf", annotations_prefix="edf")
            end

        end
    end

end


text:
using Gen

@gen (static) function foo()
    @param mu::Float64
    y = @trace(normal(mu, 1), :y)
    return y
end

load_generated_functions()

init_param!(foo, :mu, -1)

trace, = generate(foo, (), choicemap((:y, 3)))
step_size = 0.01
for iter=1:1000
    accumulate_param_gradients!(trace, 0.)
    grad_val = get_param_grad(foo, :mu)
    set_param!(foo, :mu, get_param(foo, :mu) + step_size * grad_val)
    zero_param_grad!(foo, :mu)
end

@assert abs(get_param(foo, :mu) - 3) < 1e-2

ext: jl


text:
###include("Crystal.jl")
###include("DFToutMod.jl")
#using XMLDict

####################### Wannier90 specific
"""
    module Ewald

Module for electrostatic preperation.
"""
module Ewald
"""
Electrostatics
"""

#include("Atomdata.jl")
using ..Atomdata:atoms

using LinearAlgebra
using ..CrystalMod:crystal
###using ..CalcTB:distances_etc_3bdy
using SpecialFunctions
import Base.Threads.@spawn
#using Base.Threads

"""
   function getU(types)

Get values of `U` from `Atomdata`.
"""
function getU(types)

    U = zeros(length(types))
    for (c,t) in enumerate(types)
        U[c] = atoms[t].U
    end

    return U

end

"""
    function get_onsite(crys::crystal, U::Array{Float64,1})

return Onsite terms, Coulumb `U`.
"""
function get_onsite(crys::crystal, U::Array{Float64,1})

    return diagm(U)
#    gamma_onsite = zeros(Float64, crys.nat, crys.nat)
#
#    for i = 1:crys.nat
#        gamma_onsite[i,i] = U[i]
#    end##
#
#    return gamma_onsite

end

"""
    function electrostatics_getgamma(crys::crystal;  kappa=missing, noU=false, onlyU=false, screening = 1.0)

Main function. Does Ewald calculation on `crys`. Returns `gamma`, which is used in scf calculation.
This is only run once for a given `tb_crys` object and stored.

- `kappa` is the splitting parameter between real/k-space in Ewald calculation. Will estimate best one if not provided.
- `noU=false` for testing only
- `onlyU=false` for testing only
- `screening=1.0` Not used. Purpose is to reduce U values for values < 1.
"""
function electrostatics_getgamma(crys::crystal;  kappa=missing, noU=false, onlyU=false, screening = 1.0)
#noU and onlyU are for testing purposes

#R_keep, R_keep_ab, array_ind3, array_floats3, dist_arr, c_zero = distances_etc_3bdy(crys,cutoff2X, 0.0)

#    println("estimate")
    if ismissing(kappa)
#        kappa_default = 0.25
        kappa = estimate_best_kappa(crys.A)
#        println("kappa $kappa")
    end
    kappa = Float64(kappa)

    if noU
        println("noU - FOR TESTING")
        U = zeros(Float64, crys.nat)
    else
        U = getU(crys.types)
        U = U * screening
    end

    starting_size_rspace = 1
    starting_size_kspace = 1

    gamma_onsiteU = get_onsite(crys, U)

    T = typeof(crys.coords[1,1])
    gamma_rs = zeros(T, crys.nat, crys.nat)
    gamma_U = zeros(T, crys.nat, crys.nat)
    gamma_k = zeros(T, crys.nat, crys.nat)
    gamma_self = zeros(T, crys.nat, crys.nat)

    #can do these at same time.
    # can run real space and k-space in parallel with asyncronous parallelization. Only a minor improvement.
    rs = @spawn begin
    #rs = begin
        gamma_rs, gamma_U = real_space(crys, kappa, U, starting_size_rspace)
    end

#    println("gamma_rs")
#    println(gamma_rs)
#    println("ks")
    ks = @spawn begin
    #ks = begin
        gamma_k = k_space(crys, kappa, starting_size_kspace)
    end
#    println("gamma_k")
#    println(gamma_k)

    self = @spawn begin
        #self
        for i = 1:crys.nat
            gamma_self[i,i] -= kappa / sqrt(pi) * 2.0
        end
    end

    wait(ks)
    wait(self)
    wait(rs)

    if false #for debugging
        println("gamma_rs")
        println(gamma_rs)
        println("gamma_k")
        println(gamma_k)
        println("gamma_self")
        println(gamma_self)
        println("gamma_onsiteU")
        println(gamma_onsiteU)
        println("gamma_U")
        println(gamma_U)
        println()
        println("only 1/r")
        println(gamma_rs + gamma_k + gamma_self)
    end

    #rydberg units
    e2 = 2.0

    gamma_tot = e2*(gamma_rs + gamma_k + gamma_self + gamma_U) + gamma_onsiteU

    if onlyU #for debugging
        gamma_tot = gamma_onsiteU
    end

    return gamma_tot

end

"""
    function real_space(crys::crystal, kappa::Float64, U::Array{Float64}, starting_size_rspace=2)

Real-space Ewald sum.
"""
function real_space(crys::crystal, kappa::Float64, U::Array{Float64}, starting_size_rspace=2)

    T = typeof(crys.coords[1,1])

    first_iter = true
    old_val = 0.0

    gamma_ij_tot = zeros(T, crys.nat, crys.nat)
    gamma_ij_new = zeros(T, crys.nat, crys.nat)

    gamma_U_tot = zeros(T, crys.nat, crys.nat)
    gamma_U_new = zeros(T, crys.nat, crys.nat)

    Uconst = zeros(T, crys.nat, crys.nat)

    if sum(abs.(U)) > 1e-5
        useU = true
        for i = 1:crys.nat
            Fi = sqrt( 8* log(2)/pi ) / (U[i]/2.0) #rydberg units, U in Ryd , we divide by 2, and multiply by e^2 layer. For Hartree formula, see koskinen comp mater sci 47 (2009) 237
            for j = 1:crys.nat
#                Uconst[i,j] = sqrt(pi/2 * (U[i]^2 * U[j]^2 / (U[i]^2 + U[j]^2)))

                Fj = sqrt( 8* log(2)/pi ) / (U[j]/2.0)
                Uconst[i,j] = sqrt(4 * log(2) / (Fi^2 + Fj^2))
            end
        end
    else
        useU = false
    end

    #    R = zeros(T,1, 3)
    R = zeros(T, 3)
    Ra = zeros(T, 3)
    Rb = zeros(T, 3)

    coords_cart = crys.coords * crys.A
    coords_cartT = coords_cart'

    coords_cartTij = zeros(T, 3, crys.nat, crys.nat)
    for i = 1:crys.nat
        for j = 1:crys.nat
            coords_cartTij[:,i,j] = (@view coords_cartT[:,i]) - (@view coords_cartT[:,j])
        end
    end

    converged = false
    converged_old = false
    R0 = false

    N_list = [[1,1,1]]
    a1 = sqrt(sum(crys.A[1,:].^2))
    a2 = sqrt(sum(crys.A[2,:].^2))
    a3 = sqrt(sum(crys.A[3,:].^2))
    a = (minimum([a1,a2,a3]) - 0.0010) * 0.95

    for N = 1:70

        n1 = max(Int64(ceil( a * N / a1)), 2)
        n2 = max(Int64(ceil( a * N / a2)), 2)
        n3 = max(Int64(ceil( a * N / a3)), 2)

        if [n1,n2,n3] != N_list[end]
            N_list = push!(N_list, [n1,n2,n3])
#            println("N_list ", (n1,n2,n3))
        end
    end

    newcontr = 0.0
    newcontrU = 0.0

#    for N = starting_size_rspace:20
#        Nx = N
#        Ny = N
#        Nz = N
    Nxold = 0
    Nyold = 0
    Nzold = 0

    At = crys.A'

    for (Nx, Ny, Nz) in N_list
        gamma_ij_new[:,:] .= 0.0
        gamma_U_new[:,:] .= 0.0
        for x = -Nx:Nx
            Ra[1] = x
            for y = -Ny:Ny
                Ra[2] = y
                for z = -Nz:Nz
                    Ra[3] = z
                    if x != 0 || y != 0 || z != 0
                        R0 = false
                    else
                        R0 = true
                    end
                    R[:] = At * Ra
                    if first_iter == true || (abs(x) > Nxold || abs(y) > Nyold || abs(z) > Nzold )
                        for i = 1:crys.nat
                            for j = 1:crys.nat
                                #                                Rb = (@view coords_cartT[:,i]) - (@view coords_cartT[:,j]) - (@view R[:])
                                Rb = (@view coords_cartTij[:,i, j])  - (@view R[:])
                                r = (Rb'*Rb).^0.5
#                                r = sum(( (@view coords_cartT[:,i]) - (@view coords_cartT[:,j]) - (@view R[:]) ).^2)^0.5
                                #r = sum(( (coords_cart[i,:]) - (coords_cart[j,:]) - ( R[1,:] ) ).^2)^0.5
                                if i != j ||  R0 == false
                                    @inbounds gamma_ij_new[i,j] += erfc( kappa * r) / r

                                    if useU
                                        @inbounds gamma_U_new[i,j] += -erfc( Uconst[i,j] * r) / r  #see eq 3 in prb 66 075212, or koskinen comp mater sci 47 (2009) 237
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end

        newcontr = sum(abs.(gamma_ij_new))
        newcontrU = sum(abs.(gamma_U_new))

#        println("new $Nx $Ny $Nz real_space $newcontr $newcontrU")

        if newcontr < 1e-5 && newcontrU < 1e-5
            if converged_old == true || (newcontr < 1e-7 && newcontrU < 1e-7)
                converged = true
                break
            end
            converged_old = true
        end
        first_iter = false

        gamma_ij_tot += gamma_ij_new
        gamma_U_tot += gamma_U_new

        Nxold = Nx
        Nyold = Ny
        Nzold = Nz

    end

    if converged == false
        println("WARNING, real_space EWALD !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        println("real_space NOT converged  : $newcontr    $newcontrU")
        println(crys)

    end

    return gamma_ij_tot, gamma_U_tot

end

"""
    function estimate_best_kappa(A)

Estimate best value of `kappa` for Ewald sum.
Shouldn't effect final value, only calculation speed.
There is probably a better way to do this.
"""
function estimate_best_kappa(A)

    a1 = sqrt(sum(A[1,:].^2))
    a2 = sqrt(sum(A[2,:].^2))
    a3 = sqrt(sum(A[3,:].^2))

    a = minimum([a1,a2,a3])

#    println("a ", a)
#    println(a1)
#    println(a2)
#    println(a3)

    B = inv(A)'

    b1 = sqrt(sum(B[1,:].^2))
    b2 = sqrt(sum(B[2,:].^2))
    b3 = sqrt(sum(B[3,:].^2))

    b = minimum([b1,b2,b3])

#    println("b ", b)
#    println(b1)
#    println(b2)
#    println(b3)

    tot = Float64[]
    kappa_test = [0.00002 0.0001 0.0005 0.001 0.002 0.005 0.01 0.015 0.02 0.03 0.04 0.05 0.06 0.07 0.08 0.09 0.1 0.12 0.15 0.17 0.2 0.22 0.25 0.27  0.3 0.33 0.35 0.4 0.5 0.65 0.75 0.80 0.90 1.0 1.1 1.3 2.0 5.0 10.0 20.0 50.0 100.0]
    for kappa = kappa_test
        rs = erfc( kappa * a) / a * 10
        ks = exp(-b^2 / 4.0 / kappa^2  * (2 * pi)^2 ) / (2*pi*b)^2
        push!(tot, rs + ks)
#        println("$kappa rs $rs ks $ks  tot ",rs + ks)
    end
    i = argmin(tot)
    kappa = kappa_test[i]
#    println("bestkappa :", kappa)

    return kappa
end

"""
    function k_space(crys::crystal, kappa, starting_size_kspace=2)

K-space Ewald sum.
"""
function k_space(crys::crystal, kappa, starting_size_kspace=2)

    T = typeof(crys.coords[1,1])

    first_iter = true
    old_val = 0.0

    gamma_ij_tot = zeros(T, crys.nat, crys.nat)
    gamma_ij_new = zeros(T, crys.nat, crys.nat)
    K = zeros(T, 1,3)

    coords_cart = crys.coords * crys.A

    converged = false
    converged_old = false

    B = inv(crys.A)'
    vol = abs(det(crys.A))

    N_list = [[1,1,1]]
    b1 = sqrt(sum(B[1,:].^2))
    b2 = sqrt(sum(B[2,:].^2))
    b3 = sqrt(sum(B[3,:].^2))
    b = (minimum([b1,b2,b3]) - 0.0010) * 0.95

    for N = 1:30

        n1 = max(Int64(ceil( b * N / b1)), 2)
        n2 = max(Int64(ceil( b * N / b2)), 2)
        n3 = max(Int64(ceil( b * N / b3)), 2)

        if [n1,n2,n3] != N_list[end]
            N_list = push!(N_list, [n1,n2,n3])
        end
    end

    newcontr = 0.0

    Nxold = 0
    Nyold = 0
    Nzold = 0

    for (Nx, Ny, Nz) in N_list

#    for N = starting_size_kspace:25
        gamma_ij_new[:,:] .= 0.0
        for kx = -Nx:Nx
            for ky = -Ny:Ny
                for kz = -Nz:Nz
                    if kx == 0 && ky == 0 && kz == 0
                        continue
                    end
                    K[1,:] .= [kx,ky,kz]
                    K = K * B
                    k2 = sum(K.^2) * (2*pi)^2
                    factor_k = (2*pi) * exp(-k2 / 4.0 / kappa^2  ) / k2

                    if first_iter == true || abs(kx) > Nxold || abs(ky) > Nyold || abs(kz) > Nzold
                        for i = 1:crys.nat
                            for j = 1:crys.nat
                                kr = (K * ( (@view coords_cart[i,:]) - (@view coords_cart[j,:]) ) )[1]

                                exp_c = exp(2*pi*im*kr)
                                temp = real(exp_c )
                                gamma_ij_new[i,j] += factor_k * temp

                            end
                        end
                    end

                end
            end
        end
        first_iter = false
        newcontr = sum(abs.(gamma_ij_new))

        if newcontr < 1e-5
            if converged_old == true || newcontr < 1e-7
                converged = true
                break
            end
            converged_old = true #need convergence twice

        end

        gamma_ij_tot += 2.0*gamma_ij_new /vol

        Nxold = Nx
        Nyold = Ny
        Nzold = Nz

    end

    if converged == false
        println("WARNING, k_space EWALD !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        println("newcontr $newcontr")
        println(crys)
    end

    return gamma_ij_tot

end

end #end module


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

using Markdown
using InteractiveUtils

# ╔═╡ a3865853-58ad-4e27-bfbf-7223760ca026
md"""
# Methodical testing for the go-to-definition plugin

### FAQ
**Q:** Wait.. isn't it better to have this as javascript tests, so we can actually automate them and be sure that nothing breaks instead of relying on our flimsy human eyeballs?

**A:** Yes
"""

# ╔═╡ c93bd0ce-dffe-4416-a007-79116ab1f1cb
macro >(expr)
	nothing
end

# ╔═╡ 91bd5444-154e-48f1-bd51-fbdcfc5cfc14
macro +(expr)
	return esc(expr)
end

# ╔═╡ cf229fd1-5bdc-4a44-8272-5ecc872a2fbd
md"""
### Reading guide for ["Definitions"](#Definitions)

Every section here consists of a couple of definitions, followed by a cell that references all of these. The test is pressing every reference and making sure it selects the right part in the definition cell.
"""

# ╔═╡ 8d559967-84da-45c9-9e05-ce1adc36804c
md"""
### Reading guide for ["Usages"](#Usages)

There are some "predefined" variables: `x`, `y` and `z`; `@z`, `@z_str`; `X`, `Y` and `Z`. These are then used to make sure they get underlined. I decided that the `z`-s and `y`-s are intended to reference outside of the cell. If there is a way to override the variable locally in the cell or an identifier that isn't a reference,  use `x` for that. There should be **no underlined `x`'s**.
"""

# ╔═╡ 52620cb5-7d42-4d40-a3ae-6d9d29448014
md"## Definitions"

# ╔═╡ 7527a5b8-7eeb-4ec7-abbe-34ed68415ce1
md"---"

# ╔═╡ 0cdf94d0-dd1a-4a4a-b6c2-d2d290cd3224
import Example

# ╔═╡ 2d145800-c22f-4bbe-a6f3-bca8bb79d513
import Example as ExampleAlias

# ╔═╡ df05dcd2-fb6b-4907-ad80-6d231d6b6ce5
import Example: hello, domath as domath_alias

# ╔═╡ f22e3e28-6cda-4dbd-bc30-d2f7c2411e91
import Example: hello as hello_alias

# ╔═╡ 8948913b-9917-46c2-a91a-14386d269502
import Example.domath

# ╔═╡ e152285a-9265-4a13-bf0d-cfc0ba418595
import Example.domath as another_domath_alias

# ╔═╡ 90682238-2769-4f20-afcf-56ced93242ba
Example; ExampleAlias; hello; domath;

# ╔═╡ e30de56d-5985-42e7-abde-1ec90c7484c9
hello_alias; domath_alias; another_domath_alias; # Pluto doesn't recognize aliases?

# ╔═╡ 3be4e4b7-02f8-4f4c-bcab-2d796f38a93f
md"---"

# ╔═╡ 1d2f2927-49d0-404e-a324-4ab4f677d89a
abstract type _abstract_type1 end

# ╔═╡ 654566d0-9b46-49dd-a8ac-4d3f5708cba5
abstract type _abstract_type2{T} end

# ╔═╡ 8e823381-7b06-44b1-99e0-565083bc2434
_abstract_type1; _abstract_type2;

# ╔═╡ 709300fe-3b3d-471f-8fc1-608b904d18ed
md"---"

# ╔═╡ 4f996030-0ff2-4ea2-aa48-ac92765072c8
struct _struct1 end;

# ╔═╡ 670eb47b-225f-4626-9276-4cae08d90843
struct _struct2 x end;

# ╔═╡ 5afc38cb-ae2b-495e-93f9-a7ca85773ca9
struct _struct4{T} <: _abstract_type2{T} end

# ╔═╡ 327e846b-3edc-406a-936f-799475a7b3de
md"---"

# ╔═╡ 955ae894-24a1-4f8a-999d-3e8d809e0eba
(_tuple1,) = (:x,);

# ╔═╡ 09a24659-6646-40bb-8e7c-0d340d2b45da
(_tuple2, _tuple3, _tuple4...) = (:x1, :x2, :x3, :x4);

# ╔═╡ a85d6a78-b212-46b0-bacc-ad880911f2cb
_tuple1; _tuple2; _tuple3; _tuple4;

# ╔═╡ cb226d6d-f294-42b0-a40e-cb07e58fa776
md"---"

# ╔═╡ 554ad95c-2edd-4ef9-b532-0ec7cf08243c
_baretuple1, = (:x,);

# ╔═╡ d6e27295-77eb-4a84-98be-ea338a85f455
_baretuple2, _baretuple3... = (:x1, :x2, :x3);

# ╔═╡ 06948715-85e7-4609-94b1-bb91123d928c
_baretuple1; _baretuple2; _baretuple3;

# ╔═╡ 22032cef-5889-48d8-8709-bb7eeac89638
md"""
For both `_tuple4` and `_baretuple3`: JS side has these covered (shown by example below), but Pluto doesn't recognize it (thus it not working outside cells)
"""

# ╔═╡ 31713a89-8656-4cf9-aee0-ac24c3430c0a
let
	(_tuple2, _tuple3, _tuple4...) = (:x1, :x2, :x3, :x4);
	_baretuple2, _baretuple3... = (:x1, :x2, :x3);

	_baretuple3; _tuple4; # cmd+click me!
end;

# ╔═╡ 0a8e5b66-daaa-4f32-8dd8-1e8b9ec35465
md"---"

# ╔═╡ 4e3b5537-59b7-4d87-ae54-f1153c092ffa
const _const = :const;

# ╔═╡ 5d2d43f0-5a17-407a-a9ee-a041010e5579
global _global = :global;

# ╔═╡ e14b69ca-f325-440c-acb1-c33de71e75a2
_const; _global;

# ╔═╡ f9737a60-9909-445c-8382-4dd00aa267ba
md"## Usages"

# ╔═╡ a0f62011-92f2-44ad-b420-27217bcbd89c
md"### Predefined"

# ╔═╡ 90a81bc8-cce9-49bd-a3c4-126c74a515d3
x = "Hi, I'm x";

# ╔═╡ ad3ab162-2139-45b4-8877-18285f1c2fa8
y = "Hi, x";

# ╔═╡ 18e25073-6ea7-4287-975d-84c05743930e
z = "I was expecting a joke";

# ╔═╡ b7bd9512-a119-46e5-a6df-58b04e04bda0
_function3() = (x, z);

# ╔═╡ 1c530e0e-770c-44b5-a13e-e5b905e3d84c
_function6(x) = x, z;

# ╔═╡ a6eec43c-dc37-4113-a7e5-fafde8bdbd39
@Base.kwdef struct _struct3
	x::Int = z
end;

# ╔═╡ 0135431d-4b3f-4513-9982-937cc0b60837
_struct1; _struct2; _struct3; _struct4;

# ╔═╡ 669cf7b9-b1cc-4f01-8dc3-b9a156f9de63
macro z end;

# ╔═╡ 20804f0f-ca92-463d-afb7-604aeb6057f7
macro z_str end;

# ╔═╡ 7781bce5-2dd8-4d06-8bfd-7589bea619dc
abstract type X end

# ╔═╡ f9236dd2-4fd7-4d86-8dcb-cbc011da8b03
abstract type Y end

# ╔═╡ 4a9d7bb3-9ca7-4af9-b7e1-e730fbc1c570
abstract type Z end

# ╔═╡ 5d9701eb-46a3-483a-9b51-0e376031ae34
function _function1()::Z end;

# ╔═╡ 24716c18-5a36-4b12-9ed5-43069c81f15b
function _function2(x...)::Z
	x, z
end;

# ╔═╡ 5365c683-75b7-4f29-8cbe-810167409686
_function4(x::Z) = begin x end;

# ╔═╡ a85475ce-44ee-4b87-99b6-10b8b31d99fa
_function4(::Z) = (10, z, "hi");

# ╔═╡ b2c9c60f-c107-4cbb-bc7e-3018117e9ecb
_function5(x)::Z = (x, z);

# ╔═╡ 3f3440dc-4449-4a9f-b019-f2b5a3fc97b9
_function1; _function2; _function3; _function4; _function5; _function6;

# ╔═╡ e5b3cb7a-bf0d-479e-a533-d5699abe6f2a
md"### Cases"

# ╔═╡ 94c6fccc-ae3a-40bb-bf3f-a11f6d9244e2
@> sum(x for x in [1,2,3])

# ╔═╡ 18ea6b86-9ed5-4a75-852e-13cbb7f4d8da
@> function f(x=z)::z
	x, z
end

# ╔═╡ d8de87c9-e516-4f5f-978a-4c53071c5656
@> function f(x=z; y)::z
	x, y, z
end

# ╔═╡ 6bc106ce-ebba-4487-859e-9701ebeda022
md"---"

# ╔═╡ c283d772-e007-447a-aa88-a01d75d48e2e
@> z.@x

# ╔═╡ fa2aaff3-9b8a-4cdc-919d-2de38f54018b
@> @z.x

# ╔═╡ da330714-5130-470b-bd2a-6343c932f808
@> @z

# ╔═╡ 196a17eb-1ded-40be-a00b-b5e4d7fa1070
md"---"

# ╔═╡ 3e773c63-a8b7-4ad8-9dbb-74d519054f0f
@> struct x
	x1
	x2::z
	x3::z = z
	x4 = z
end

# ╔═╡ eac662c8-369e-4bf3-9aec-5f9c46fc4f77
@> x::z = z

# ╔═╡ 1d7c5c41-f5c8-40b3-b0b1-7da19572bef2
@> :(1 + $(z...) + 1)

# ╔═╡ 7521b3c9-a1bd-4786-aa19-44fd1db90149
@> (x, y...) = z

# ╔═╡ 4eb71552-16da-42bb-a9dc-f18e1b3e7190
@> z.@macro x = z

# ╔═╡ 3f009423-6a36-41c2-ad74-de596a8519c1
@> begin
	struct x end
	struct x{y} end
	struct x <: z end
	struct x <: z{x,z} end
	struct x{y} <: z{x,y,z} end

	x, y, z
end

# ╔═╡ fba1aebf-c00a-4750-89d3-a1b6e27c07a4
@> var"x" = 10;

# ╔═╡ 0b34c665-9c93-40ab-9965-d0ae36a2190d
@> var"x";

# ╔═╡ 29f7c36e-9b4d-4a91-83b3-162f1cd9e00c
@> (x=z) -> (x, z)

# ╔═╡ a478a0e0-51ba-4261-ac9d-5b3333b5a8bf
@> f(z)

# ╔═╡ 8ea3cd9e-c831-4140-977c-cc6cee5367b9
@> f(x=z)

# ╔═╡ 499392fa-4c50-4d0b-b584-098355e9af48
md"---"

# ╔═╡ 5749eb9d-f18f-4a91-8bb1-17d1c680d397
@> quote $(z) end

# ╔═╡ 259004c3-770f-4fe2-9c84-4c6f5d5c879c
@> quote quote $(x) end end

# ╔═╡ 91f8bfee-102a-437d-8c40-d7c694f0a384
@> quote quote $$(z) end end

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Example = "7876af07-990d-54b4-ab0e-23690620f79a"

[compat]
Example = "~0.5.3"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Example]]
git-tree-sha1 = "46e44e869b4d90b96bd8ed1fdcf32244fddfb6cc"
uuid = "7876af07-990d-54b4-ab0e-23690620f79a"
version = "0.5.3"
"""

# ╔═╡ Cell order:
# ╟─a3865853-58ad-4e27-bfbf-7223760ca026
# ╟─c93bd0ce-dffe-4416-a007-79116ab1f1cb
# ╟─91bd5444-154e-48f1-bd51-fbdcfc5cfc14
# ╟─cf229fd1-5bdc-4a44-8272-5ecc872a2fbd
# ╟─8d559967-84da-45c9-9e05-ce1adc36804c
# ╟─52620cb5-7d42-4d40-a3ae-6d9d29448014
# ╠═5d9701eb-46a3-483a-9b51-0e376031ae34
# ╠═24716c18-5a36-4b12-9ed5-43069c81f15b
# ╠═b7bd9512-a119-46e5-a6df-58b04e04bda0
# ╠═5365c683-75b7-4f29-8cbe-810167409686
# ╠═a85475ce-44ee-4b87-99b6-10b8b31d99fa
# ╠═b2c9c60f-c107-4cbb-bc7e-3018117e9ecb
# ╠═1c530e0e-770c-44b5-a13e-e5b905e3d84c
# ╠═3f3440dc-4449-4a9f-b019-f2b5a3fc97b9
# ╟─7527a5b8-7eeb-4ec7-abbe-34ed68415ce1
# ╠═0cdf94d0-dd1a-4a4a-b6c2-d2d290cd3224
# ╠═2d145800-c22f-4bbe-a6f3-bca8bb79d513
# ╠═df05dcd2-fb6b-4907-ad80-6d231d6b6ce5
# ╠═f22e3e28-6cda-4dbd-bc30-d2f7c2411e91
# ╠═8948913b-9917-46c2-a91a-14386d269502
# ╠═e152285a-9265-4a13-bf0d-cfc0ba418595
# ╠═90682238-2769-4f20-afcf-56ced93242ba
# ╠═e30de56d-5985-42e7-abde-1ec90c7484c9
# ╟─3be4e4b7-02f8-4f4c-bcab-2d796f38a93f
# ╠═1d2f2927-49d0-404e-a324-4ab4f677d89a
# ╠═654566d0-9b46-49dd-a8ac-4d3f5708cba5
# ╠═8e823381-7b06-44b1-99e0-565083bc2434
# ╟─709300fe-3b3d-471f-8fc1-608b904d18ed
# ╠═4f996030-0ff2-4ea2-aa48-ac92765072c8
# ╠═670eb47b-225f-4626-9276-4cae08d90843
# ╠═a6eec43c-dc37-4113-a7e5-fafde8bdbd39
# ╠═5afc38cb-ae2b-495e-93f9-a7ca85773ca9
# ╠═0135431d-4b3f-4513-9982-937cc0b60837
# ╟─327e846b-3edc-406a-936f-799475a7b3de
# ╠═955ae894-24a1-4f8a-999d-3e8d809e0eba
# ╠═09a24659-6646-40bb-8e7c-0d340d2b45da
# ╠═a85d6a78-b212-46b0-bacc-ad880911f2cb
# ╟─cb226d6d-f294-42b0-a40e-cb07e58fa776
# ╠═554ad95c-2edd-4ef9-b532-0ec7cf08243c
# ╠═d6e27295-77eb-4a84-98be-ea338a85f455
# ╠═06948715-85e7-4609-94b1-bb91123d928c
# ╟─22032cef-5889-48d8-8709-bb7eeac89638
# ╠═31713a89-8656-4cf9-aee0-ac24c3430c0a
# ╟─0a8e5b66-daaa-4f32-8dd8-1e8b9ec35465
# ╠═4e3b5537-59b7-4d87-ae54-f1153c092ffa
# ╠═5d2d43f0-5a17-407a-a9ee-a041010e5579
# ╠═e14b69ca-f325-440c-acb1-c33de71e75a2
# ╟─f9737a60-9909-445c-8382-4dd00aa267ba
# ╟─a0f62011-92f2-44ad-b420-27217bcbd89c
# ╠═90a81bc8-cce9-49bd-a3c4-126c74a515d3
# ╠═ad3ab162-2139-45b4-8877-18285f1c2fa8
# ╠═18e25073-6ea7-4287-975d-84c05743930e
# ╠═669cf7b9-b1cc-4f01-8dc3-b9a156f9de63
# ╠═20804f0f-ca92-463d-afb7-604aeb6057f7
# ╠═7781bce5-2dd8-4d06-8bfd-7589bea619dc
# ╠═f9236dd2-4fd7-4d86-8dcb-cbc011da8b03
# ╠═4a9d7bb3-9ca7-4af9-b7e1-e730fbc1c570
# ╟─e5b3cb7a-bf0d-479e-a533-d5699abe6f2a
# ╠═94c6fccc-ae3a-40bb-bf3f-a11f6d9244e2
# ╠═18ea6b86-9ed5-4a75-852e-13cbb7f4d8da
# ╠═d8de87c9-e516-4f5f-978a-4c53071c5656
# ╟─6bc106ce-ebba-4487-859e-9701ebeda022
# ╠═c283d772-e007-447a-aa88-a01d75d48e2e
# ╠═fa2aaff3-9b8a-4cdc-919d-2de38f54018b
# ╠═da330714-5130-470b-bd2a-6343c932f808
# ╟─196a17eb-1ded-40be-a00b-b5e4d7fa1070
# ╠═3e773c63-a8b7-4ad8-9dbb-74d519054f0f
# ╠═eac662c8-369e-4bf3-9aec-5f9c46fc4f77
# ╠═1d7c5c41-f5c8-40b3-b0b1-7da19572bef2
# ╠═7521b3c9-a1bd-4786-aa19-44fd1db90149
# ╠═4eb71552-16da-42bb-a9dc-f18e1b3e7190
# ╠═3f009423-6a36-41c2-ad74-de596a8519c1
# ╠═fba1aebf-c00a-4750-89d3-a1b6e27c07a4
# ╠═0b34c665-9c93-40ab-9965-d0ae36a2190d
# ╠═29f7c36e-9b4d-4a91-83b3-162f1cd9e00c
# ╠═a478a0e0-51ba-4261-ac9d-5b3333b5a8bf
# ╠═8ea3cd9e-c831-4140-977c-cc6cee5367b9
# ╟─499392fa-4c50-4d0b-b584-098355e9af48
# ╠═5749eb9d-f18f-4a91-8bb1-17d1c680d397
# ╠═259004c3-770f-4fe2-9c84-4c6f5d5c879c
# ╠═91f8bfee-102a-437d-8c40-d7c694f0a384
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
"""
    TableMetaData(kwargs...)

Has the following fields, which default to `nothing` if not specified with a keyword:
- `name` - a name for the table
- `id` - if a mort.SOA.org sourced table, will be the identifying table ID
- `provider` - Where the rates came from
- `reference` - Source for more info on table
- `content_type`
- `description`
- `comments`
- `source_path`

When you call a `MortalityTable` interactively, it will nicely print this summary infomration.

# Example content from mort.SOA.org:

- **Table Identity**: 1076
- **Provider Domain**: actuary.org
- **Provider Name**: American Academy of Actuaries
- **Table Reference**: Tillinghast, “American Council of Life Insurers: Preferred Version of 2001 CSO Mortality Tables”, ...
- **Content Type**: CSO/CET
- **Table Name**: 2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB
- **Table Description**: 2001 Commissioners Standard Ordinary (CSO) Super Preferred Select and Ultimate Table – Male Nonsmoker. Basis: Age Nearest Birthday. Minimum Select Age: 0. Maximum Select Age: 99. Minimum Ultimate Age: 16. Maximum Ultimate Age: 120
- **Comments**: Study Data: A preferred version of the 2001 Commissioners Standard Ordinary (CSO) table ...

And the `source_path` would be: [https://mort.soa.org/ViewTable.aspx?&TableIdentity=1076](https://mort.soa.org/ViewTable.aspx?&TableIdentity=1076)

# Example usage:
```
julia-repl> TableMetaData(name="My Table Name")
TableMetaData("My Table Name", nothing, nothing, nothing, nothing, nothing, nothing, nothing)
```

"""
Base.@kwdef struct TableMetaData
    name::Union{Nothing,String} = nothing
    id::Union{Nothing,String} = nothing
    provider::Union{Nothing,String} = nothing
    reference::Union{Nothing,String} = nothing
    content_type::Union{Nothing,String} = nothing
    description::Union{Nothing,String} = nothing
    comments::Union{Nothing,String} = nothing
    source_path::Union{Nothing,String} = nothing
end


text:
# Rectangle methods

RectanglePoints(x1::Int, y1::Int, x2::Int, y2::Int)  = RectanglePoints(Point(x1, y1), Point(x2, y2))
RectanglePoints(p1::CartesianIndex{2}, p2::CartesianIndex{2}) = RectanglePoints(Point(p1[1], p1[2]), Point(p2[1], p2[2]))

draw!(img::AbstractArray{T, 2}, rectangle::RectanglePoints, color::T) where {T<:Colorant} = draw!(img, Polygon(rectangle), color)


text:
#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5

# Short description of this file: Chain rule applied to NeuralPDE system

Pkg.add("ModelingToolkit")
Pkg.add("DiffEqFlux")
Pkg.add("Flux")
Pkg.add("Plots")
Pkg.add("Test")
Pkg.add("Optim")
Pkg.add("CUDA")

using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux, CUDA
using Plots, PyPlot

@parameters t x y η θ
@variables u1(..), u2(..), u3(..), u4(..), u5(..), u6(..), u7(..)
@derivatives Dt'~t
@derivatives Dx'~x
@derivatives Dy'~y
@derivatives Dη'~η
@derivatives Dxx''~x
@derivatives Dyy''~y

#functions
u  = u1(t,x,y,η)
v  = u2(t,x,y,η)
w  = u3(t,x,y,η)
ω  = u4(t,x,y,η)
θm = u5(t,x,y,η)
qm = u6(t,x,y,η)
z  = u7(t,x,y,η)

#physical constants
g  = 9.81 #constant
γ  = 1.4 #air cp/cv=1.4
ρd = 1 #dry air density
p0 = 100000 #reference pressure 10^5 Pa
Rd = 8.31 #gas constant for dry air
ηc = 0.2 #page 8 of Advanced Research WRF V4 (http://dx.doi.org/10.5065/1dfh-6p97)

#simulation inputs
sum_of_qs = 0.1 #read from file
#dx =
#dy =
#dz =
Δx = 1 #dummy
Δy = 1 #dummy

#scale factors
dist = x^2 + y^2 #dummy for distance on earth
Dxd  = 2*x #Dx(dist)
Dyd  = 2*y #Dy(dist)
mx   = Δx/dist
my   = Δy/dist

#definitions
αd = 1/ρd
p  = p0*(Rd*θm/(p0*αd))^γ
c1 = 2*ηc*ηc/((1 - ηc)*(1 - ηc)*(1 - ηc))
c2 = -ηc*(4 + ηc + ηc*ηc)/((1 - ηc)*(1 - ηc)*(1 - ηc))
c3 = 2*(1 + ηc + ηc*ηc)/((1 - ηc)*(1 - ηc)*(1 - ηc))
c4 = -(1 + ηc)/((1 - ηc)*(1 - ηc)*(1 - ηc))
B  = c1 + c2*η + c3*η*η + c4*η*η*η
ps = 10 #dummy
pt = 1 #dummy
pd = B*(ps - pt) + (η - B)*(p0 - pt) + pt
μd = (ps - pt)*(c2 + 2*c3*η + 3*c4*η*η) + (p0 - pt)*(1 - c2 - 2*c3*η - 3*c4*η*η) #[μd = Dη(pd)]
D2pd = (ps - pt)*(2*c3 + 6*c4*η) + (p0 - pt)*(-2*c3 - 6*c4*η) # (maybe find η as a function of pd)
#ω  = Dt(η)
U  = μd*u/my
V  = μd*v/mx
W  = μd*w/my
Ω  = μd*ω/my
Θm = μd*θm
Qm = μd*qm
ϕ  = g*z
α  = αd/(1 + sum_of_qs)

#divergences
#divVu  = Dx(U*u)  + Dy(V*u)  + Dη(Ω*u)
#divVv  = Dx(U*v)  + Dy(V*v)  + Dη(Ω*v)
#divVw  = Dx(U*w)  + Dy(V*w)  + Dη(Ω*w)
#divVθm = Dx(U*θm) + Dy(V*θm) + Dη(Ω*θm)
#divV   = Dx(V)    + Dy(V)    + Dη(V)
#divVqm = Dx(U*qm) + Dy(V*qm) + Dη(Ω*qm)
divVu1 = μd*u/Δy*(u*Dxd + 2*dist*Dx(u))
divVu2 = μd/Δx*(u*v*Dyd + dist*u*Dy(v) + dist*v*Dy(u))
divVu3 = dist/Δy*(D2pd*u*ω + μd*Dη(u)*ω + μd*u*Dη(ω))
divVu  = divVu1 + divVu2 + divVu3

divVv1 = μd/Δy*(u*v*Dxd + dist*u*Dx(v) + dist*v*Dx(u))
divVv2 = μd/Δx*(v*v*Dyd + dist*2*v*Dy(v))
divVv3 = 1/my*(D2pd*w*v + μd*Dη(ω)*v + μd*ω*Dη(v))
divVv  = divVv1 + divVv2 + divVv3

divVw1 = μd/Δy*(Dxd*u*w + dist*Dx(u)*w + dist*u*Dx(w))
divVw2 = μd/Δx*(Dyd*v*w + dist*Dy(v)*w + dist*v*Dy(w))
divVw3 = 1/my*(D2pd*ω*w + μd*Dη(ω)*w   + μd*ω*Dη(w))
divVw  = divVw1 + divVw2 + divVw3

divVθm1 = μd/Δy*(Dxd*u*θm + dist*Dx(u)*θm + dist*u*Dx(θm))
divVθm2 = μd/Δx*(Dyd*v*θm + dist*Dy(v)*θm + dist*v*Dy(θm))
divVθm3 = 1/my*(D2pd*ω*θm + μd*Dη(ω)*θm   + μd*ω*Dη(θm))
divVθm  = divVθm1 + divVθm2 + divVθm3

divV1 = μd/Δx*(Dxd*v + dist*Dx(v))
divV2 = μd/Δx*(Dyd*v + dist*Dy(v))
divV3 = 1/mx*(D2pd*v + μd*Dη(v))
divV  = divV1 + divV2 + divV3

divVqm1 = μd/Δy*(Dxd*u*qm + dist*Dx(u)*qm + dist*u*Dx(qm))
divVqm2 = μd/Δx*(Dyd*v*qm + dist*Dy(v)*qm + dist*v*Dy(qm))
divVqm3 = 1/my*(D2pd*ω*qm + μd*Dη(ω)*qm   + μd*ω*Dη(qm))
divVqm  = divVqm1 +divVqm2 + divVqm3

#inner products
innerV∇ϕ = U*g*Dx(z) + V*g*Dy(z) + Ω*g*Dη(z) #U*Dx(ϕ) + V*Dy(ϕ) + Ω*Dη(ϕ)

#geographic parametrization
Ωe = 1 #dummy #angular rotation rate of earth
ψ  = 1 #dummy #latitude on earth
re = 1000 #dummy #radius of earth
αr = 1 #dummy #local rotation angle between y-axis and the meridians
f  = 2*Ωe*sin(ψ)
e  = 2*Ωe*cos(ψ)

#external forcings
Dymx = -Δx/(dist*dist)*Dyd
Dxmy = -Δy/(dist*dist)*Dxd

FU  =  (mx/my)*(f + u*(my/mx)*Dymx - v*Dxmy)*V - (u/re + e*cos(αr))*W
FV  = -(my/mx)*((f + u*(my/mx)*Dymx - v*Dxmy*U + (v/re - e*sin(αr))*W))
FW  = e*(U*cos(αr) - (mx/my)*V*sin(αr) + 1/re*(u*U + (mx/my)*v*V))
FΘm = 1 #dummy
FQm = 1 #dummy

#equations (Flux-Form Euler)
#eq1 = Dt(U) + divVu + μd*α*Dx(p) + (α/αd)*Dη(p)*Dx(ϕ) ~ FU
#eq2 = Dt(V) + divVv + μd*α*Dy(p) + (α/αd)*Dη(p)*Dy(ϕ) ~ FV
#eq3 = Dt(W) + divVw - g*((α/αd)*Dη(p) - μd)           ~ FW
#eq4 = Dt(Θm) + divVθm                                 ~ FΘm
#eq5 = Dt(μd) + divV                                   ~ 0
#eq6 = Dt(ϕ)  + (1/μd)*(innerV∇ϕ - g*W)                ~ 0
#eq7 = Dt(Qm) + divVqm                                 ~ FQm

eq1 = μd/my*Dt(u) + divVu + μd*α*p0*(Rd/(p0*αd))^γ*γ*θm^(γ-1)*Dx(θm) + (α/αd)*p0*(Rd/(p0*αd))^γ*γ*θm^(γ-1)*Dη(θm)*g*Dx(z) ~ FU
eq2 = μd/mx*Dt(v) + divVv + μd*α*p0*(Rd/(p0*αd))^γ*γ*θm^(γ-1)*Dy(θm) + (α/αd)*p0*(Rd/(p0*αd))^γ*γ*θm^(γ-1)*Dη(θm)*g*Dy(z) ~ FV
eq3 = μd/my*Dt(w) + divVw - g*((α/αd)*p0*(Rd/(p0*αd))^γ*γ*θm^(γ-1)*Dη(θm) - μd)                                           ~ FW

eq4 = μd*Dt(θm) + divVθm                 ~ FΘm #coupling with SFIRE
eq5 = divV                               ~ 0
eq6 = g*Dt(z)  + (1/μd)*(innerV∇ϕ - g*W) ~ 0
eq7 = μd*Dt(qm) + divVqm                 ~ FQm  #coupling with SFIRE

#system of equations
eqs = [eq1,eq2,eq3,eq4,eq5,eq6,eq7]

# Initial and boundary conditions
bcs = [u1(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5
       u2(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5
       u3(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5
       u4(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5
       u5(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5
       u6(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5
       u7(0,x,y,η,θ) ~ (x^2 + y^2 + η^2)^0.5]

# Space and time domains
domains = [t ∈ IntervalDomain(0.0,2.0),
           x ∈ IntervalDomain(0.0,2.0),
           y ∈ IntervalDomain(0.0,2.0),
           η ∈ IntervalDomain(0.0,2.0)]# Discretization

dx = 0.125; dy= 0.125; dt = 0.125; dη = 0.125

# Neural network
chain = FastChain(FastDense(7,32,Flux.σ),FastDense(32,32,Flux.σ),FastDense(32,7))

#=High level
discretization = NeuralPDE.PhysicsInformedNN([dt,dx,dy,dη],
                                             chain,
                                             strategy= NeuralPDE.GridTraining())

pde_system = PDESystem(eqs,bcs,domains,[t,x,y,η],[p])

prob = NeuralPDE.discretize(pde_system,discretization)
# https://github.com/SciML/NeuralPDE.jl/blob/5c837fcbcfcd2a1153ee5278df940da01bcd44fb/src/pinns_pde_solve.jl#L399

res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=2000)
phi = discretization.phi
=#

#strategy = GridTraining()
strategy = StochasticTraining()
discretization = PhysicsInformedNN([dt,dx,dy,dη],chain,strategy=strategy)

indvars = [t,x,y,η]
depvars = [u1,u2,u3,u4,u5,u6,u7]

dim = length(domains)

expr_pde_loss_function = build_loss_function(eqs,indvars,depvars)
expr_bc_loss_functions = [build_loss_function(bc,indvars,depvars) for bc in bcs]
train_sets = generate_training_sets(domains,[dt,dx,dy,dη],bcs,indvars,depvars)

train_domain_set, train_bound_set, train_set = train_sets

phi = discretization.phi
autodiff = discretization.autodiff
derivative = discretization.derivative
initθ = discretization.initθ

pde_loss_function = get_loss_function(eval(expr_pde_loss_function),
                                      train_domain_set,
                                      phi,
                                      derivative,
                                      strategy)
bc_loss_function = get_loss_function(eval.(expr_bc_loss_functions),
                                     train_bound_set,
                                     phi,
                                     derivative,
                                     strategy)

function loss_function(θ,p)
    return  bc_loss_function(θ) + pde_loss_function(θ)
end

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

f = NeuralPDE.OptimizationFunction(loss_function, GalacticOptim.AutoZygote())

prob = GalacticOptim.OptimizationProblem(f, initθ)

# optimizer
opt = GalacticOptim.ADAM(0.01)
res = GalacticOptim.solve(prob, opt; cb = cb, maxiters=2)
phi = discretization.phi

ts,xs,ys,ηs = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
#u_predict_contourf = reshape([first(phi([t,x,y],res.minimizer)) for t in ts for x in xs for y in ys] ,length(ys),length(xs),length(ts))
#Plots.plot(ts, xs, ys, u_predict_contourf, linetype=:contourf,title = "predict")

#u_predict = [[first(phi([t,x,y],res.minimizer)) for x in xs for y in ys] for t in ts ]
u_predict = [reshape([first(phi([t,x,y],res.minimizer)) for x in xs for y in ys], (length(xs),length(ys))) for t in ts ]
p1= Plots.plot(xs, u_predict[2],title = "t = 0.1");
p2= Plots.plot(xs, u_predict[6],title = "t = 0.5");
p3= Plots.plot(xs, u_predict[end],title = "t = 1");
Plots.plot(p1,p2,p3)


text:
module PMTK4Stats

    export hdiFromIcdf
    include("hdiFromIcdf.jl")
    export normalizeLogspace
    include("tools.jl")

end


text:
using Plots
using StatsPlots
using Distributions
gr()
display(plot(Normal(3,5),lw=3))
savefig("normal.png")

max_stars_repo_name: sergev/vak-opensource

max_issues_repo_name: sergev/vak-opensource

max_forks_repo_name: sergev/vak-opensource


text:
using Flux: Dense

in_channel = 3
out_channel = 5
N = 4
adj = T[0. 1. 0. 1.;
       1. 0. 1. 0.;
       0. 1. 0. 1.;
       1. 0. 1. 0.]

adj_single_vertex =   T[0. 0. 0. 1.;
                        0. 0. 0. 0.;
                        0. 0. 0. 1.;
                        1. 0. 1. 0.]

@testset "layer" begin
    @testset "GCNConv" begin
        X = rand(in_channel, N)
        @testset "layer with graph" begin
            gc = GCNConv(adj, in_channel=>out_channel)
            @test size(gc.weight) == (out_channel, in_channel)
            @test size(gc.bias) == (out_channel,)
            @test graph(gc.fg) === adj

            Y = gc(X)
            @test size(Y) == (out_channel, N)

            # Test with transposed features
            Xt = rand(N, in_channel)
            Y = gc(transpose(Xt))
            @test size(Y) == (out_channel, N)

            g = Zygote.gradient(x -> sum(gc(x)), X)[1]
            @test size(g) == size(X)

            g = Zygote.gradient(model -> sum(model(X)), gc)[1]
            @test size(g.weight) == size(gc.weight)
            @test size(g.bias) == size(gc.bias)
        end

        @testset "layer without graph" begin
            gc = GCNConv(in_channel=>out_channel)
            @test size(gc.weight) == (out_channel, in_channel)
            @test size(gc.bias) == (out_channel,)
            @test !has_graph(gc.fg)

            fg = FeaturedGraph(adj, X)
            fg_ = gc(fg)
            @test size(node_feature(fg_)) == (out_channel, N)
            @test_throws AssertionError gc(X)

            # Test with transposed features
            Xt = rand(N, in_channel)
            fgt = FeaturedGraph(adj, transpose(Xt))
            fgt_ = gc(fgt)
            @test size(node_feature(fgt_)) == (out_channel, N)

            g = Zygote.gradient(x -> sum(node_feature(gc(x))), fg)[1]
            @test size(g[].nf) == size(X)

            g = Zygote.gradient(model -> sum(node_feature(model(fg))), gc)[1]
            @test size(g.weight) == size(gc.weight)
            @test size(g.bias) == size(gc.bias)
        end
    end

    @testset "ChebConv" begin
        k = 6
        X = rand(in_channel, N)
        @testset "layer with graph" begin
            cc = ChebConv(adj, in_channel=>out_channel, k)
            @test size(cc.weight) == (out_channel, in_channel, k)
            @test size(cc.bias) == (out_channel,)
            @test graph(cc.fg) === adj
            @test cc.k == k
            @test cc.in_channel == in_channel
            @test cc.out_channel == out_channel

            Y = cc(X)
            @test size(Y) == (out_channel, N)

            # Test with transposed features
            Xt = rand(N, in_channel)
            Y = cc(transpose(Xt))
            @test size(Y) == (out_channel, N)

            g = Zygote.gradient(x -> sum(cc(x)), X)[1]
            @test size(g) == size(X)

            g = Zygote.gradient(model -> sum(model(X)), cc)[1]
            @test size(g.weight) == size(cc.weight)
            @test size(g.bias) == size(cc.bias)
        end

        @testset "layer without graph" begin
            cc = ChebConv(in_channel=>out_channel, k)
            @test size(cc.weight) == (out_channel, in_channel, k)
            @test size(cc.bias) == (out_channel,)
            @test !has_graph(cc.fg)
            @test cc.k == k
            @test cc.in_channel == in_channel
            @test cc.out_channel == out_channel

            fg = FeaturedGraph(adj, X)
            fg_ = cc(fg)
            @test size(node_feature(fg_)) == (out_channel, N)
            @test_throws AssertionError cc(X)

            # Test with transposed features
            Xt = rand(N, in_channel)
            fgt = FeaturedGraph(adj, transpose(Xt))
            fgt_ = cc(fgt)
            @test size(node_feature(fgt_)) == (out_channel, N)

            g = Zygote.gradient(x -> sum(node_feature(cc(x))), fg)[1]
            @test size(g[].nf) == size(X)

            g = Zygote.gradient(model -> sum(node_feature(model(fg))), cc)[1]
            @test size(g.weight) == size(cc.weight)
            @test size(g.bias) == size(cc.bias)
        end
    end

    @testset "GraphConv" begin
        @testset "layer with graph" begin
            gc = GraphConv(adj, in_channel=>out_channel)
            @test graph(gc.fg) == [[2,4], [1,3], [2,4], [1,3]]
            @test size(gc.weight1) == (out_channel, in_channel)
            @test size(gc.weight2) == (out_channel, in_channel)
            @test size(gc.bias) == (out_channel,)

            X = rand(in_channel, N)
            Y = gc(X)
            @test size(Y) == (out_channel, N)

            # Test with transposed features
            Xt = rand(N, in_channel)
            Y = gc(transpose(Xt))
            @test size(Y) == (out_channel, N)

            g = Zygote.gradient(x -> sum(gc(x)), X)[1]
            @test size(g) == size(X)

            g = Zygote.gradient(model -> sum(model(X)), gc)[1]
            @test size(g.weight1) == size(gc.weight1)
            @test size(g.weight2) == size(gc.weight2)
            @test size(g.bias) == size(gc.bias)
        end

        @testset "layer without graph" begin
            gc = GraphConv(in_channel=>out_channel)
            @test size(gc.weight1) == (out_channel, in_channel)
            @test size(gc.weight2) == (out_channel, in_channel)
            @test size(gc.bias) == (out_channel,)

            X = rand(in_channel, N)
            fg = FeaturedGraph(adj, X)
            fg_ = gc(fg)
            @test size(node_feature(fg_)) == (out_channel, N)
            @test_throws AssertionError gc(X)

            # Test with transposed features
            Xt = rand(N, in_channel)
            fgt = FeaturedGraph(adj, transpose(Xt))
            fgt_ = gc(fgt)
            @test size(node_feature(fgt_)) == (out_channel, N)

            g = Zygote.gradient(x -> sum(node_feature(gc(x))), fg)[1]
            @test size(g[].nf) == size(X)

            g = Zygote.gradient(model -> sum(node_feature(model(fg))), gc)[1]
            @test size(g.weight1) == size(gc.weight1)
            @test size(g.weight2) == size(gc.weight2)
            @test size(g.bias) == size(gc.bias)
        end
    end

    @testset "GATConv" begin

        X = rand(in_channel, N)

        for adj_gat in [adj, adj_single_vertex]

            @testset "layer with graph" begin
                @testset "concat=true" begin
                    for heads = [1, 6]
                        gat = GATConv(adj_gat, in_channel=>out_channel, heads=heads, concat=true)

                        if adj_gat == adj
                            @test graph(gat.fg) == [[2,4], [1,3], [2,4], [1,3]]
                        end
                        if adj_gat == adj_single_vertex
                            @test graph(gat.fg) == [[4], Int64[], [4], [1, 3]]
                        end

                        @test size(gat.weight) == (out_channel * heads, in_channel)
                        @test size(gat.bias) == (out_channel * heads,)
                        @test size(gat.a) == (2*out_channel, heads)

                        Y = gat(X)
                        @test size(Y) == (out_channel * heads, N)

                        # Test with transposed features
                        Xt = rand(N, in_channel)
                        Y = gat(transpose(Xt))
                        @test size(Y) == (out_channel * heads, N)

                        g = Zygote.gradient(x -> sum(gat(x)), X)[1]
                        @test size(g) == size(X)

                        g = Zygote.gradient(model -> sum(model(X)), gat)[1]
                        @test size(g.weight) == size(gat.weight)
                        @test size(g.bias) == size(gat.bias)
                        @test size(g.a) == size(gat.a)
                    end
                end

                @testset "concat=false" begin
                    for heads = [1, 6]
                        gat = GATConv(adj_gat, in_channel=>out_channel, heads=heads, concat=false)
                        if adj_gat == adj
                            @test graph(gat.fg) == [[2,4], [1,3], [2,4], [1,3]]
                        end
                        if adj_gat == adj_single_vertex
                            @test graph(gat.fg) == [[4], Int64[], [4], [1, 3]]
                        end
                        @test size(gat.weight) == (out_channel * heads, in_channel)
                        @test size(gat.bias) == (out_channel * heads,)
                        @test size(gat.a) == (2*out_channel, heads)

                        Y = gat(X)
                        @test size(Y) == (out_channel, N)

                        # Test with transposed features
                        Xt = rand(N, in_channel)
                        Y = gat(transpose(Xt))
                        @test size(Y) == (out_channel, N)

                        g = Zygote.gradient(x -> sum(gat(x)), X)[1]
                        @test size(g) == size(X)

                        g = Zygote.gradient(model -> sum(model(X)), gat)[1]
                        @test size(g.weight) == size(gat.weight)
                        @test size(g.bias) == size(gat.bias)
                        @test size(g.a) == size(gat.a)
                    end
                end
            end

            @testset "layer without graph" begin
                fg = FeaturedGraph(adj_gat, X)

                @testset "concat=true" begin
                    for heads = [1, 6]
                        gat = GATConv(in_channel=>out_channel, heads=heads, concat=true)
                        @test size(gat.weight) == (out_channel * heads, in_channel)
                        @test size(gat.bias) == (out_channel * heads,)
                        @test size(gat.a) == (2*out_channel, heads)

                        fg_ = gat(fg)
                        Y = node_feature(fg_)
                        @test size(Y) == (out_channel * heads, N)
                        @test_throws AssertionError gat(X)

                        # Test with transposed features
                        Xt = rand(N, in_channel)
                        fgt = FeaturedGraph(adj_gat, transpose(Xt))
                        fgt_ = gat(fgt)
                        @test size(node_feature(fgt_)) == (out_channel * heads, N)

                        g = Zygote.gradient(x -> sum(node_feature(gat(x))), fg)[1]
                        @test size(g[].nf) == size(X)

                        g = Zygote.gradient(model -> sum(node_feature(model(fg))), gat)[1]
                        @test size(g.weight) == size(gat.weight)
                        @test size(g.bias) == size(gat.bias)
                        @test size(g.a) == size(gat.a)
                    end
                end

                @testset "concat=false" begin
                    for heads = [1, 6]
                        gat = GATConv(in_channel=>out_channel, heads=heads, concat=false)
                        @test size(gat.weight) == (out_channel * heads, in_channel)
                        @test size(gat.bias) == (out_channel * heads,)
                        @test size(gat.a) == (2*out_channel, heads)

                        fg_ = gat(fg)
                        Y = node_feature(fg_)
                        @test size(Y) == (out_channel, N)
                        @test_throws AssertionError gat(X)

                        # Test with transposed features
                        Xt = rand(N, in_channel)
                        fgt = FeaturedGraph(adj_gat, transpose(Xt))
                        fgt_ = gat(fgt)
                        @test size(node_feature(fgt_)) == (out_channel, N)

                        g = Zygote.gradient(x -> sum(node_feature(gat(x))), fg)[1]
                        @test size(g[].nf) == size(X)

                        g = Zygote.gradient(model -> sum(node_feature(model(fg))), gat)[1]
                        @test size(g.weight) == size(gat.weight)
                        @test size(g.bias) == size(gat.bias)
                        @test size(g.a) == size(gat.a)
                    end
                end
            end
        end
    end

    @testset "GatedGraphConv" begin
        num_layers = 3
        @testset "layer with graph" begin
            ggc = GatedGraphConv(adj, out_channel, num_layers)
            @test graph(ggc.fg) == [[2,4], [1,3], [2,4], [1,3]]
            @test size(ggc.weight) == (out_channel, out_channel, num_layers)

            X = rand(in_channel, N)
            Y = ggc(X)
            @test size(Y) == (out_channel, N)

            # Test with transposed features
            Xt = rand(N, in_channel)
            Y = ggc(transpose(Xt))
            @test size(Y) == (out_channel, N)

            g = Zygote.gradient(x -> sum(ggc(x)), X)[1]
            @test size(g) == size(X)

            g = Zygote.gradient(model -> sum(model(X)), ggc)[1]
            @test size(g.weight) == size(ggc.weight)
        end

        @testset "layer without graph" begin
            ggc = GatedGraphConv(out_channel, num_layers)
            @test size(ggc.weight) == (out_channel, out_channel, num_layers)

            X = rand(in_channel, N)
            fg = FeaturedGraph(adj, X)
            fg_ = ggc(fg)
            @test size(node_feature(fg_)) == (out_channel, N)
            @test_throws AssertionError ggc(X)

            # Test with transposed features
            Xt = rand(N, in_channel)
            fgt = FeaturedGraph(adj, transpose(Xt))
            fgt_ = ggc(fgt)
            @test size(node_feature(fgt_)) == (out_channel, N)

            g = Zygote.gradient(x -> sum(node_feature(ggc(x))), fg)[1]
            @test size(g[].nf) == size(X)

            g = Zygote.gradient(model -> sum(node_feature(model(fg))), ggc)[1]
            @test size(g.weight) == size(ggc.weight)
        end
    end

    @testset "EdgeConv" begin
        @testset "layer with graph" begin
            ec = EdgeConv(adj, Dense(2*in_channel, out_channel))
            @test graph(ec.fg) == [[2,4], [1,3], [2,4], [1,3]]

            X = rand(in_channel, N)
            Y = ec(X)
            @test size(Y) == (out_channel, N)

            # Test with transposed features
            Xt = rand(N, in_channel)
            Y = ec(transpose(Xt))
            @test size(Y) == (out_channel, N)

            g = Zygote.gradient(x -> sum(ec(x)), X)[1]
            @test size(g) == size(X)

            g = Zygote.gradient(model -> sum(model(X)), ec)[1]
            @test size(g.nn.W) == size(ec.nn.W)
            @test size(g.nn.b) == size(ec.nn.b)
        end

        @testset "layer without graph" begin
            ec = EdgeConv(Dense(2*in_channel, out_channel))

            X = rand(in_channel, N)
            fg = FeaturedGraph(adj, X)
            fg_ = ec(fg)
            @test size(node_feature(fg_)) == (out_channel, N)
            @test_throws AssertionError ec(X)

            # Test with transposed features
            Xt = rand(N, in_channel)
            fgt = FeaturedGraph(adj, transpose(Xt))
            fgt_ = ec(fgt)
            @test size(node_feature(fgt_)) == (out_channel, N)

            g = Zygote.gradient(x -> sum(node_feature(ec(x))), fg)[1]
            @test size(g[].nf) == size(X)

            g = Zygote.gradient(model -> sum(node_feature(model(fg))), ec)[1]
            @test size(g.nn.W) == size(ec.nn.W)
            @test size(g.nn.b) == size(ec.nn.b)
        end
    end
end


text:
using Pkg
Pkg.UPDATED_REGISTRY_THIS_SESSION[] = true   # Tell Julia package manager not to download an updated registry for
Pkg.activate(".")
Pkg.instantiate()  # Not needed if we've already installed packages into this project when we prepared the inputs


text:
using SimplePCHIP

# function lineStyles = linspecer(N)
# This function creates an Nx3 array of N [R B G] colors
# These can be used to plot lots of lines with distinguishable & nice
# looking colors.
#
# lineStyles = linspecer(N);  makes N colors for you to use: lineStyles[ii,:]
#
# colormap(linspecer); set your colormap to have easily distinguishable
#                      colors & a pleasing aesthetic
#
# lineStyles = linspecer(N,"qualitative"); forces the colors to all be distinguishable [up to 12]
# lineStyles = linspecer(N,"sequential"); forces the colors to vary along a spectrum
#

function linspecer(n=128)

    if n<=0 # its empty; nothing else to do here
        lineStyles=[]
        return
    end

    if n==1
        return [0.2005 0.5593 0.7380]
    end
    if n==2
        return  [0.2005 0.5593 0.7380; 0.9684 0.4799 0.2723]
    end

    frac=.85; # Slight modification from colorbrewer here to make the yellows in the center just a bit darker
    cmapp = [158 1 66; 213 62 79; 244 109 67; 253 174 97; 254 224 139; 255*frac 255*frac 191*frac; 230 245 152; 171 221 164; 102 194 165; 50 136 189; 94 79 162]
    x = range(1, n, length=size(cmapp,1))
    xi = 1:n
    cmap = zeros(n,3)
    for ii=1:3
        cmap[:,ii] = [SimplePCHIP.interpolate(x,cmapp[:,ii])(xii) for xii in xi]
    end
    cmap = reverse(cmap/255, dims = 1)

    return cmap
end

max_stars_repo_name: rhoadesScholar/Drugowitsch-Rotation

max_issues_repo_name: rhoadesScholar/Drugowitsch-Rotation

max_forks_repo_name: rhoadesScholar/Drugowitsch-Rotation


text:
## extend logical operators so that we can have a notation like

#f(x,y) < 0 or f(x,y) == 0

#import Base: <, <=, ==, !=, !==, >=, >
import Base: &, |, !

abstract type Predicate end

"""

A predicate is defined in terms of a function of two variables, an
inquality, and either another function or a real number.  They are
conveniently created by the functions `Lt`, `Le`, `Eq`, `Ne`, `Ge`,
and `Gt`. The equivalent unicode operators:

* `≪` (`\\ll[tab]`),
* `≦` (`\\leqq[tab]`),
* `⩵` (`\\Equal[tab]`),
* `≶` (`\\lessgtr[tab]`)  or `≷` (`\\gtrless[tab]`),
* `≧` (`\\geqq[tab]`),
* `≫` (`\\leqq[tab]`) may also be used.

The use of Julia's usual comparison operators is no longer supported.

To combine predicates, `&` and `|` can be used.

To negate a predicate, `!` is used.

"""
mutable struct Pred <: Predicate
    f
    op
    val
end

## meta these
preds = [(:Lt, :≪, :<), # \ll
         (:Le, :≦, :<=), # \leqq
         (:Eq, :⩵, :(==)), # \Equal
         (:Ge, :≧, :>=), # \gegg
         (:Gt, :≫, :>) # \gg
         ]

    # type piracy of CommonEq
    for (fn,uop, op) ∈ (
        (:Lt, :≪, :<), # \ll
        (:Le, :≦, :<=), # \leqq
        (:Eq, :⩵, :(==)), # \Equal
        (:Ge, :≧, :>=), # \gegg
        (:Gt, :≫, :>) # \gg
    )
    fnname =  string(fn)
    @eval begin
        (CommonEq.$fn)(f::Function, x::Real) = Pred(f, $op, x)
        (CommonEq.$fn)(f::Function, g::Function) = $(fn)((x,y) -> f(x,y) - g(x,y), 0)
    end
    eval(Expr(:export, fn))
end

CommonEq.Ne(f::Function, x::Real) = Pred(f, !== , x)
CommonEq.Ne(f::Function, g::Function) = Ne((x,y) -> f(x,y) - g(x,y), 0)
const Neq = Ne # deprecate Neq

"""

Predicates can be joined together with either `&` or `|`. Individual
predicates can be negated with `!`. The parsing rules require the
individual predicates to be enclosed with parentheses, as in `(f==0) | (g==0)`.

"""
mutable struct Preds <: Predicate
    ps
    ops
end

## Some algebra for Pred and Preds
(&)(r1::Pred, r2::Pred) = Preds([r1,r2], Any[&])
(|)(r1::Pred, r2::Pred) = Preds([r1,r2], Any[|])

(&)(ps::Preds, r1::Pred) = Preds(vcat(ps.ps, r1), vcat(ps.ops, &))
(&)(r1::Pred, ps::Preds) = ps & r1
(|)(ps::Preds, r1::Pred) = Preds(vcat(ps.ps, r1), vcat(ps.ops, |))
(|)(r1::Pred, ps::Preds) = ps | r1


text:
using Pkg
Pkg.add("Distributions")
Pkg.add("Random")
Pkg.add("Combinatorics")
Pkg.add("DataFrames")
using Distributions
using Random
using Combinatorics
using DataFrames

function Amat(s,c,alpha)
    """
    Genereate adjacency matrix of size s,
    with connectance c
    and interaction strength alpha
    """
A = zeros(s, s)

for i in 1:s
    for j in 1:s
    if rand(Uniform(0,1))< c/2
        A[i,j] = rand(Normal(0,alpha))
    end
end
end

return A
end

function Board(s)
    """
    Create s*s matrix for simulation
    to take place on
    """
    return zeros(Int, s, s)
end

function intc(x)
    """
    get coordinates for interactions
    """
    intCoords = []
    for i in -1:1
        for j in -1:1
            if i * j != i+j
                push!(intCoords, (i,j))
            end
        end
    end
    return intCoords
end

function fitnessModifier(board, A, intCoords,focal)
    """
    Calculate the change in fitness of focal individual
    due to local interactions
    """
    modifier = 0

    partners = [Int(board[p[1],p[2]])  for p in intCoords]

    for partner in partners
        if partner != 0
        modifier += A[focal,partner] #
    end
    end
    return modifier
end

function setBoard(board, density, nSpecies)
    """
    set out initial species distributions
    """

    for i in size(board)[1]
        for j in size(board)[1]
        if rand(Uniform(0,1))< density
            board[i,j] = rand(1:nSpecies)
        end
    end
    end
    return board
end
function checkSpace(board, intCoords)

    if sum([board[i[1],i[2]] == 0 for i in intCoords])!= 0
        return true
    else
        return false
    end
end

function Birth(board, intCoords,focal)

coords = [(i[1],i[2]) for i in intCoords if board[i[1],i[2]] == 0]
coord = coords[rand(1:length(coords))]
board[coord[1],coord[2]] = focal
return board

end

function cull(board, maxDens)
    """
    density regulation
    """
    Size = size(board)[1]

    density = sum([i != 0 for i in board])# current density
    while density > maxDens # if/while density is too high
        board[rand(1:Size),rand(1:Size)] = 0 #kill at random
        density = sum([i != 0 for i in board])# current density
    end
    return board
end

function disperse(board, sections)
    """
    Shuffle the peices on the board
    """
    Size = size(board)[1]
    deme = Size ÷ sections # length of individual demes
    demes = [i for i in 1:deme:Size] # list of starting points for demes
    for i in demes
        for j in demes
            # shuffle the individual demes
            #each deme is the value from the list + the deme length
            board[i: i+deme-1,j:j+deme-1] = shuffle(board[i: i+deme-1,j:j+deme-1])
        end
    end
    return board
end

function sim(c, nSpecies, density, boardSize, dispersal, demes, alpha,bLim,maxDens )
    # set up board and adjacency matrix
    board = Board(boardSize)
    board = setBoard(board, density, nSpecies)
    Amatrix = Amat(nSpecies, c, alpha)
    intCoords = intc(1)
    moves = combinations(hcat(1:boardSize,1:boardSize),2) |> collect

    # initialise counters
    births = 0
    richness = length(unique(board))
    gens = 0

    #run the simulation
    while births < bLim
        gens+=1
        #get randomly ordered vector of board indeces
        for (i,j) in shuffle(unique(moves))
                if board[i,j] != 0 #not empty
                    #evaluate fitness
                    focal = board[i,j] # species id

                    locality = [(i,j) .+ ic for ic in intCoords]

                    locality = locality[[l for l in 1:8 if checkbounds(Bool,board,locality[l][1],locality[l][2])]]

                    if rand(Uniform(-30.5,-30))< fitnessModifier(board, Amatrix, locality,focal)
                        # then give birth if space available
                        if checkSpace(board, locality) == true
                        board = Birth(board,locality,focal)
                        births += 1
                        end
                    end
                end

        end
        #dispersal
        if rand(Uniform(0,1)) < dispersal
            board = disperse(board, demes)
        end
        #density regulation
        board = cull(board,maxDens)

        # check for extinctions
        if length(unique(board)) < richness # if an extinction has occurred
            richness = length(unique(board)) #update richness
            births = 0 # reset counter
        end
    end
        # return data!
        return [density c nSpecies boardSize dispersal demes^2 alpha maxDens richness gens]
end

function threadSim(c, nSpecies, density, boardSize, dispersal, demes, alpha,bLim,maxDens)
# initialise dictionary for storing data
D = Dict()
print(Threads.nthreads())
Threads.@threads for i in 1:Threads.nthreads()
    D[Threads.threadid()] = sim(c, nSpecies, density, boardSize, dispersal, demes, alpha,bLim,maxDens)
end
return [vcat(V) for (K, V) in D ]
end

#x = sim(0.5, 100, 0.5, 100, 1, 4, 1,10,800 )
#x
y = threadSim(0.5, 100, 0.5, 100, 1, 4, 1,10,800 )
y

# run the project
dat = Array{Float64}(undef, 0, 10)
for c in LinRange(0.1,1,10)
    for nSpecies in [50,100,200]
        for density in [0.2,0.5,0.8]
            for boardSize in [100,1000]
                for dispersal in LinRange(0.1,1,10)
                    for demes in [boardSize ÷ d for d in 1:boardSize if boardSize/d == boardSize ÷ d]
                        for alpha in LinRange(0.1,2,20)
                            for bLim in boardSize^2 .* [0.1,0.5,1,2,3,]

                                dat = vcat(dat,threadSim(c, nSpecies, density, boardSize, dispersal, demes, alpha, bLim ,density))
                            end
                        end
                    end
                end
                df = DataFrame(dat)
                names!(df, [:c, :nSpecies, :density, :boardSize, :dispersal, :demes, :alpha,:bLim ,:maxDens, :gens])
                CSV.write("space.csv", df)

            end
        end
    end
end


text:
"""
Analytic axisymmetric small tokamak equilibrium in circular coordinates.
"""
module TokamakSmallToroidal

    import ElectromagneticFields.AxisymmetricTokamakToroidal

    export initial_conditions_barely_passing, initial_conditions_barely_trapped,
           initial_conditions_deeply_passing, initial_conditions_deeply_trapped,
           initial_conditions_pauli

    export hamiltonian, toroidal_momentum

    AxisymmetricTokamakToroidal.@code() # inject magnetic field code

    const xᵢ = [1.05, 0., 0.]
    const uᵢ = -0.00045135897235326736
    const qᵢ = [from_cartesian(0, xᵢ)..., uᵢ]

    const parameters = (μ = 2.314593645825811e-6,)

    initial_conditions_barely_passing() = ([1.05, 0., 0., 8.117E-4], (μ = 2.448E-6,))
    initial_conditions_barely_trapped() = ([1.05, 0., 0., 7.610E-4], (μ = 2.250E-6,))
    initial_conditions_deeply_passing() = ([1.05, 0., 0., 1.623E-3], (μ = 2.448E-6,))
    initial_conditions_deeply_trapped() = ([1.05, 0., 0., 4.306E-4], (μ = 2.250E-6,))
    initial_conditions_pauli() = ([1.05, 0., 0., 4.3E-4], (μ = 2.310E-6,))

    u_loop() = 4.0E-4
    μ_loop() = 2.5E-6
    u_surface() = 4.0E-4
    μ_surface() = 2.5E-6

    function f_loop(t)
        R0 = 1.0
        Z0 = 0.0
        φ0 = 0.0
        u0 = u_loop()
        r0 = 0.05

        Rt = R0 + r0*cos(2π*t)
        Zt = Z0 + r0*sin(2π*t)

        qt = [Rt, Zt, φ0, u0]

        return qt
    end

    function f_surface(s,t)
        R0 = 1.0
        Z0 = 0.0
        φ0 = 0.0
        u0 = u_surface()
        r0 = 0.1

        Rt = R0 + 2r0*(s-0.5)
        Zt = Z0 + 2r0*(t-0.5)

        qt = [Rt, Zt, φ0, u0]

        return qt
    end

    include("guiding_center_4d_common.jl")
    include("guiding_center_4d_equations.jl")
    include("guiding_center_4d_loop.jl")
    include("guiding_center_4d_surface.jl")

    function toroidal_momentum(t,q)
        R(t,q) * ϑ3(t,q)
    end

    include("guiding_center_4d_diagnostics.jl")

end


text:
import CollectiveSpins
using PyPlot

# Band structure (Omega) for 2D square lattice for in-plane circularly polarized
# and perpendicular (x-axis) polarized atoms.
# See Shahmoon et al 10.1103/PhysRevLett.118.113601 Fig. 4a.

polarization_par = [0, 1, im]
polarization_perp = [1, 0, 0]
a_vec1 = [0.2, 0]
a_vec2 = [0, 0.2]
k_max = 200

k_listy = [iii*pi/(a_vec1[1]*k_max) for iii=k_max:-1:0]
k_listz = copy(k_listy)
k_listy = append!(k_listy, [iii*pi/(a_vec1[1]*k_max) for iii=0:k_max])
k_listz = append!(k_listz, zeros(k_max+1))
k_listy = append!(k_listy, ones(k_max+1)*pi/a_vec1[1])
k_listz = append!(k_listz, [iii*pi/(a_vec1[1]*k_max) for iii=0:k_max])
k_dim = length(k_listy)

bandstructure_par = zeros(k_dim)
bandstructure_perp = zeros(k_dim)
for iii=1:k_dim
    k = [k_listy[iii], k_listz[iii]]
    bandstructure_par[iii] = CollectiveSpins.collective_modes.Omega_k_2D(k, a_vec1, a_vec2, polarization_par)
    bandstructure_perp[iii] = CollectiveSpins.collective_modes.Omega_k_2D(k, a_vec1, a_vec2, polarization_perp)
end

positions = [0, k_max, 2*k_max, 3*k_max]
labels = ["M", "Gamma", "X", "M"]

plot(collect(1:k_dim), bandstructure_par, "k--")
plot(collect(1:k_dim), bandstructure_perp, color="gold")
xticks(positions, labels)
ylim([-5.2, 5.2])

tight_layout()
savefig("band_structure.svg")

# Collective decay rate (Gamma) for 2D square lattice for y-axis (in-plane)
# and x-axis (perpendicular to plane) polarized atoms.
# See Asenjo-Garcia et al 10.1103/PhysRevX.7.031024 Figs. 2c and 2d.

polarization_par = [0, 1, 0]
polarization_perp = [1, 0, 0]
a_vec1 = [0.2, 0]
a_vec2 = [0, 0.2]
k_max = 200

k_list = [iii*pi/(a_vec1[1]*k_max) for iii=-k_max:k_max]
k_dim = length(k_list)
decay_par = zeros(k_dim, k_dim)
decay_perp = zeros(k_dim, k_dim)
y = 1
for ky in k_list
    z = 1
    for kz in k_list
        k = [ky, kz]
        decay_par[z, y] = CollectiveSpins.collective_modes.Gamma_k_2D(k, a_vec1, a_vec2, polarization_par)
        decay_perp[z, y] = CollectiveSpins.collective_modes.Gamma_k_2D(k, a_vec1, a_vec2, polarization_perp)
        z += 1
    end
    global y += 1
end

positions = [0, k_max, 2*k_max]
labels = [-1, 0, 1]

subplot(211)
imshow(decay_par, cmap="jet")
colorbar()
clim(0, 22)
xlabel("k_y d/pi")
ylabel("k_z d/pi")
plt.xticks(positions, labels)
plt.yticks(positions, labels)

subplot(212)
imshow(decay_perp, cmap="jet")
colorbar()
clim(0, 22)
xlabel("k_y d/pi")
ylabel("k_z d/pi")
plt.xticks(positions, labels)
plt.yticks(positions, labels)

tight_layout()
savefig("decay.svg")


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

name = "finufft"
version = v"1.1.2"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/flatironinstitute/finufft.git", "5b92daf4244c92844dd2640ef80457c030e1bf25")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/finufft/
# Overwrite LIBSFFT such that we do not require fftw3_threads or fftw3_omp for OMP support. Since the libraries in FFTW_jll already provide for threading, we do not loose anything.
make lib -j${nproc} CFLAGS="-fPIC -O3 -funroll-loops -fcx-limited-range" LIBSFFT="-lfftw3 -lfftw3f -lm"
cp lib/libfinufft.so "${libdir}/libfinufft.${dlext}"
"""

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

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

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="FFTW_jll", uuid="f5851436-0d7a-5f13-b9de-f02708fd171a")),
    Dependency(PackageSpec(name="CompilerSupportLibraries_jll", uuid="e66e0078-7015-5450-92f7-15fbd957f2ae"))
]

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

max_stars_repo_name: mortenpi/Yggdrasil

max_issues_repo_name: mortenpi/Yggdrasil

max_forks_repo_name: mortenpi/Yggdrasil


text:
# Default tzdata version to use if the environmental variable "JULIA_TZ_VERSION" is not set.
# We want to use a specific version here to ensure that specific revisions of the
# TimeZones.jl package always use the same revision of tzdata. Doing so ensure that we can
# always use older revisions of this package and always reproduce the same results.
const DEFAULT_TZDATA_VERSION = "2021c"  # Do not use floating revision "latest" here

# Note: A tz code or data version consists of a year and letter while a release consists of
# a pair of tz code and data versions. In recent releases the tz code and data use the same
# version.
#
# > Typically a release R consists of a pair of tarball files,
# > tzcodeR.tar.gz and tzdataR.tar.gz.  However, some releases (e.g.,
# > code2010a, data2012c) consist of just one or the other tarball, and a
# > few (e.g., code2012c-data2012d) have tarballs with mixed version
# > numbers.
#
# ―NEWS file (tzdata2017a)

# Parse the tzdata version from things such as the archive filename
const TZDATA_VERSION_REGEX = r"(?<!\d)(?:\d{2}){1,2}[a-z]?\b"

# Parse release lines from the NEWS file
const TZDATA_NEWS_REGEX = r"""
    ^Release\s+
    (?:code(?:\d{2}){1,2}[a-z]?-)?
    (?:data)?
    (?<version>(?:\d{2}){1,2}[a-z]?)
    \b
"""x

const ACTIVE_VERSION_FILE = joinpath(DEPS_DIR, "active_version")

"""
    read_news(news, [limit]) -> Vector{AbstractString}

Reads all of the tzdata versions from the NEWS file in the order in which they appear. Note
that since the NEWS file is in reverse chronological order the versions will also be in that
order. Useful for identifying the version of the tzdata.
"""
function read_news(news::IO, limit::Integer=0)
    count = 0
    revs = sizehint!(AbstractString[], limit)
    while !eof(news) && (limit == 0 || count < limit)
        line = readline(news)

        m = match(TZDATA_NEWS_REGEX, line)
        if m !== nothing
            push!(revs, m[:version])
            count += 1
        end
    end
    return revs
end

function read_news(news::AbstractString, limit::Integer=0)
    open(news, "r") do fp
        read_news(fp, limit)
    end
end

"""
    tzdata_version_dir(dir::AbstractString) -> AbstractString

Determines the tzdata version by inspecting various files in a directory.
"""
function tzdata_version_dir(dir::AbstractString)
    return cd(dir) do
        if isfile("version")  # Added in release 2016h
            readchomp("version")
        elseif isfile("NEWS")
            # Find the archive version by determining the latest version in the change log
            first(read_news("NEWS", 1))  # Added in release 2014g
        else
            error("Unable to determine tzdata version")
        end
    end
end

"""
    tzdata_version_archive(archive::AbstractString) -> AbstractString

Determines the tzdata version by inspecting the contents within the archive. Useful when
downloading the latest archive "tzdata-latest.tar.gz".
"""
function tzdata_version_archive(archive::AbstractString)
    # Attempting to extract files that do not exist in the archive will result in an
    # exception.
    files = readarchive(archive)
    available_files = intersect(Set(["NEWS", "version"]), Set(files))
    isempty(available_files) && error("Unable to determine tzdata release")

    mktempdir() do temp_dir
        extract(archive, temp_dir, available_files)
        tzdata_version_dir(temp_dir)
    end
end

tzdata_version() = get(ENV, "JULIA_TZ_VERSION", DEFAULT_TZDATA_VERSION)

function active_version()
    if !isfile(ACTIVE_VERSION_FILE)
        error("No active tzdata version. Try re-building TimeZones")
    end
    read(ACTIVE_VERSION_FILE, String)
end

function active_archive()
    version = active_version()
    archive = joinpath(ARCHIVE_DIR, "tzdata$version.tar.gz")
    !isfile(archive) && error("Missing $version tzdata archive")
    return archive
end


text:
## Pendulum
# https://github.com/HarvardAgileRoboticsLab/unscented-dynamic-programming/blob/master/pendulum_dynamics.m
function pendulum_dynamics!(ẋ::AbstractVector{T},x::AbstractVector{T},u::AbstractVector{T}) where T
    m = 1.
    l = 0.5
    b = 0.1
    lc = 0.5
    I = 0.25
    g = 9.81
    ẋ[1] = x[2]
    ẋ[2] = (u[1] - m*g*lc*sin(x[1]) - b*x[2])/I
end

n,m = 2,1
pendulum = Model(pendulum_dynamics!,n,m) # inplace model

function pendulum_dynamics2!(ẋ::AbstractVector{T},x::AbstractVector{T},u::AbstractVector{T}) where T
    ẋ[1] = x[2];
    ẋ[2] = sin(x[1]) + u[1];

end

pendulum2 = Model(pendulum_dynamics2!,n,m) # inplace model

# unknown mass
function pendulum_dynamics_uncertain!(ẋ,x,u,w)
    m = 1. + w[1]
    l = 0.5
    b = 0.1
    lc = 0.5
    I = 0.25
    g = 9.81

    ẋ[1] = x[2]
    ẋ[2] = u[1]/(m*lc*lc) - g*sin(x[1])/lc - b*x[2]/(m*lc*lc)
    return nothing
end

n = 2; m = 1; r = 1
pendulum_uncertain = UncertainModel(pendulum_dynamics_uncertain!,n,m,r)

max_stars_repo_name: MehrdadZareian/GMARA

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
"""
The `Format` module provides a public interface that can be used to define custom formatters
aside from the predefined HTML and LaTeX outputs supported by `Highlights`.

The [Formatting](@ref) section of the manual provides a example of how to go about extending
`Highlights` to support additional output formats.

The following functions and types are exported from the module for public use:

$(EXPORTS)
"""
module Format

using DocStringExtensions

import ..Highlights
import ..Highlights.Compiler: Context
import ..Highlights.Themes: RGB, Style, Theme
import ..Highlights.Tokens

export TokenIterator, render

"""
The `render` function is used to define custom output formats for tokenised source code.
Methods with signatures of the form

```julia
function Format.render(io::IO, mime::MIME, tokens::Format.TokenIterator)
    # ...
    for (str, id, style) in tokens
        # ...
    end
    # ...
end
```

can be defined for any `mime` type to allow the [`highlight`](@ref Highlights.highlight)
function to support new output formats.
"""
function render end

# RGB colours.

render(io::IO, ::MIME"text/css", c::RGB) = (print(io, "rgb("); _tup(io, c); print(io, ")"))
render(io::IO, ::MIME"text/latex", c::RGB) = (print(io, "[RGB]{"); _tup(io, c); print(io, "}"))

_tup(io::IO, c::RGB) = print(io, Int(c.r), ",", Int(c.g), ",", Int(c.b))

# Styles.

function render(io::IO, mime::MIME"text/css", style::Style)
    if style.fg.active
        print(io, "color: ")
        render(io, mime, style.fg)
        print(io, "; ")
    end
    if style.bg.active
        print(io, "background-color: ")
        render(io, mime, style.bg)
        print(io, "; ")
    end
    style.bold      && print(io, "font-weight: bold; ")
    style.italic    && print(io, "font-style: italic; ")
    style.underline && print(io, "text-decoration: underline; ")
    return nothing
end

function render(io::IO, mime::MIME"text/latex", style::Style)
    print(io, "[1]{")
    if style.fg.active
        print(io, "\\textcolor")
        render(io, mime, style.fg)
        print(io, "{")
    end
    if style.bg.active
        print(io, "\\colorbox")
        render(io, mime, style.bg)
        print(io, "{")
    end
    style.bold      && print(io, "\\textbf{")
    style.italic    && print(io, "\\textit{")
    style.underline && print(io, "\\underline{")
    print(io, "#1")
    style.underline && print(io, "}")
    style.italic    && print(io, "}")
    style.bold      && print(io, "}")
    style.bg.active && print(io, "}")
    style.fg.active && print(io, "}")
    print(io, "}")
end

# Stylesheets.

function render(io::IO, mime::MIME"text/css", theme::Theme)
    print(io,
        """
        pre.hljl {
            border: 1px solid #ccc;
            margin: 5px;
            padding: 5px;
            overflow-x: auto;
            """
    )
    render(io, mime, theme.base)
    println(io, "}")
    for (nth, style) in enumerate(theme.styles)
        print(io, "pre.hljl > span.hljl-", Tokens.__SHORTNAMES__[nth], " { ")
        render(io, mime, style)
        println(io, "}")
    end
end

function render(io::IO, ::MIME"text/html", theme::Theme)
    println(io, "\n\n")
end

function render(io::IO, mime::MIME"text/latex", theme::Theme)
    println(io, "\\usepackage{upquote}")
    println(io, "\\usepackage{listings}")
    println(io, "\\usepackage{xcolor}")
    print(io,
        """
        \\lstset{
            basicstyle=\\ttfamily\\footnotesize,
            upquote=true,
            breaklines=true,
            breakindent=0pt,
            keepspaces=true,
            showspaces=false,
            columns=fullflexible,
            showtabs=false,
            showstringspaces=false,
            escapeinside={(*@}{@*)},
            extendedchars=true,
        }
        """
    )
    for (nth, style) in enumerate(theme.styles)
        print(io, "\\newcommand{\\HLJL", Tokens.__SHORTNAMES__[nth], "}")
        render(io, mime, style)
        println(io)
    end
end

# Token Iterator.

"""
$(TYPEDEF)

An iterator type used in user-defined [`render`](@ref) methods to provide custom output
formats. This type supports the basic Julia iterator protocol: namely `iterate` which enables `for`-loop interation over tokenised text.

The iterator returns a 3-tuple of `str`, `id`, and `style` where

  * `str` is the `SubString` covered by the token;
  * `id` is the shorthand name of the token type as a `Symbol`;
  * `style` is the `Style` applied to the token.

"""
struct TokenIterator
    ctx::Context
    theme::Theme
end

Base.iterate(t::TokenIterator) = isempty(t.ctx.tokens) ? nothing : iterate(t::TokenIterator, 1)
Base.length(t::TokenIterator) = length(t.ctx.tokens)

function Base.iterate(t::TokenIterator, state)
    (state > length(t)) && return nothing

    token = t.ctx.tokens[state]
    result = (
        SubString(t.ctx.source, token.first, token.last),
        Tokens.__SHORTNAMES__[token.value.value],
        t.theme.styles[token.value.value],
    )

    return result, state + 1
end

# Code blocks.

render(io::IO, mime::MIME, ctx::Context, theme::Theme) =
    render(io, mime, TokenIterator(ctx, theme))

function render(io::IO, mime::MIME"text/html", tokens::TokenIterator)
    println(io, "
")
    for (str, id, style) in tokens
        print(io, "")
        escape(io, mime, str)
        print(io, "")
    end
    println(io, "\n
")
end

const CONSECUTIVE_WHITESPACE = r"\s+"

function print_formatted(io::IO,mime::MIME"text/latex",str,id,style)
    id === :t || print(io, "(*@\\HLJL", id, "{")
    escape(io, mime, str; charescape=(id === :t))
    id === :t || print(io, "}@*)")
end

function render_nonwhitespace(io::IO,mime::MIME"text/latex",str,id,style)
    # Whitespace chars within an escapeinside are not correctly printed in a
    # lstlisting env. So in order to have the correct highlighting
    # and correct characters, that are recognized by listings, they need to be
    # added outside of the escapeinside.
    offset = 1
    for m in eachmatch(CONSECUTIVE_WHITESPACE,str)
        whitespace = m.match
        nonwhitespace = str[offset:prevind(str,m.offset)]
        offset = nextind(str,m.offset,length(whitespace))
        length(nonwhitespace) > 0 && print_formatted(io,mime,nonwhitespace,id,style)
        print(io,whitespace)
    end
    nonwhitespace = str[offset:end]
    length(nonwhitespace) > 0 && print_formatted(io,mime,str[offset:end],id,style)
end

function render(io::IO, mime::MIME"text/latex", tokens::TokenIterator)
    println(io, "\\begin{lstlisting}")
    for (str, id, style) in tokens
        render_nonwhitespace(io,mime,str,id,style)
    end
    println(io, "\n\\end{lstlisting}")
end

# Character escapes.

function escape(io::IO, ::MIME"text/html", str::AbstractString)
    for char in str
        char === '&'  ? print(io, "&") :
        char === '<'  ? print(io, "<") :
        char === '>'  ? print(io, ">") :
        char === '"'  ? print(io, """) :
        char === '\'' ? print(io, "'") :
            print(io, char)
    end
end

function escape(io::IO, ::MIME"text/latex", str::AbstractString; charescape=false)
    for char in str
        char === '`'   ? printe(io, charescape, "{\\textasciigrave}") :
        char === '\''  ? printe(io, charescape, "{\\textquotesingle}") :
        char === '$'   ? printe(io, charescape, "{\\\$}") :
        char === '%'   ? printe(io, charescape, "{\\%}") :
        char === '#'   ? printe(io, charescape, "{\\#}") :
        char === '&'   ? printe(io, charescape, "{\\&}") :
        char === '\\'  ? printe(io, charescape, "{\\textbackslash}") :
        char === '^'   ? printe(io, charescape, "{\\textasciicircum}") :
        char === '_'   ? printe(io, charescape, "{\\_}") :
        char === '{'   ? printe(io, charescape, "{\\{}") :
        char === '}'   ? printe(io, charescape, "{\\}}") :
        char === '~'   ? printe(io, charescape, "{\\textasciitilde}") :
        char === '"'   ? printe(io, charescape, "\"{}") :
            print(io, char)
    end
end

function printe(io::IO, charescape::Bool, s)
    charescape && print(io, "(*@{")
    print(io, s)
    charescape && print(io, "}@*)")
end

end # module