text:
using Oceananigans.Diagnostics: AbstractDiagnostic
using Oceananigans.Fields: FieldSlicer
using Oceananigans.OutputWriters: fetch_output
using Oceananigans.Utils: AbstractSchedule, prettytime

import Oceananigans: run_diagnostic!
import Oceananigans.Utils: TimeInterval
import Oceananigans.Fields: location

"""
    mutable struct AveragedTimeInterval <: AbstractSchedule

Container for parameters that configure and handle time-averaged output.
"""
mutable struct AveragedTimeInterval <: AbstractSchedule
                       interval :: Float64
                         window :: Float64
                         stride :: Int
    previous_interval_stop_time :: Float64
                     collecting :: Bool
end

"""
    AveragedTimeInterval(interval; window=interval, stride=1)

Returns a `schedule` that specifies periodic time-averaging of output.
The time `window` specifies the extent of the time-average, which
reoccurs every `interval`.

`output` is computed and accumulated into the average every `stride` iterations
during the averaging window. For example, `stride=1` computs output every iteration,
whereas `stride=2` computes output every other iteration. Time-averages with
longer `stride`s are faster to compute, but less accurate.

The time-average of ``a`` is a left Riemann sum corresponding to

`` ⟨a⟩ = 1/T \\int_{tᵢ-T}^T a \\mathrm{d} t ,``

where ``⟨a⟩`` is the time-average of ``a``, ``T`` is the time-window for averaging,
and the ``tᵢ`` are discrete times separated by the time `interval`. The ``tᵢ`` specify
both the end of the averaging window and the time at which output is written.

Example
=======

```jldoctest averaged_time_interval
using Oceananigans.OutputWriters: AveragedTimeInterval
using Oceananigans.Utils: year, years

schedule = AveragedTimeInterval(4years, window=1year)

# output
AveragedTimeInterval(window=1 year, stride=1, interval=4 years)
```

An `AveragedTimeInterval` schedule directs an output writer
to time-average its outputs before writing them to disk:

```jldoctest averaged_time_interval
using Oceananigans
using Oceananigans.OutputWriters: JLD2OutputWriter
using Oceananigans.Utils: minutes

model = NonhydrostaticModel(grid=RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1)))

simulation = Simulation(model, Δt=10minutes, stop_time=30years)

simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities,
                                                          prefix = "averaged_velocity_data",
                                                          schedule = AveragedTimeInterval(4years, window=1year, stride=2))

# output
JLD2OutputWriter scheduled on TimeInterval(4 years):
├── filepath: ./averaged_velocity_data.jld2
├── 3 outputs: (:u, :v, :w) averaged on AveragedTimeInterval(window=1 year, stride=2, interval=4 years)
├── field slicer: FieldSlicer(:, :, :, with_halos=false)
├── array type: Array{Float32}
├── including: [:grid, :coriolis, :buoyancy, :closure]
└── max filesize: Inf YiB
```
"""
function AveragedTimeInterval(interval; window=interval, stride=1)

    window > interval && throw(ArgumentError("Averaging window $window is greater than the output interval $interval."))

    return AveragedTimeInterval(Float64(interval), Float64(window), stride, 0.0, false)
end

# Determines whether or not to call run_diagnostic
(schedule::AveragedTimeInterval)(model) =
    schedule.collecting || model.clock.time >= schedule.previous_interval_stop_time + schedule.interval - schedule.window

TimeInterval(schedule::AveragedTimeInterval) = TimeInterval(schedule.interval)

Base.copy(schedule::AveragedTimeInterval) = AveragedTimeInterval(schedule.interval, window=schedule.window, stride=schedule.stride)

"""
    WindowedTimeAverage{OP, R, FS} <: AbstractDiagnostic

An object for computing 'windowed' time averages, or moving time-averages
of a `operand` over a specified `window`, collected on `interval`.
"""
mutable struct WindowedTimeAverage{OP, R, FS} <: AbstractDiagnostic
                      result :: R
                     operand :: OP
           window_start_time :: Float64
      window_start_iteration :: Int
    previous_collection_time :: Float64
                field_slicer :: FS
                    schedule :: AveragedTimeInterval
end

"""
    WindowedTimeAverage(operand, model=nothing; schedule, field_slicer=FieldSlicer())

Returns an object for computing running averages of `operand` over `schedule.window` and
recurring on `schedule.interval`, where `schedule` is an `AveragedTimeInterval`.
During the collection period, averages are computed every `schedule.stride` iteration.

`operand` may be a `Oceananigans.Field` or a function that returns an array or scalar.

Calling `wta(model)` for `wta::WindowedTimeAverage` object returns `wta.result`.
"""
function WindowedTimeAverage(operand, model=nothing; schedule, field_slicer=FieldSlicer())

    output = fetch_output(operand, model, field_slicer)
    result = similar(output) # convert views to arrays
    result .= output # initialize `result` with initial output

    return WindowedTimeAverage(result, operand, 0.0, 0, 0.0, field_slicer, schedule)
end

# Time-averaging doesn't change spatial location
location(wta::WindowedTimeAverage) = location(wta.operand)

function accumulate_result!(wta, model)

    # Time increment:
    Δt = model.clock.time - wta.previous_collection_time

    # Time intervals:
    T_current = model.clock.time - wta.window_start_time
    T_previous = wta.previous_collection_time - wta.window_start_time

    # Accumulate left Riemann sum
    integrand = fetch_output(wta.operand, model, wta.field_slicer)

    @. wta.result = (wta.result * T_previous + integrand * Δt) / T_current

    # Save time of integrand collection
    wta.previous_collection_time = model.clock.time

    return nothing
end

function run_diagnostic!(wta::WindowedTimeAverage, model)

    if model.clock.iteration == 0 # initialize previous interval stop time
        wta.schedule.previous_interval_stop_time = model.clock.time
    end

    # Don't start collecting if we are *only* "initializing" run_diagnostic! at the beginning
    # of a Simulation.
    #
    # Note: this can be false at the zeroth iteration if interval == window (which
    # implies we are always collecting)

    unscheduled = model.clock.iteration == 0 &&
        model.clock.time < wta.schedule.previous_interval_stop_time + wta.schedule.interval - wta.schedule.window

    if unscheduled
        # This is an "unscheduled" call to run_diagnostic! --- which occurs when run_diagnostic!
        # is called at the beginning of a run (and schedule.interval != schedule.window).
        # In this case we do nothing.
    elseif !(wta.schedule.collecting)
        # run_diagnostic! has been called on schedule, but we are not currently collecting data.
        # Initialize data collection:

        # Start averaging period
        wta.schedule.collecting = true

        # Zero out result
        wta.result .= 0

        # Save averaging start time and the initial data collection time
        wta.window_start_time = model.clock.time
        wta.window_start_iteration = model.clock.iteration
        wta.previous_collection_time = model.clock.time

    elseif model.clock.time >= wta.schedule.previous_interval_stop_time + wta.schedule.interval
        # Output is imminent. Finalize averages and cease data collection.
        accumulate_result!(wta, model)

        # Averaging period is complete.
        wta.schedule.collecting = false

        # Reset the "previous" interval time, subtracting a sliver that presents overshoot from accumulating.
        wta.schedule.previous_interval_stop_time = model.clock.time - rem(model.clock.time, wta.schedule.interval)

    elseif mod(model.clock.iteration - wta.window_start_iteration, wta.schedule.stride) == 0
        # Collect data as usual
        accumulate_result!(wta, model)
    end

    return nothing
end

function (wta::WindowedTimeAverage)(model)

    # For the paranoid
    wta.schedule.collecting &&
        model.clock.iteration > 0 &&
        @warn "Returning a WindowedTimeAverage before the collection period is complete."

    return wta.result
end

Base.show(io::IO, schedule::AveragedTimeInterval) = print(io, summary(schedule))

Base.summary(schedule::AveragedTimeInterval) = string("AveragedTimeInterval(",
                                                      "window=", prettytime(schedule.window), ", ",
                                                      "stride=", schedule.stride, ", ",
                                                      "interval=", prettytime(schedule.interval),  ")")

show_averaging_schedule(schedule) = ""
show_averaging_schedule(schedule::AveragedTimeInterval) = string(" averaged on ", summary(schedule))

output_averaging_schedule(output::WindowedTimeAverage) = output.schedule


text:
export partial_trace, computational_basis_vectors

export n_product_id

# validation methods
export commutes, is_hermitian, is_positive_semidefinite, is_orthonormal_basis, is_complete

"""
    partial_trace(
        ρ :: AbstractMatrix, dims :: Vector{Int64}, id :: Int64
    ) :: Matrix

Performs the partial trace on matrix `ρ` with respect to the subsystem at the specified `id`
and returns resulting reduced matrix.
The returned matrix is square with dimension equal to the product of `dims` with the element
corresponding to `id` removed.

*Inputs:*
* `ρ` : The matrix on which the partial trace is performed. This matrix should be square
        and have dimension equal to the product of the `dims`.
* `dims` :  A vector containing positive integer elements which specify the
        dimension of each subsystem. *E.g.* A 3-qubit system has `dims = [2,2,2]`,
        a system containing a qubit and a qutrit has `dims = [2,3]`.
* `id` : A positive integer indexing the `dims` vector to signify
        the subsystem to be traced out.

The partial trace can be understood as a quantum operation ``\\mathcal{E}`` mapping the input
Hilbert space to the output Hilbert space,
``\\mathcal{E} \\; : \\; \\mathcal{H}_X \\rightarrow \\mathcal{H}_Y``. A quantum
operation admits the operator sum representation, ``\\mathcal{E}(\\rho) = \\sum_i E_i  \\rho E_i^{\\dagger}``.
For example, given a 3-qubit system with density matrix ``\\rho_{ABC}``, the partial trace with
respect to system ``B``, is explicitly,

```math
\\rho_{AC} = \\text{Tr}_B[\\rho_{ABC}] = \\mathcal{E}_B(\\rho_{ABC}) = \\sum_i E_i \\rho_{ABC} E_i^{\\dagger}
```

where ``E_i = \\mathbb{I}_A \\otimes \\langle i |_B \\otimes\\mathbb{I}_C``, represents
the Kraus operators for the quantum operation and ``\\rho_{AC}`` is the reduced density
operator remaining after system ``B`` is traced out.

A `DomainError` is thrown if `ρ` is not square or of proper dimension.
"""
function partial_trace(ρ::AbstractMatrix, dims::Vector{Int64}, id::Int64) :: Matrix
    dim = *(dims...)

    if size(ρ) != (dim, dim)
        throw(DomainError(ρ, "Matrix ρ should be square and have dimension equal to the product of subsystem dimensions."))
    end

    trace_dim = dims[id]

    pre_dim = (id > 1) ? *(dims[1:(id-1)]...) : 1
    pre_id = Matrix(1I, pre_dim, pre_dim)

    post_dim = (id < length(dims)) ? *(dims[(id + 1):end]...) : 1
    post_id = Matrix(1I, post_dim, post_dim)

    # TODO: This method can be sped up significantly by avoiding kraus operator construction
    # and coputing the partial trace by a more direct means.
    basis_vectors = computational_basis_vectors(trace_dim)
    kraus_operators = map(bv -> kron(pre_id, bv', post_id), basis_vectors)

    sum(map( E -> E*ρ*E', kraus_operators))
end

"""
    computational_basis_vectors( dim::Int64 ) ::Vector{ Vector{Int64} }

Returns the set of orthonormal column vectors ``\\{|1\\rangle,\\dots,|i\\rangle,\\dots,|d\\rangle \\}``
spanning a vector space of dimension `dim`.
Note that ``|1\\rangle = ( 1, 0, \\dots, 0 )^T``.
"""
function computational_basis_vectors(dim::Int64) :: Vector{Vector{Int64}}
    _basis_vec(i) = begin
        vec = zeros(Int64, dim)
        vec[i] = 1
        return vec
    end

    map(_basis_vec, 1:dim)
end

"""
    is_hermitian( M :: AbstractMatrix{<:Number}; atol=ATOL :: Float64 ) :: Bool

Returns `true` if the supplied matrix is hermitian (self-adjoint), `M == M'`.
"""
function is_hermitian(M :: AbstractMatrix{<:Number}; atol=ATOL :: Float64) :: Bool
    indsm, indsn = axes(M)
    if indsm != indsn
        return false
    end

    for i = indsn, j = i:last(indsn)
        if M[i,j] != M[j,i]' && !isapprox(M[i,j], M[j,i]', atol=atol)
            return false
        end
    end

    return true
end

"""
    is_positive_semidefinite(M :: AbstractMatrix) :: Bool

Returns `true` if the supplied matrix is square and positive semidefinite (all eigen values
are real and greater than or equal to 0).
"""
function is_positive_semidefinite(M :: AbstractMatrix{<:Number}; atol=ATOL) :: Bool
    if !isequal(size(M)...)
        return false
    end

    for λ in eigvals(M)
        if (real(λ) < 0) && !isapprox(real(λ), 0, atol=atol)
            return false
        elseif !isapprox(imag(λ), 0, atol=atol)
            return false
        end
    end

    return true
end

"""
    is_orthonormal_basis(basis_vecs ::  Vector;  atol=ATOL ::  Float64) :: Bool

Returns `true` if `basis_vecs` ``\\{|\\psi_j\\rangle\\}_{j=1}^n forms an orthonormal basis:

```math
\\langle \\psi_j |\\psi_j \\rangle = 1 \\quad \\text{and} \\quad \\langle \\psi_j | \\psi_{k} \\rangle = 0
```

for all ``j`` and ``k`` where ``j\\neq k``.
"""
function is_orthonormal_basis(basis_vecs :: Vector{<:AbstractVector}; atol=ATOL :: Float64) :: Bool
    num_vecs = length(basis_vecs)
    if num_vecs != length(basis_vecs[1])
        return false
    end

    for i = 1:num_vecs, j = i:num_vecs
        inner_prod = basis_vecs[j]'basis_vecs[i]
        if i == j &&  !isapprox(inner_prod,1,atol=atol)
            return false
        elseif j > i  && !isapprox(inner_prod,0,atol=atol)
            return false
        end
    end
    return true
end

"""
    n_product_id( tensor_coord :: Vector{Int64}, dims :: Vector{Int64} ) :: Int64

For a given tensor product of subspace dimension `dims` and tensor coordingate `tensor_coord`
returns the corresponding matrix index of the kronecker product.
The tenosr product `dims` is a vector of positive definite integers specifying
the dimension of each of the matrices in the tensor product.
The `tensor_coord` is a vector of positive definite integers specifying the target
index in each matrix in the tensor product.
The `dims` and `tensor_coord` should describe either row or column indices of the
tensor product.

A `DomainError` is thrown if any index in `tensor_coord` is not valid or the number of
elements in `tensor_coord` does not equal that of `dims`.
"""
function n_product_id(tensor_coord::Vector{Int64}, dims::Vector{Int64}) :: Int64
    num_ids = length(tensor_coord)
    if num_ids != length(dims)
        throw(DomainError(tensor_coord, "id does not contain the right number of elements"))
    elseif !all(x -> dims[x] ≥ tensor_coord[x] ≥ 1, 1:num_ids)
        throw(DomainError(dims, "not all num inputs are valid"))
    end

    id = 1
    multiplier = 1
    for i in num_ids:-1:1
        multiplier *= (i+1 > num_ids) ? 1 : dims[i+1]
        id += (tensor_coord[i] - 1) * multiplier
    end

    id
end

"""
    commutes(A :: Matrix, B :: Matrix; atol=ATOL :: Float64) :: Bool

Returns true if matrix `A` commutes with matrix `B`. Two matrices commute if
`A*B == B*A`. The matrices must have compatible dimensions:
* A: Matrix, mxn dimensions
* B: Matrix, nxm dimensions

A `DomainError` is thrown if the matrices are not compatible.
"""
function commutes(A :: AbstractMatrix{<:Number}, B :: Matrix; atol=ATOL :: Float64) :: Bool
    if (size(A)[1] != size(B)[2]) || (size(A)[2] != size(B)[1])
        throw(DomainError((A,B), "size(A) is not compatible with size(B)."))
    end

    all(isapprox.(A*B, B*A, atol=atol))
end

"""
    is_complete(Π :: Vector{<:AbstractMatrix}; atol=ATOL :: Float64) :: Bool

Returns `true` if the provided set of matrices sums to the identity.
"""
function is_complete(Π :: Vector{<:AbstractMatrix}; atol=ATOL :: Float64) :: Bool
    if !isapprox(sum(Π), I, atol=atol)
        return false
    end

    return true
end

ext: jl


text:
const SupervisedNetwork = Union{DeterministicNetwork,ProbabilisticNetwork}

# to suppress inclusion in models():
MLJBase.is_wrapper(::Type{DeterministicNetwork}) = true
MLJBase.is_wrapper(::Type{ProbabilisticNetwork}) = true

## FALL-BACKS FOR LEARNING NETWORKS EXPORTED AS MODELS

function MLJBase.update(model::Union{SupervisedNetwork,UnsupervisedNetwork},
                        verbosity, yhat, cache, args...)

    # If any `model` field has been replaced (and not just mutated)
    # then we actually need to fit rather than update (which will
    # force build of a new learning network). If `model` has been
    # created using a learning network export macro, the test used
    # below is perfect. In any other case it is at least conservative:
    network_model_ids = objectid.(models(yhat))
    fields = [getproperty(model, name) for
        name in fieldnames(typeof(model))]
    submodels = filter(f->f isa Model, fields)
    submodel_ids = objectid.(submodels)
    if !issubset(submodel_ids, network_model_ids)
        return fit(model, verbosity, args...)
    end

    is_anonymised = cache isa NamedTuple{(:sources, :data)}

    if is_anonymised
        sources, data = cache.sources, cache.data
        for k in eachindex(sources)
            rebind!(sources[k], data[k])
        end
    end

    fit!(yhat; verbosity=verbosity)
    if is_anonymised
        for s in sources
            rebind!(s, nothing)
        end
    end

    return yhat, cache, nothing
end

MLJBase.predict(composite::SupervisedNetwork, fitresult, Xnew) =
    fitresult(Xnew)

MLJBase.transform(composite::UnsupervisedNetwork, fitresult, Xnew) =
    fitresult(Xnew)

function fitted_params(yhat::Node)
    machs = machines(yhat)
    fitted = [fitted_params(m) for m in machs]
    return (machines=machs, fitted_params=fitted)
end

fitted_params(composite::Union{SupervisedNetwork,UnsupervisedNetwork}, yhat) =
    fitted_params(yhat)

## FOR EXPORTING LEARNING NETWORKS BY HAND

"""
    anonymize!(sources...)

Returns a named tuple `(sources=..., data=....)` whose values are the
provided source nodes and their contents respectively, and clears the
contents of those source nodes.

"""
function anonymize!(sources...)
    data = Tuple(s.data for s in sources)
    [MLJ.rebind!(s, nothing) for s in sources]
    return (sources=sources, data=data)
end

function report(yhat::Node)
    machs = machines(yhat)
    reports = [report(m) for m in machs]
    return (machines=machs, reports=reports)
end

# what should be returned by a fit method for an exported learning
# network:
function fitresults(yhat)
    inputs = sources(yhat, kind=:input)
    targets = sources(yhat, kind=:target)
    length(inputs) == 1 ||
        error("Improperly exported supervised network does "*
              "not have a unique input source. ")
    if length(targets) == 1
        cache = anonymize!(inputs[1], targets[1])
    elseif length(targets) == 0
        cache = anonymize!(inputs[1])
    else
        error("Improperly exported network has multiple target sources. ")
    end
    r = report(yhat)
    return yhat, cache, r
end

## EXPORTING LEARNING NETWORKS AS MODELS WITH @from_network

"""
    replace(W::MLJ.Node, a1=>b1, a2=>b2, ...)

Create a deep copy of a node `W`, and thereby replicate the learning
network terminating at `W`, but replacing any specified sources and
models `a1, a2, ...` of the original network with `b1, b2, ...`.
"""
function Base.replace(W::Node, pairs::Pair...)

    # Note: We construct nodes of the new network as values of a
    # dictionary keyed on the nodes of the old network. Additionally,
    # there are dictionaries of models keyed on old models and
    # machines keyed on old machines. The node and machine
    # dictionaries must be built simultaneously.

    # build model dict:
    model_pairs = filter(collect(pairs)) do pair
        first(pair) isa Model
    end
    models_ = models(W)
    models_to_copy = setdiff(models_, first.(model_pairs))
    model_copy_pairs = [model=>deepcopy(model) for model in models_to_copy]
    newmodel_given_old = IdDict(vcat(model_pairs, model_copy_pairs))

    # build complete source replacement pairs:
    source_pairs = filter(collect(pairs)) do pair
        first(pair) isa Source
    end
    sources_ = sources(W)
    sources_to_copy = setdiff(sources_, first.(source_pairs))
    isempty(sources_to_copy) ||
        @warn "No replacement specified for one or more source nodes. "*
    "Data there will be duplicated. "
    source_copy_pairs = [source=>deepcopy(source) for source in sources_to_copy]
    all_source_pairs = vcat(source_pairs, source_copy_pairs)

    # drop source nodes from all nodes of network terminating at W:
    nodes_ = filter(nodes(W)) do N
        !(N isa Source)
    end
    isempty(nodes_) && error("All nodes in network are source nodes. ")
    # instantiate node and machine dictionaries:
    newnode_given_old =
        IdDict{AbstractNode,AbstractNode}(all_source_pairs)
    newmach_given_old = IdDict{NodalMachine,NodalMachine}()

    # build the new network:
    for N in nodes_
       args = [newnode_given_old[arg] for arg in N.args]
         if N.machine === nothing
             newnode_given_old[N] = node(N.operation, args...)
         else
             if N.machine in keys(newmach_given_old)
                 mach = newmach_given_old[N.machine]
             else
                 train_args = [newnode_given_old[arg] for arg in N.machine.args]
                 mach = machine(newmodel_given_old[N.machine.model], train_args...)
                 newmach_given_old[N.machine] = mach
             end
             newnode_given_old[N] = N.operation(mach, args...)
        end
    end

    return newnode_given_old[nodes_[end]]

 end

# closure for later:
function fit_method(network, models...)

    network_Xs = sources(network, kind=:input)[1]

    function fit(model::M, verbosity, X, y) where M <: Supervised
        replacement_models = [getproperty(model, fld)
                              for fld in fieldnames(M)]
        model_replacements = [models[j] => replacement_models[j]
                          for j in eachindex(models)]
        network_ys = sources(network, kind=:target)[1]
        Xs = source(X)
        ys = source(y, kind=:target)
        source_replacements = [network_Xs => Xs, network_ys => ys]
        replacements = vcat(model_replacements, source_replacements)
        yhat = replace(network, replacements...)

        Set([Xs, ys]) == Set(sources(yhat)) ||
            error("Failed to replace sources in network blueprint. ")

        fit!(yhat, verbosity=verbosity)

        return fitresults(yhat)
    end

    function fit(model::M, verbosity, X) where M <:Unsupervised
        replacement_models = [getproperty(model, fld)
                              for fld in fieldnames(M)]
        model_replacements = [models[j] => replacement_models[j]
                          for j in eachindex(models)]
        Xs = source(X)
        source_replacements = [network_Xs => Xs,]
        replacements = vcat(model_replacements, source_replacements)
        Xout = replace(network, replacements...)
        Set([Xs]) == Set(sources(Xout)) ||
            error("Failed to replace sources in network blueprint. ")

        fit!(Xout, verbosity=verbosity)

        return fitresults(Xout)
    end

    return fit

end

net_alert(message) = throw(ArgumentError("Learning network export error.\n"*
                                     string(message)))
net_alert(k::Int) = throw(ArgumentError("Learning network export error $k. "))

# returns Model supertype - or `missing` if arguments are incompatible
function kind_(is_supervised, is_probabilistic)
    if is_supervised
        if ismissing(is_probabilistic) || !is_probabilistic
            return :DeterministicNetwork
        else
            return :ProbabilisticNetwork
        end
    else
        if ismissing(is_probabilistic) || !is_probabilistic
            return :UnsupervisedNetwork
        else
            return missing
        end
    end
end

function from_network_preprocess(modl, ex,
                                 is_probabilistic::Union{Missing,Bool})

    ex isa Expr || net_alert(1)
    ex.head == :call || net_alert(2)
    ex.args[1] == :(<=) || net_alert(3)
    ex.args[2] isa Expr || net_alert(4)
    ex.args[2].head == :call || net_alert(5)
    modeltype_ex = ex.args[2].args[1]
    modeltype_ex isa Symbol || net_alert(6)
    if length(ex.args[2].args) == 1
        kw_exs = []
    else
        kw_exs = ex.args[2].args[2:end]
    end
    fieldname_exs = []
    model_exs = []
    for ex in kw_exs
        ex isa Expr || net_alert(7)
        ex.head == :kw || net_alert(8)
        variable_ex = ex.args[1]
        value_ex = ex.args[2]
        variable_ex isa Symbol || net_alert(9)
        push!(fieldname_exs, variable_ex)
        value = modl.eval(value_ex)
        value isa Model ||
            net_alert("Got $value but expected something of type `Model`.")
        push!(model_exs, value_ex)
    end
    N_ex = ex.args[3]
    N = modl.eval(N_ex)
    N isa AbstractNode ||
        net_alert("Got $N but expected something of type `AbstractNode`. ")

    inputs = sources(N, kind=:input)
    targets = sources(N, kind=:target)

    length(inputs) == 0 &&
        net_alert("Network has no source with `kind=:input`.")
    length(inputs) > 1  &&
        net_alert("Network has multiple sources with `kind=:input`.")
    length(targets) > 1 &&
        net_alert("Network has multiple sources with `kind=:target`.")

    is_supervised = length(targets) == 1

    kind = kind_(is_supervised, is_probabilistic)
    ismissing(kind) &&
        net_alert("Network appears unsupervised (has no source with "*
                  "`kind=:target`) and so `is_probabilistic=true` "*
                  "declaration is not allowed. ")

    models_ = [modl.eval(e) for e in model_exs]
    issubset(models_, models(N)) ||
        net_alert("One or more specified models are not in the learning network "*
              "terminating at $N_ex.\n Use models($N_ex) to inspect models. ")

    nodes_  = nodes(N)

    return modeltype_ex, fieldname_exs, model_exs, N_ex, kind

end

from_network_preprocess(modl, ex) = from_network_preprocess(modl, ex, missing)

function from_network_preprocess(modl, ex, kw_ex)
    kw_ex isa Expr || net_alert(10)
    kw_ex.head == :(=) || net_alert(11)
    kw_ex.args[1] == :is_probabilistic ||
        net_alert("Unrecognized keywork `$(kw_ex.args[1])`.")
    value = kw_ex.args[2]
    if value isa Bool
        return from_network_preprocess(modl, ex, value)
    else
        net_alert("`is_probabilistic` can only be `true` or `false`.")
    end
end

function from_network_(modl, modeltype_ex, fieldname_exs, model_exs,
                          N_ex, kind)

    args = gensym(:args)

    # code defining the composite model struct and fit method:
    program1 = quote

        mutable struct $modeltype_ex <: MLJ.$kind
            $(fieldname_exs...)
        end

        MLJ.fit(model::$modeltype_ex, verbosity::Integer, $args...) =
            MLJ.fit_method(
                $N_ex, $(model_exs...))(model, verbosity, $args...)

    end

    program2 = quote
        defaults =
            MLJ.@set_defaults $modeltype_ex deepcopy.([$(model_exs...)])

    end

    modl.eval(program1)
    modl.eval(program2)

end

"""
    @from_network(NewCompositeModel(fld1=model1, fld2=model2, ...) <= N
    @from_network(NewCompositeModel(fld1=model1, fld2=model2, ...) <= N is_probabilistic=false

Create a new stand-alone model type called `NewCompositeModel`, using
a learning network as a blueprint. Here `N` refers to the terminal
node of the learning network (from which final predictions or
transformations are fetched).

*Important.* If the learning network is supervised (has a source with
`kind=:target`) and makes probabilistic predictions, then one must
declare `is_probabilistic=true`. In the deterministic case the keyword
argument can be omitted.

The model type `NewCompositeModel` is equipped with fields named
`:fld1`, `:fld2`, ..., which correspond to component models `model1`,
`model2`, ...,  appearing in the network (which must therefore be elements of
`models(N)`).  Deep copies of the specified component models are used
as default values in an automatically generated keyword constructor
for `NewCompositeModel`.

### Return value

 A new `NewCompositeModel` instance, with default field values.

For details and examples refer to the "Learning Networks" section of
the documentation.

"""
macro from_network(exs...)

    args = from_network_preprocess(__module__, exs...)
    modeltype_ex = args[1]

    from_network_(__module__, args...)

    esc(quote
        $modeltype_ex()
        end)

end

## A COMPOSITE FOR TESTING PURPOSES

"""
    SimpleDeterministicCompositeModel(;regressor=ConstantRegressor(),
                              transformer=FeatureSelector())

Construct a composite model consisting of a transformer
(`Unsupervised` model) followed by a `Deterministic` model. Mainly
intended for internal testing .

"""
mutable struct SimpleDeterministicCompositeModel{L<:Deterministic,
                             T<:Unsupervised} <: DeterministicNetwork
    model::L
    transformer::T

end

function SimpleDeterministicCompositeModel(; model=DeterministicConstantRegressor(),
                          transformer=FeatureSelector())

    composite =  SimpleDeterministicCompositeModel(model, transformer)

    message = MLJ.clean!(composite)
    isempty(message) || @warn message

    return composite

end

MLJBase.is_wrapper(::Type{<:SimpleDeterministicCompositeModel}) = true

function MLJBase.fit(composite::SimpleDeterministicCompositeModel,
                     verbosity::Integer, Xtrain, ytrain)
    X = source(Xtrain) # instantiates a source node
    y = source(ytrain, kind=:target)

    t = machine(composite.transformer, X)
    Xt = transform(t, X)

    l = machine(composite.model, Xt, y)
    yhat = predict(l, Xt)

    fit!(yhat, verbosity=verbosity)

    return fitresults(yhat)
end

# MLJBase.predict(composite::SimpleDeterministicCompositeModel, fitresult, Xnew) = fitresult(Xnew)

MLJBase.load_path(::Type{<:SimpleDeterministicCompositeModel}) = "MLJ.SimpleDeterministicCompositeModel"
MLJBase.package_name(::Type{<:SimpleDeterministicCompositeModel}) = "MLJ"
MLJBase.package_uuid(::Type{<:SimpleDeterministicCompositeModel}) = ""
MLJBase.package_url(::Type{<:SimpleDeterministicCompositeModel}) = "https://github.com/alan-turing-institute/MLJ.jl"
MLJBase.is_pure_julia(::Type{<:SimpleDeterministicCompositeModel}) = true
MLJBase.input_scitype(::Type{<:SimpleDeterministicCompositeModel{L,T}}) where {L,T} =
    MLJBase.input_scitype(T)
MLJBase.target_scitype(::Type{<:SimpleDeterministicCompositeModel{L,T}}) where {L,T} =
    MLJBase.target_scitype(L)


text:
using PkgCite
using PkgCite: collect_citations, bibliography, cited_packages, make_sentence
using InteractiveUtils
using Bibliography
using Test
import Pkg

include("cite_str.jl")

@testset "PkgCite.jl" begin
    @testset "Empty env" begin
        @test_logs (:warn, "No citations found in current environment") get_citations()
    end

    test_env = "test_env"
    Pkg.activate(test_env)
    Pkg.instantiate()
    Pkg.status()

    citations = collect_citations(false)

    @testset "Packages" begin
        pkgs = cited_packages(citations)
        @test "LabelledArrays" ∈ pkgs
        @test "QuadGK" ∈ pkgs
        @test "Symbolics" ∈ pkgs
        @test "RecursiveArrayTools" ∈ pkgs
        @test "ArrayInterface" ∈ pkgs
        # There is an error in the CITATION.bib for Distributions
        @test_broken haskey(citations, "Distributions")
    end

    @testset "Citations" begin
        bib = bibliography(citations)
        @test haskey(bib, "DifferentialEquations.jl-2017")
        @test haskey(bib, "quadgk")
        @test haskey(bib, "gowda2021high")
    end

    get_citations()

    @testset "Citations in .bib" begin
        bib_path = joinpath(@__DIR__, "julia_citations.bib")
        @test isfile(bib_path)
        bib = import_bibtex(bib_path)

        @test haskey(bib, "DifferentialEquations.jl-2017")
        @test haskey(bib, "quadgk")
        @test haskey(bib, "gowda2021high")
    end

    @testset "Only direct dependencies" begin
        citations = collect_citations(true)
        pkgs = cited_packages(citations)
        @test only(pkgs) == "Symbolics"
        @test make_sentence(citations) == CITE_STR_JL_SINGLE
        @test make_sentence(citations, texttt=true) == CITE_STR_JL_SINGLE_TT
    end

    @testset "PkgCite sentence" begin
        citations = collect_citations(false)
        str = make_sentence(citations)
        @test str == CITE_STR_JL

        str = make_sentence(citations, jl=false)
        @test str == CITE_STR

        str = make_sentence(citations, cite_commands=Dict{String,String}("Symbolics" => "\\autocite"))
        @test str == CITE_STR_JL_AUTO

        @testset "Clipboard" begin
            io = IOBuffer()
            get_tool_citation(io)
            str = String(take!(io))

            @test str == CITE_STR_JL * '\n'
            @test clipboard() == CITE_STR_JL

            get_tool_citation(io, jl=false)
            str = String(take!(io))

            @test str == CITE_STR * '\n'
            @test clipboard() == CITE_STR
        end
    end

    @testset "PkgCite using badge" begin
        zenodo_env = "badge_env"
        Pkg.activate(zenodo_env)
        Pkg.instantiate()
        Pkg.status()

        citations = collect_citations(true, badge=true)
        pkgs = cited_packages(citations)
        @test issubset(["WriteVTK", "Plots"], pkgs)
        @test make_sentence(citations) == CITE_BADGE
    end
end


text:
@testset "utils" begin
    tree, character_labels, taxa_labels = load_tree("data/output")

    @testset "remove_blanks" begin
        d = Dict("a" => "-", "b" => "2")
        d1 = remove_blanks(d)
        d2 = remove_blanks(d, change_to_N=true)
        @test d1["a"] == ""
        @test d2["a"] == "N"
    end

    @testset "get_max_depth" begin
        tree, character_labels, taxa_labels = generate_caos_rules("data/S10593.nex", "data/output")
        tree_out = load_tree("data/output")
        @test length(tree_out[1].children) == 2
        @test get_max_depth(tree_out[1], 10) == 33
        @test typeof(tree_out[1]) == CAOS.Node
    end

    @testset "find_sequence" begin
        result = find_sequence(tree, "Tubificoides_diazi_CE3411")
        @test typeof(result) == CAOS.Node
    end

    @testset "get_neighbors" begin
        neighbors = get_neighbors(tree, "Tubificoides_diazi_CE3411")

        @test length(neighbors) == 61
        @test neighbors[1:3] == ["Tubificoides_insularis_CE3417", "Tubificoides_insularis_I_CE3424", "Tubificoides_insularis_II_CE4438"]
        @test neighbors[58:61] ==  ["Tubificoides_heterochaetus_CE2266", "Tubificoides_heterochaetus_CE2447", "Tubificoides_heterochaetus_CE2260", "Tubificoides_insularis_CE3418"]
    end

    @testset "get_all_neighbors" begin
        @test get_all_neighbors(tree, character_labels, "Tubificoides_diazi_CE3411") == ["Tubificoides_diazi_CE3410"]
    end

    @testset "get_first_taxa_from_tree" begin
        #println(get_first_taxa_from_tree(tree))
        # needs to fix fefn
    end

    @testset "get_descendents" begin
        desc = get_descendents(tree)
        @test desc[1:3] == ["Tubificoides_insularis_CE3417", "Tubificoides_insularis_I_CE3424", "Tubificoides_insularis_II_CE4438"]
    end

    @testset "downsample_taxa" begin
        taxa = ["1", "2", "3", "4", "5"]
        @test length(downsample_taxa(taxa, 0.2)) == 1
    end

    @testset "get_adjusted_start" begin
        s = "--S1234"
        @test get_adjusted_start(1, s) == 3
    end
end


text:
#  Copyright 2022 California Institute of Technology
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
# Author: Philip G. Brodrick, philip.g.brodrick@jpl.nasa.gov
module SpectralUnmixing

using ArchGDAL
using EllipsisNotation
using DelimitedFiles
using Logging
using Statistics
using Distributed
using Printf
using LinearAlgebra
using Combinatorics
using Random

include("Datasets.jl")
include("EndmemberLibrary.jl")
include("Solvers.jl")
include("Plotting.jl")

# Endmember Library Functions
export SpectralLibrary, load_data!, filter_by_class!, read_envi_wavelengths, interpolate_library_to_new_wavelengths!, remove_wavelength_region_inplace!, scale_library!
export reduce_endmembers_nmf!, reduce_endmembers_kmeans!, reduce_endmembers_pca!, brightness_normalize!, split_library

# Plotting Functions
export plot_mean_endmembers, plot_endmembers, plot_endmembers_individually

# Dataset functions
export initiate_output_datasets, set_band_names, write_results

# Unmixing and simulation functions
export unmix_line, unmix_pixel, simulate_pixel

function wl_index(wavelengths::Vector{Float64}, target)
    argmin(abs.(wavelengths .- target))
end

function scale_data(refl::Matrix{Float64}, wavelengths::Vector{Float64}, criteria::String)

    if criteria == "none"
        return refl
    elseif criteria == "brightness"
        bad_regions_wl = [[1300,1500],[1800,2000]]
        good_bands = convert(Vector{Bool}, ones(length(wavelengths)))
        for br in bad_regions_wl
            good_bands[wl_index(wavelengths, br[1]):wl_index(wavelengths, br[2])] .= false
        end
        if length(size(refl)) == 2
            norm = sqrt.(mean(refl[:,good_bands].^2, dims=2))
        else
            norm = sqrt.(mean(refl[good_bands].^2))
        end
    else
        try
            target_wl = parse(Float64,criteria)
            if length(size(refl)) == 2
                norm = refl[:,wl_index(wavelengths, target_wl)] ./ 0.5
            else
                norm = refl[wl_index(wavelengths, target_wl)] ./ 0.5
            end
        catch e
            throw(ArgumentError(string("normalization must be [none, brightness, or a specific wavelength].  Provided:", criteria)))
        end
    end

    return refl ./ norm
end

function get_sma_permutation(class_idx, num_endmembers::Vector{Int64}, combination_type::String, library_length::Int64)
    if num_endmembers[1] != -1
        if combination_type == "class-even"

            perm_class_idx = []
            for class_subset in class_idx
                push!(perm_class_idx, Random.shuffle(class_subset))
            end

            perm = []
            selector = 1
            while selector <= num_endmembers[1]
                _p = mod(selector, length(perm_class_idx)) + 1
                push!(perm, perm_class_idx[_p][1])
                deleteat!(perm_class_idx[_p],1)

                if length(perm_class_idx[_p]) == 0
                    deleteat!(perm_class_idx,_p)
                end
                selector += 1
            end

        else
            perm = randperm(library_length)[1:num_endmembers[1]]
        end
    else
        perm = convert(Vector{Int64},1:library_length)
    end

    return perm
end

function results_from_mc(results::Matrix{Float64}, cost::Vector{Float64}, mode::String)

    if size(results)[1] == 1
        output_var = nothing
    else
        output_var = std(results, dims=1)[1,:]
    end

    if occursin("best", mode)
        best_idx = argmin(cost)
        output = results[best_idx,:]
    else
        output = mean(results, dims=1)[1,:]
    end

    return output, output_var

end

function unmix_pixel(library::SpectralLibrary, img_dat_input::Array{Float64}, unc_dat, class_idx, options, mode::String, n_mc::Int64, num_endmembers::Vector{Int64}, normalization::String, optimization::String, max_combinations::Int64, combination_type::String)

    if length(size(img_dat_input)) == 1
        img_dat = reshape(img_dat_input, 1, length(img_dat_input))
    else
        img_dat = img_dat_input
    end

    mc_comp_frac = zeros(n_mc, size(library.spectra)[1]+1)
    scores = zeros(n_mc)
    for mc in 1:n_mc #monte carlo loop
        Random.seed!(mc)

        d = img_dat
        if isnothing(unc_dat) == false
            d += (rand(size(d)) .* 2 .- 1) .* unc_dat
        end

        if occursin("sma", mode)
            perm = get_sma_permutation(class_idx, num_endmembers, combination_type, size(library.spectra)[1])
            G = library.spectra[perm, library.good_bands]

            G = scale_data(G, library.wavelengths[library.good_bands], normalization)'

            if occursin("pinv", optimization)
                inverse_method = "pinv"
            elseif occursin("qr", optimization)
                inverse_method = "qr"
            else
                inverse_method = "default"
            end
            x0 = dolsq(G, d', method=inverse_method)

            x0 = x0[:]
            res = nothing
            if occursin("bvls", optimization)
                res, cost = bvls(G, d[:], x0, zeros(size(x0)), ones(size(x0)), 1e-3, 100, 1, inverse_method)
            elseif occursin("ldsqp", optimization)
                res, cost = opt_solve(G, d[:], x0, zeros(length(x0)), ones(length(x0)) )
            elseif occursin("inverse", optimization)
                res = x0
                r = G * x0 - d[:]
                cost = dot(r,r)
            end
            mc_comp_frac[mc, perm] = res
            scores[mc] = cost

        elseif occursin("mesma", mode)
            solutions = []
            costs = zeros(size(options)[1]).+1e12

            if max_combinations != -1 && length(options) > max_combinations
                perm = randperm(length(options))[1:max_combinations]
            else
                perm = convert(Vector{Int64},1:length(options))
            end

            for (_comb, comb) in enumerate(options[perm])
                comb = [c for c in comb]
                #G = hcat(library.spectra[comb,:], ones(size(library.spectra[comb,:])[1],1))
                G = scale_data(library.spectra[comb, library.good_bands], library.wavelengths[library.good_bands], normalization)'

                x0 = dolsq(G, d')
                ls = nothing
                if optimization == "bvls"
                    ls, lc = bvls(G, d[:], x0, zeros(size(x0)), ones(size(x0)), 1e-3, 10, 1)
                    costs[_comb] = lc
                elseif optimization == "ldsqp"
                    ls, lc = opt_solve(G, d[:], x0, 0, 1)
                    costs[_comb] = lc
                elseif optimization == "inverse"
                    ls = x0
                    r = G * x0 - d[:]
                    costs[_comb] = dot(r,r)
                end
                push!(solutions,ls)

            end
            best = argmin(costs)
            scores[mc] = best

            mc_comp_frac[mc, [ind for ind in options[perm][best]]] = solutions[best]
        else
            error("Invalid mode provided")
        end
    end

    # Calculate the sum of values (inverse of shade), and then normalize
    mc_comp_frac[mc_comp_frac .< 0] .= 0
    mc_comp_frac[:,end] = sum(mc_comp_frac,dims=2)
    mc_comp_frac[:,1:end-1] = mc_comp_frac[:,1:end-1] ./ mc_comp_frac[:,end]

    # Aggregate results from per-library to per-unique-class
    mixture_results = zeros(size(mc_comp_frac)[1], length(library.class_valid_keys) + 1)
    for _i in 1:size(mc_comp_frac)[1]
        for (_class, cl) in enumerate(library.class_valid_keys)
            mixture_results[_i, _class] = sum(mc_comp_frac[_i,1:end-1][cl .== library.classes])
        end
        mixture_results[_i, end] = mc_comp_frac[_i,end]
    end

    output_mixture, output_mixture_var = results_from_mc(mixture_results, scores, mode)
    output_comp_frac, output_comp_frac_var = results_from_mc(mc_comp_frac, scores, mode)

    return output_mixture, output_mixture_var, output_comp_frac, output_comp_frac_var

end

function unmix_line(line::Int64, reflectance_file::String, mode::String, refl_nodata::Float64,
                    refl_scale::Float64, normalization::String, library::SpectralLibrary,
                    reflectance_uncertainty_file::String = "", n_mc::Int64 = 1,
                    combination_type::String = "all", num_endmembers::Vector{Int64} = [2,3],
                    max_combinations::Int64 = -1, optimization="bvls")

    Random.seed!(13)
    println(line)

    img_dat, unc_dat, good_data = load_line(reflectance_file, reflectance_uncertainty_file, line, library.good_bands, refl_nodata)
    mixture_results = fill(-9999.0, sum(good_data), size(library.class_valid_keys)[1] + 1)
    complete_fractions = zeros(size(img_dat)[1], size(library.spectra)[1] + 1)

    if n_mc > 1
        mixture_results_std = fill(-9999.0, sum(good_data), size(library.class_valid_keys)[1] + 1)
        complete_fractions_std = zeros(size(img_dat)[1], size(library.spectra)[1] + 1)
    else
        mixture_results_std = nothing
        complete_fractions_std = nothing
    end

    if isnothing(img_dat)
        return line, nothing, good_data, nothing, nothing
    end
    scale_data(img_dat, library.wavelengths[library.good_bands], normalization)
    img_dat = img_dat ./ refl_scale

    class_idx = []
    if combination_type == "class-even"
        for uc in library.class_valid_keys
            push!(class_idx, (1:size(library.classes)[1])[library.classes .== uc])
        end
    end

    # Prepare combinations if relevant
    options = []
    if mode == "mesma"
        if combination_type == "class-even"
            options = collect(Iterators.product(class_idx...))[:]
        elseif combination_type == "all"
            for num in num_endmembers
                combo = [c for c in combinations(1:length(library.classes), num)]
                push!(options,combo...)
            end
        else
            error("Invalid combiation string")
        end
    end

    # Solve for each pixel
    for _i in 1:size(img_dat)[1] # Pixel loop

        lid = img_dat[_i:_i,:]
        if isnothing(unc_dat)
            lud = nothing
        else
            lud = unc_dat[_i:_i,:]
        end

        loc_mixture_res, loc_mixture_var, loc_cf_res, loc_cf_var  = unmix_pixel(library, lid,
            lud, class_idx, options, mode, n_mc, num_endmembers, normalization, optimization,
            max_combinations, combination_type)

        complete_fractions[_i,:] = loc_cf_res
        mixture_results[_i,:] = loc_mixture_res

        if n_mc > 1
            complete_fractions_std[_i,:] = loc_cf_var
            mixture_results_std[_i,:] = loc_mixture_var
        end

    end

    return line, mixture_results, good_data, mixture_results_std, complete_fractions

end

function class_assign_fractions(complete_fractions, library::SpectralLibrary)
    # Aggregate results from per-library to per-unique-class
    if length(size(complete_fractions)) == 1
        cf = reshape(complete_fractions, (1, length(complete_fractions)))
    else
        cf = complete_fractions
    end
    mixture_results = zeros(size(cf)[1], length(library.class_valid_keys))

    for _i in 1:size(cf)[1]
        for (_class, cl) in enumerate(library.class_valid_keys)
            mixture_results[_i, _class] = sum(complete_fractions[_i,1:end-1][cl .== library.classes])
        end
    end

    if length(size(complete_fractions)) == 1
        return mixture_results[1,:]
    else
        return mixture_results
    end

end

function simulate_pixel(library::SpectralLibrary, max_components::Int64, combination_type::String, seed::Int64)

    Random.seed!(seed)

    output_mixture = zeros(size(library.spectra)[2])
    output_mixture[:] .= NaN

    class_idx = []
    if combination_type == "class-even"
        for uc in library.class_valid_keys
            push!(class_idx, (1:size(library.classes)[1])[library.classes .== uc])
        end
    end
    perm = get_sma_permutation(class_idx, [max_components], combination_type, size(library.spectra)[1])

    G = library.spectra[perm,:]

    distribution = rand(max_components)
    distribution = distribution ./ sum(distribution)

    output_distribution = zeros(size(library.spectra)[1])
    output_distribution[perm] = distribution

    output_distribution_classes = zeros(size(library.class_valid_keys))
    for (_class, cl) in enumerate(library.class_valid_keys)
        output_distribution_classes[_class] = sum(output_distribution[cl .== library.classes])
    end

    simulated_rfl = G' * distribution
    return simulated_rfl, output_distribution, output_distribution_classes

end

end

max_stars_repo_name: fotxoa-geo/SpectralUnmixing

max_issues_repo_name: fotxoa-geo/SpectralUnmixing

max_forks_repo_name: fotxoa-geo/SpectralUnmixing


text:
include("../examples/lorenz.jl")
using LinearAlgebra
using JLD
using SharedArrays
using Distributed
function sens(s,nSteps)
    d = 3
    n_runup = 2000
    u = rand(d)
    for i = 1:n_runup
	u = next(u, s)
    end
    u_trj = step(u, s, nSteps)
    x, y, z = view(u_trj,1,:), view(u_trj,2,:), view(u_trj,3,:)

    du_trj = dstep(u_trj, s)
    ddu_trj = d2step(u_trj,s)
    vs = zeros(d)
    q = rand(d)
    q /= norm(q)

    pp = pert(u_trj, s)
    dpp = dpert(u_trj, s)
    q1 = rand(d)
    vs1 = zeros(d)

    dJds_st = 0.
    dJds_ust = 0.
    N = 12
    nSteps = nSteps + 1
    g = zeros(nSteps)
    J = copy(z)

    # Df represents the derivative of f
    # on the
    # unstable manifold.
    # Xu = a q
    Xu = zeros(d, nSteps)
    a = zeros(nSteps)
    Da = zeros(nSteps)
    Dq = zeros(d)
    Dvs = zeros(d)
    Dvs1 = zeros(d)
    rho = s[2]
    # nSteps large number.
    for i = 1:nSteps-1
	dui = du_trj[:,:,i] # for large systems, can't store Jacobian.
    	ppi = pp[:,i]
	dppi = dpp[:,:,i]
	xi, yi, zi = x[i], y[i], z[i]

    	q1 .= dui*q
    	alpha = norm(q1)
    	q1 ./= alpha
    	alpha2 = alpha*alpha

    	vs1 .= dui*vs + ppi
    	a[i+1] = dot(vs1, q1)
    	vs1 .-= a[i+1]*q1

    	d2q = reshape([dot(ddu_trj[:,j,i],
    				q) for j=1:d*d],d,d)
    	Dq = d2q*q/alpha2 + dui*Dq/alpha2
    	dalphadx = dot(alpha2*Dq, q1)
    	Dq .= Dq .- dot(Dq,q1)*q1

    	Dvs1 .= d2q*vs/alpha + dui*Dvs/alpha + dppi*q1/alpha
    	Da[i+1] = dot(vs1, Dq) + dot(Dvs1, q1)
    	Dvs1 .= Dvs1 - Da[i+1]*q1 - a[i+1]*Dq
    	Delta_Da = dot(Dvs1, q1) + dot(vs1, Dq)
    	Dvs1 .-= Delta_Da*q
    	Da[i+1] += Delta_Da

    	g[i+1] = g[i]/alpha - dalphadx/alpha2
    	dJdu = [0., 0., 1.]
    	dJds_st += dot(dJdu, vs)/nSteps
    	vs .= vs1
    	q .= q1
    	Dvs .= Dvs1
    end

    for n = 1:N
    	J_shift = J[n+1:end]
    	nJ = length(J_shift)
    	g_shift = g[2:nJ+1]
    	Da_shift = Da[2:nJ+1]
    	a_shift = a[2:nJ+1]
    	dJds_ust -= dot(J_shift,Da_shift)/nJ
    	dJds_ust -= dot(J_shift,g_shift.*a_shift)/nJ
    end

    return dJds_st + dJds_ust
end
function get_sens(s)
    nSteps = 500000
    n_exps = size(s)[2]
    dJds = zeros(n_exps)
    var_dJds = zeros(n_exps)
    n_rep = 16
    dJds_proc = SharedArray{Float64}(n_rep)
    for k=1:n_exps
    	sk = s[:,k]
    	dJds_proc .= 0.
    	t = @distributed for j=1:n_rep
    		dJds_proc[j] = sens(sk,nSteps)/n_rep
    	end
    	wait(t)
    	dJds[k] = sum(dJds_proc)
        var_dJds[k] = n_rep*sum(dJds_proc.*dJds_proc) - (dJds[k])^2
    end
    save("../data/sens/lorenz/dJds.jld", "rho",
         s[2,:], "dJds", dJds, "var_dJds", var_dJds)
end

max_stars_repo_name: nishaChandramoorthy/linearResponse

max_issues_repo_name: nishaChandramoorthy/linearResponse

max_forks_repo_name: nishaChandramoorthy/linearResponse


text:
#=

This file provides helper functions to create a MieModel object

=#

"""
    $(FUNCTIONNAME)(computation_type::NAI2, aerosol::AbstractAerosolType, λ::Real, polarization::AbstractPolarizationType, truncation_type::AbstractTruncationType)

Convenience function to create Mie Model with NAI2 computation type
"""
function make_mie_model(computation_type::NAI2,
                        aerosol::Aerosol,
                        λ::Real,
                        polarization::AbstractPolarizationType,
                        truncation_type::AbstractTruncationType,
                        r_max::Real,
                        nquad_radius::Integer)
    return MieModel(computation_type, aerosol, λ, polarization, truncation_type, r_max, nquad_radius, zeros(1, 1, 1), zeros(1, 1, 1))
end

"""
    $(FUNCTIONNAME)(computation_type::PCW, aerosol::AbstractAerosolType, λ::Real, polarization::AbstractPolarizationType, truncation_type::AbstractTruncationType, wigner_filepath::String)

Convenience function to load Wigner matrices from file and create Mie Model with PCW computation type
"""
function make_mie_model(computation_type::PCW,
                        aerosol::Aerosol,
                        λ::Real,
                        polarization::AbstractPolarizationType,
                        truncation_type::AbstractTruncationType,
                        r_max::Real,
                        nquad_radius::Integer,
                        wigner_filepath::String)
    wigner_A, wigner_B = Scattering.load_wigner_values(wigner_filepath)
    return MieModel(computation_type, aerosol, λ, polarization, truncation_type, r_max, nquad_radius, wigner_A, wigner_B)
end

"""
    $(FUNCTIONNAME)(computation_type::PCW, aerosol::AbstractAerosolType, λ::Real, polarization::AbstractPolarizationType, truncation_type::AbstractTruncationType, wigner_A, wigner_B)

Convenience function to take in Wigner matrices and create Mie Model with PCW computation type
"""
function make_mie_model(computation_type::PCW,
                        aerosol::AbstractAerosolType,
                        λ::Real,
                        polarization::AbstractPolarizationType,
                        truncation_type::AbstractTruncationType,
                        r_max::Real,
                        nquad_radius::Integer,
                        wigner_A, wigner_B)
return MieModel(computation_type, aerosol, λ, polarization, truncation_type, r_max, nquad_radius, wigner_A, wigner_B)
end


text:
# This uses the `compute_forces(term, ψ, occ; kwargs...)` function defined by all terms
"""
Compute the forces of an obtained SCF solution. Returns the forces wrt. the fractional
lattice vectors. To get cartesian forces use [`compute_forces_cart`](@ref).
Returns a list of lists of forces (as SVector{3}) in the same order as the `atoms`
and `positions` in the underlying [`Model`](@ref).
"""
@timing function compute_forces(basis::PlaneWaveBasis{T}, ψ, occupation; kwargs...) where {T}
    # no explicit symmetrization is performed here, it is the
    # responsability of each term to return symmetric forces
    forces_per_term = [compute_forces(term, basis, ψ, occupation; kwargs...)
                       for term in basis.terms]
    sum(filter(!isnothing, forces_per_term))
end

"""
Compute the cartesian forces of an obtained SCF solution in Hartree / Bohr.
Returns a list of lists of forces
`[[force for atom in positions] for (element, positions) in atoms]`
which has the same structure as the `atoms` object passed to the underlying [`Model`](@ref).
"""
function compute_forces_cart(basis::PlaneWaveBasis, ψ, occupation; kwargs...)
    forces_reduced = compute_forces(basis, ψ, occupation; kwargs...)
    covector_red_to_cart.(basis.model, forces_reduced)
end

function compute_forces(scfres)
    compute_forces(scfres.basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ)
end
function compute_forces_cart(scfres)
    compute_forces_cart(scfres.basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ)
end


text:
module StaticCompiler

using GPUCompiler: GPUCompiler
using LLVM: LLVM
using Libdl: Libdl
using Base: RefValue
using Serialization: serialize, deserialize
using Clang_jll: clang

export compile, load_function
export native_code_llvm, native_code_typed, native_llvm_module

"""
    compile(f, types, path::String = tempname()) --> (compiled_f, path)

   !!! Warning: this will fail on programs that heap allocate any memory, or have dynamic dispatch !!!

Statically compile the method of a function `f` specialized to arguments of the type given by `types`.

This will create a directory at the specified path (or in a temporary directory if you exclude that argument)
that contains the files needed for your static compiled function. `compile` will return a
`StaticCompiledFunction` object and `obj_path` which is the absolute path of the directory containing the
compilation artifacts. The `StaticCompiledFunction` can be treated as if it is a function with a single
method corresponding to the types you specified when it was compiled.

To deserialize and instantiate a previously compiled function, simply execute `load_function(path)`, which
returns a callable `StaticCompiledFunction`.

### Example:

Define and compile a `fib` function:
```julia
julia> using StaticCompiler

julia> fib(n) = n <= 1 ? n : fib(n - 1) + fib(n - 2)
fib (generic function with 1 method)

julia> fib_compiled, path = compile(fib, Tuple{Int}, "fib")
(f = fib(::Int64) :: Int64, path = "fib")

julia> fib_compiled(10)
55
```
Now we can quit this session and load a new one where `fib` is not defined:
```julia
julia> fib
ERROR: UndefVarError: fib not defined

julia> using StaticCompiler

julia> fib_compiled = load_function("fib.cjl")
fib(::Int64) :: Int64

julia> fib_compiled(10)
55
```
Tada!

### Details:

Here is the structure of the directory created by `compile` in the above example:
```julia
shell> tree fib
path
├── obj.cjl
├── obj.o
└── obj.so

0 directories, 3 files
````
* `obj.so` (or `.dylib` on MacOS) is a shared object file that can be linked to in order to execute your
compiled julia function.
* `obj.cjl` is a serialized `LazyStaticCompiledFunction` object which will be deserialized and instantiated
with `load_function(path)`. `LazyStaticcompiledfunction`s contain the requisite information needed to link to the
`obj.so` inside a julia session. Once it is instantiated in a julia session (i.e. by
`instantiate(::LazyStaticCompiledFunction)`, this happens automatically in `load_function`), it will be of type
`StaticCompiledFunction` and may be called with arguments of type `types` as if it were a function with a
single method (the method determined by `types`).
"""
function compile(f, _tt, path::String = tempname();  name = GPUCompiler.safe_name(repr(f)), kwargs...)
    tt = Base.to_tuple_type(_tt)
    isconcretetype(tt) || error("input type signature $_tt is not concrete")

    rt = only(native_code_typed(f, tt))[2]
    isconcretetype(rt) || error("$f on $_tt did not infer to a concrete type. Got $rt")

    # Would be nice to use a compiler pass or something to check if there are any heap allocations or references to globals
    # Keep an eye on https://github.com/JuliaLang/julia/pull/43747 for this

    f_wrap!(out::Ref, args::Ref{<:Tuple}) = (out[] = f(args[]...); nothing)

    generate_shlib(f_wrap!, Tuple{RefValue{rt}, RefValue{tt}}, path, name; kwargs...)

    lf = LazyStaticCompiledFunction{rt, tt}(Symbol(f), path, name)
    cjl_path = joinpath(path, "obj.cjl")
    serialize(cjl_path, lf)
    (; f = instantiate(lf), path=abspath(path))
end

"""
    load_function(path) --> compiled_f

load a `StaticCompiledFunction` from a given path. This object is callable.
"""
load_function(path) = instantiate(deserialize(joinpath(path, "obj.cjl")) :: LazyStaticCompiledFunction)

struct LazyStaticCompiledFunction{rt, tt}
    f::Symbol
    path::String
    name::String
end

function instantiate(p::LazyStaticCompiledFunction{rt, tt}) where {rt, tt}
    StaticCompiledFunction{rt, tt}(p.f, generate_shlib_fptr(p.path::String, p.name))
end

struct StaticCompiledFunction{rt, tt}
    f::Symbol
    ptr::Ptr{Nothing}
end

function Base.show(io::IO, f::StaticCompiledFunction{rt, tt}) where {rt, tt}
    types = [tt.parameters...]
    print(io, String(f.f), "(", join(("::$T" for T ∈ tt.parameters), ',')  ,") :: $rt")
end

function (f::StaticCompiledFunction{rt, tt})(args...) where {rt, tt}
    Tuple{typeof.(args)...} == tt || error("Input types don't match compiled target $((tt.parameters...,)). Got arguments of type $(typeof.(args))")
    out = RefValue{rt}()
    refargs = Ref(args)
    ccall(f.ptr, Nothing, (Ref{rt}, Ref{tt}), out, refargs)
    out[]
end

instantiate(f::StaticCompiledFunction) = f

module TestRuntime
    # dummy methods
    signal_exception() = return
    # HACK: if malloc returns 0 or traps, all calling functions (like jl_box_*)
    #       get reduced to a trap, which really messes with our test suite.
    malloc(sz) = Ptr{Cvoid}(Int(0xDEADBEEF))
    report_oom(sz) = return
    report_exception(ex) = return
    report_exception_name(ex) = return
    report_exception_frame(idx, func, file, line) = return
end

struct TestCompilerParams <: GPUCompiler.AbstractCompilerParams end
GPUCompiler.runtime_module(::GPUCompiler.CompilerJob{<:Any,TestCompilerParams}) = TestRuntime

function native_job(@nospecialize(func), @nospecialize(types); kernel::Bool=false, name=GPUCompiler.safe_name(repr(func)), kwargs...)
    source = GPUCompiler.FunctionSpec(func, Base.to_tuple_type(types), kernel, name)
    target = GPUCompiler.NativeCompilerTarget(always_inline=true)
    params = TestCompilerParams()
    GPUCompiler.CompilerJob(target, source, params), kwargs
end

"""
```julia
generate_shlib(f, tt, path::String, name::String; kwargs...)
```
Low level interface for compiling a shared object / dynamically loaded library
 (`.so` / `.dylib`) for function `f` given a tuple type `tt` characterizing
the types of the arguments for which the function will be compiled.

See also `StaticCompiler.generate_shlib_fptr`.

### Examples
```julia
julia> function test(n)
           r = 0.0
           for i=1:n
               r += log(sqrt(i))
           end
           return r/n
       end
test (generic function with 1 method)

julia> path, name = StaticCompiler.generate_shlib(test, Tuple{Int64}, "./test")
("./test", "test")

shell> tree \$path
./test
|-- obj.o
`-- obj.so

0 directories, 2 files

julia> test(100_000)
5.256496109495593

julia> ccall(StaticCompiler.generate_shlib_fptr(path, name), Float64, (Int64,), 100_000)
5.256496109495593
```
"""
function generate_shlib(f, tt, path::String = tempname(), name = GPUCompiler.safe_name(repr(f)); kwargs...)
    mkpath(path)
    obj_path = joinpath(path, "obj.o")
    lib_path = joinpath(path, "obj.$(Libdl.dlext)")
    open(obj_path, "w") do io
        job, kwargs = native_job(f, tt; name, kwargs...)
        obj, _ = GPUCompiler.codegen(:obj, job; strip=true, only_entry=false, validate=false)

        write(io, obj)
        flush(io)

        # Pick a Clang
        cc = Sys.isapple() ? `cc` : clang()
        # Compile!
        run(`$cc -shared -o $lib_path $obj_path`)
    end
    path, name
end

function generate_shlib_fptr(f, tt, path::String=tempname(), name = GPUCompiler.safe_name(repr(f)); temp::Bool=true, kwargs...)
    generate_shlib(f, tt, path, name; kwargs...)
    lib_path = joinpath(abspath(path), "obj.$(Libdl.dlext)")
    ptr = Libdl.dlopen(lib_path, Libdl.RTLD_LOCAL)
    fptr = Libdl.dlsym(ptr, "julia_$name")
    @assert fptr != C_NULL
    if temp
        atexit(()->rm(path; recursive=true))
    end
    fptr
end

"""
```julia
generate_shlib_fptr(path::String, name)
```
Low level interface for obtaining a function pointer by `dlopen`ing a shared
library given the `path` and `name` of a `.so`/`.dylib` already compiled by
`generate_shlib`.

See also `StaticCompiler.enerate_shlib`.

### Examples
```julia
julia> function test(n)
           r = 0.0
           for i=1:n
               r += log(sqrt(i))
           end
           return r/n
       end
test (generic function with 1 method)

julia> path, name = StaticCompiler.generate_shlib(test, Tuple{Int64}, "./test");

julia> test_ptr = StaticCompiler.generate_shlib_fptr(path, name)
Ptr{Nothing} @0x000000015209f600

julia> ccall(test_ptr, Float64, (Int64,), 100_000)
5.256496109495593

julia> @ccall \$test_ptr(100_000::Int64)::Float64 # Equivalently
5.256496109495593

julia> test(100_000)
5.256496109495593
```
"""
function generate_shlib_fptr(path::String, name)
    lib_path = joinpath(abspath(path), "obj.$(Libdl.dlext)")
    ptr = Libdl.dlopen(lib_path, Libdl.RTLD_LOCAL)
    fptr = Libdl.dlsym(ptr, "julia_$name")
    @assert fptr != C_NULL
    fptr
end

function native_code_llvm(@nospecialize(func), @nospecialize(types); kwargs...)
    job, kwargs = native_job(func, types; kwargs...)
    GPUCompiler.code_llvm(stdout, job; kwargs...)
end

function native_code_typed(@nospecialize(func), @nospecialize(types); kwargs...)
    job, kwargs = native_job(func, types; kwargs...)
    GPUCompiler.code_typed(job; kwargs...)
end

# Return an LLVM module
function native_llvm_module(f, tt, name = GPUCompiler.safe_name(repr(f)); kwargs...)
    job, kwargs = native_job(f, tt; name, kwargs...)
    m, _ = GPUCompiler.codegen(:llvm, job; strip=true, only_entry=false, validate=false)
    return m
end

end # module


text:
"""
   probability_check(pomdp::POMDP)
Checks if the transition and observation function of the discrete `pomdp`
have probability mass that sums up to unity for all state-action pairs.
"""
function probability_check(pomdp::POMDP)
    obs_prob_consistency_check(pomdp)
    trans_prob_consistency_check(pomdp)
end

"""
    obs_prob_consistency_check(pomdp::POMDP)
Checks if the observation function of the discrete `pomdp`
has probability mass that sums up to unity for all state-action pairs.
"""
function obs_prob_consistency_check(pomdp::POMDP)
    # initalize space
    sspace = states(pomdp)
    aspace = actions(pomdp)
    ospace = observations(pomdp)
    # iterate through all s-a pairs
    for s in sspace
        for a in aspace
            obs = observation(pomdp, a, s)
            psum = 0.0
            for o in ospace
                p = pdf(obs, o)
                @assert p ≥ 0 "Probability is negative for state: $s, action: $a, observation: $o"
                psum += p
            end
            @assert isapprox(psum, 1.0) "Observation probability does not sum to unity for state: $s, action: $a"
        end
    end
end

"""
    trans_prob_consistency_check(pomdp::Union{MDP, POMDP})
Checks if the transition function of the discrete problem
has probability mass that sums up to unity for all state-action pairs.
"""
function trans_prob_consistency_check(pomdp::Union{MDP, POMDP})
    # initalize space
    sspace = states(pomdp)
    aspace = actions(pomdp)
    # iterate through all s-a pairs
    for s in sspace
        for a in aspace
            tran = transition(pomdp, s, a)
            psum = 0.0
            for sp in sspace
                p = pdf(tran, sp)
                @assert p ≥ 0 "Probability is negative for state: $s, action: $a, next state: $sp"
                psum += p
            end
            @assert isapprox(psum, 1.0) "Transition probability does not sum to unity for state: $s, action: $a"
        end
    end
end


text:
using DataFrames, DataFramesMeta, Statistics
import FileIO:load

# Helper functions to load dataframe files

function load_dataframe(filepath, filename)
    file = load(joinpath(filepath, filename))
    df = file["df"]
end

function df_loadcat(;filepath = "/path/to/jldfiles_GAUSSIAN/")
    jldfiles = searchdirendings(filepath, ".jld2")
    dfmat = []
    for i in 1:length(jldfiles)
        dflearner = load_dataframe(filepath, jldfiles[i])
        if in(:nposteriors, names(dflearner))
            deletecols!(dflearner, :nposteriors) # Delete because it takes too much space
        end
        if in(:seedparticlefilter, names(dflearner))
            rename!(dflearner, :seedparticlefilter => :seedlearner) # Rename for compatibility with older naming
        end
        if in(:nresamplings, names(dflearner))
            deletecols!(dflearner, :nresamplings) # Delete because it takes too much space
        end
        push!(dfmat, dflearner)
    end
    df = vcat(dfmat...)
end

function df_loadseparately(;filepath = "/path/to/jldfiles_GAUSSIAN/")
    jldfiles = searchdirendings(filepath, ".jld2")
    #jldfiles = ["multi_leaky.jld2", "multi_pf10.jld2"]
    dfall = []
    for i in 1:length(jldfiles)
        df = load_dataframe(filepath, jldfiles[i])
        push!(dfall, df)
    end
    dfall
end


text:
using Test
import Pluto
import Pluto: update_save_run!, WorkspaceManager, ClientSession, ServerSession, Notebook, Cell

import UUIDs: UUID, uuid1

function get_unique_short_id()
    string(uuid1())[1:8]
end

function stringify_keys(d::Dict)
    Dict(string(k) => stringify_keys(v) for (k, v) in d)
end
stringify_keys(x::Any) = x

import Pluto.Firebasey

function await_with_timeout(check::Function, timeout::Real=60.0, interval::Real=.05)
    starttime = time()
    while !check()
        sleep(interval)
        if time() - starttime >= timeout
            error("Timeout after $(timeout) seconds")
        end
    end
end

@testset "Communication protocol" begin

    @testset "Functionality sweep" begin
        u = [uuid1() for _ in 1:100]

        buffer = IOBuffer()
        client = ClientSession(:buffery, buffer)

        🍭 = ServerSession()
        # 🍭.connected_clients[client.id] = client

        notebook = Notebook([
            Cell(
                u[1],
                ""
            ),
        ])
        🍭.notebooks[notebook.notebook_id] = notebook
        update_save_run!(🍭, notebook, notebook.cells)
        client.connected_notebook = notebook
        read(buffer)

        local_state = Pluto.notebook_to_js(notebook)

        function send(type, body, metadata = Dict(:notebook_id => string(notebook.notebook_id)))
            request_id = get_unique_short_id()
            Pluto.process_ws_message(🍭, Dict(
                "type" => string(type),
                "client_id" => string(client.id),
                "request_id" => request_id,
                "body" => body,
                metadata...
            ) |> stringify_keys, client.stream)
        end

        function send_new_state(new_state)
            patches::Array{Dict} = Firebasey.diff(new_state, local_state)
            # @info "patches" patches
            send(:update_notebook, Dict("updates" => patches))

            new_state
        end

        # function update_local_notebook(mutate_fn::Function)
        #     mutable_notebook = deepcopy(notebook)

        #     mutate_fn(mutable_notebook)

        #     new_state = Pluto.notebook_to_js(mutable_notebook)
        #     send_new_state(new_state)
        # end

        last_position = position(buffer)
        function wait_for_updates(process=true)
            # @info "Waiting for updates"
            await_with_timeout() do
                position(buffer) != last_position
            end
            seek(buffer, last_position)
            response = Pluto.unpack(buffer)
            last_position = position(buffer)

            if process
                if response["type"] == "notebook_diff"
                    message = response["message"]
                    patches = [Base.convert(Firebasey.JSONPatch, update) for update in message["patches"]]
                    # @show patches
                    for patch in patches
                        Firebasey.applypatch!(local_state, patch)
                    end
                end
            end

            # @info "Response received" response

            local_state
        end

        # function patch_and_check(mutate_fn::Function)
        #     desired = update_local_notebook(mutate_fn)
        #     result = wait_for_updates()
        #     desired == result
        # end

        @test_nowarn send(:connect, Dict())
        wait_for_updates()

        @test_nowarn send_new_state(local_state)
        wait_for_updates(false)

        read(buffer)

        #=
        We would also like to test:
        - add cell
        - set code and run
        - fold cell
        - move cell
        - delete cell

        - run multiple cells
        - move cells
        - set bond

        - move notebook file
        - search for docs
        - show more items of an array =#
        @testset "Extension of response" begin
            Pluto.responses[:custom_response] = function (🙋::Pluto.ClientRequest)
                return "Yes this can be extended!"
            end
           @test send(:custom_response, Dict()) == "Yes this can be extended!"
        end

        @test_nowarn send(:shutdown_notebook, Dict("keep_in_session" => false))

        @test_nowarn await_with_timeout() do
            !haskey(🍭.notebooks, notebook.notebook_id)
        end
    end

    @testset "Docs" begin
        @test occursin("square root", Pluto.PlutoRunner.doc_fetcher("sqrt", Main)[1])
        @test occursin("square root", Pluto.PlutoRunner.doc_fetcher("Base.sqrt", Main)[1])
        @test occursin("No documentation found", Pluto.PlutoRunner.doc_fetcher("Base.findmeta", Main)[1])

        # Issue #1128
        # Ref https://docs.julialang.org/en/v1/manual/documentation/#Dynamic-documentation
        m = Module()
        Core.eval(m, :(
            module DocTest
            "Normal docstring"
            struct MyType
                value::String
            end
            Docs.getdoc(t::MyType) = "Documentation for MyType with value $(t.value)"
            const x = MyType("x")
            "A global variable"
            global y
            end
        ))

        @test occursin("Normal docstring", Pluto.PlutoRunner.doc_fetcher("MyType", m.DocTest)[1])
        @test occursin("Normal docstring", Pluto.PlutoRunner.doc_fetcher("DocTest.MyType", m)[1])
        @test occursin("Documentation for MyType with value", Pluto.PlutoRunner.doc_fetcher("x", m.DocTest)[1])
        @test occursin("Documentation for MyType with value", Pluto.PlutoRunner.doc_fetcher("DocTest.x", m)[1])
        @test occursin("A global variable", Pluto.PlutoRunner.doc_fetcher("y", m.DocTest)[1])
        @test occursin("A global variable", Pluto.PlutoRunner.doc_fetcher("DocTest.y", m)[1])
    end

    @testset "PlutoRunner API" begin
        fakeclient = ClientSession(:fake, nothing)
        🍭 = ServerSession()
        🍭.options.evaluation.workspace_use_distributed = true
        🍭.connected_clients[fakeclient.id] = fakeclient

        notebook = Notebook([
            Cell("PlutoRunner.notebook_id[] |> Text"),
            Cell("""
            let
                # not actually public API but we test it anyways
                a = PlutoRunner._publish(Dict(
                    "hello" => "world",
                    "xx" => UInt8[6,7,8],
                ))
                b = PlutoRunner._publish("cool")
                Text((a, b))
            end
            """),
            Cell("3"),
            Cell("PlutoRunner.publish_to_js(Ref(4))"),
            Cell("PlutoRunner.publish_to_js((ref=4,))"),
            Cell("x = Dict(:a => 6)"),
            Cell("PlutoRunner.publish_to_js(x)"),
        ])
        fakeclient.connected_notebook = notebook

        update_save_run!(🍭, notebook, notebook.cells)
        @test notebook.cells[1].output.body == notebook.notebook_id |> string

        @test !notebook.cells[2].errored
        a, b = Meta.parse(notebook.cells[2].output.body) |> eval
        p = notebook.cells[2].published_objects
        @test sort(collect(keys(p))) == sort([a,b])
        @test isempty(notebook.cells[3].published_objects)

        @test p[a] == Dict(
            "hello" => "world",
            "xx" => UInt8[6,7,8],
        )
        @test p[b] == "cool"

        old_pa = p[a]
        old_pb = p[b]
        update_save_run!(🍭, notebook, notebook.cells)
        p = notebook.cells[2].published_objects
        a, b = Meta.parse(notebook.cells[2].output.body) |> eval
        @test p[a] == old_pa
        @test p[b] == old_pb

        @test !isempty(notebook.cells[2].published_objects)

        setcode(notebook.cells[2], "2")
        update_save_run!(🍭, notebook, notebook.cells)
        @test isempty(notebook.cells[2].published_objects)

        @test notebook.cells[4].errored
        @test !notebook.cells[5].errored
        @test !isempty(notebook.cells[5].published_objects)

        p = notebook.cells[7].published_objects
        @test length(p) == 1
        old_x = values(p) |> first
        @test old_x == Dict(:a => 6)

        update_save_run!(🍭, notebook, notebook.cells[7])
        p = notebook.cells[7].published_objects
        new_x = values(p) |> first
        @test new_x == old_x
        @test new_x === old_x # did not change, because we don't resync the same object

        update_save_run!(🍭, notebook, notebook.cells[6])
        p = notebook.cells[7].published_objects
        new_x = values(p) |> first
        @test new_x == old_x
        @test new_x !== old_x # changed, because a new (mutable) Dict was created

        @test isempty(notebook.cells[2].published_objects)
        @test !isempty(notebook.cells[5].published_objects)

        WorkspaceManager.unmake_workspace((🍭, notebook))
    end
end


text:
import Base.sum
export sum

function sum(x::AffineExpr)
  return ones(1, x.size[1]) * x * ones(x.size[2], 1)
end

function sum(x::AffineExpr, dim::Int64)
  if dim == 1
    return ones(1, x.size[1]) * x
  elseif dim == 2
    return x * ones(x.size[2], 1)
  else
    error("dimension must be 1 or 2, but got $dimension")
  end
end


text:
using MLStyle

"""
    @support_matching XXX is_xxx

# Example
```julia
struct MyFunction
    ex :: Expr
    test_field :: Any
end

is_xxx(ex::Expr) = Meta.isexpr(ex, :function)
is_xxx(_) = false

julia> MyFunction(ex::Expr) = MyFunction(ex, :aaa)
julia> @support_matching(MyFunction, is_xxx)

julia> @match :(function f() end) begin
        MyFunction(;test_field) => test_field
    end
:aaa
```
"""
macro support_matching(XXX, is_xxx)
    ex = quote
        $__source__
        function $MLStyle.pattern_uncall(
            t::Type{$XXX},
            self::Function,
            type_params::Any,
            type_args::Any,
            args::Any,
        )
            $__source__
            type_infer(_...) = Union{Expr, t}  # type of input
            transform(expr) = :($($is_xxx)($expr) ? $($XXX)($expr) : nothing)
            $MLStyle.AbstractPatterns.P_fast_view(
                type_infer,
                transform,
                $MLStyle.Record._compile_record_pattern(t, self, type_params, type_args, args)
            )
        end
    end
    esc(ex)
end


text:
module NestedStructInitialiser
using RuntimeGeneratedFunctions
using Requires
RuntimeGeneratedFunctions.init(@__MODULE__)

export parameters, initialiser

free_param_length(::Type{<:Number}) = 1
free_param_type(::Type{<:Number}) = true
free_param_length(::Any) = NaN
free_param_type(::Any) = false
free_param_length(::Type{<:NTuple{N}}) where N = N
free_param_type(::Type{<:NTuple}) = true

extendable(ft) = isconcretetype(ft) && !free_param_type(ft)
extendable(ft::Type{<:AbstractArray}) = false
extendable(ft::Type{<:Dict}) = false
extendable(ft::Type{<:Set}) = false

function extend!(free, fixed, val; kwargs...)
    fp = parameters(val; kwargs...)
    if fp !== nothing
        append!(free, fp.free)
        append!(fixed, fp.fixed)
    end
end

struct Parameters
    type
    free
    fixed
end
parameters(::Any; kwargs...) = nothing
function parameters(s::Type; kwargs...)
    free = []
    fixed = []
    for (i, f) in enumerate(fieldnames(s))
        if haskey(kwargs, f)
            val = kwargs[f]
            val === s && error("Value for field $f of $s is $val.")
            push!(fixed, f => val)
            extend!(free, fixed, val; kwargs...)
        else
            ft = fieldtype(s, i)
            if extendable(ft)
                extend!(free, fixed, ft; kwargs...)
            else
                push!(free, f => ft)
            end
        end
    end
    Parameters(s, free, fixed)
end
function Base.show(io::IO, ::MIME"text/plain", p::Parameters)
    println(io, "$(p.type) Parameters")
    println(io, "Fixed")
    if length(p.fixed) == 0
        println(io, " none")
    end
    for (name, value) in p.fixed
        println(io, " $name = $value")
    end
    println(io, "Free")
    if length(p.free) == 0
        println(io, " none")
    end
    for (name, value) in p.free
        d = free_param_length(value)
        ds = isnan(d) ? " (indeterminable dimensionality)" : ""
        println(io, " $name: $value$ds")
    end
end

isleaf(::Number) = true
isleaf(::AbstractArray{<:Number}) = true
isleaf(::NTuple) = true
isleaf(f) = applicable(f)

function free_param(::Type{<:Number}, x, k)
    k[] += 1
    :($x[$(k[]-1)])
end
function free_param(::Type{<:NTuple{N}}, x, k) where N
    k[] += N
    args = [:($x[$i]) for i in k[] - N:k[] - 1]
    :(tuple($(args...)))
end

function field_initialiser(name, type, x, k; kwargs...)
    if haskey(kwargs, name)
        val = kwargs[name]
        isleaf(val) && return applicable(val) ? :($val()) : val
        type_initialiser(val, x, k; kwargs...)
    else
        if free_param_type(type)
            free_param(type, x, k)
        else
            type_initialiser(type, x, k; kwargs...)
        end
    end
end
function type_initialiser(s, x, k; kwargs...)
    :($s($([field_initialiser(f, fieldtype(s, i), x, k; kwargs...)
            for (i, f) in enumerate(fieldnames(s))]...)))
end
function initialiser(s::Type; silent = true, kwargs...)
    k = Ref(1)
    rhs = type_initialiser(s, :x, k; kwargs...)
    silent || println("Number of Free Parameters: $(k[] - 1)")
    @RuntimeGeneratedFunction(:(x -> $rhs))
end
initialiser(p::Parameters) = initialiser(p.type; p.fixed...)

function __init__()
    @require StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" begin
        function free_param(::Type{<:StaticArrays.SArray{S,T,N,L}}, x, k) where {S,T,N,L}
            k[] += L
            :(StaticArrays.SArray{$S}(view($x, $(k[] - L:k[] - 1))))
        end
        free_param_length(::Type{<:StaticArrays.SArray{S,T,N,L}}) where {S,T,N,L} = L
        free_param_type(::Type{<:StaticArrays.SArray}) = true
    end
    @require Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" begin
        function free_param(t::Type{<:Unitful.Quantity{<:Number}}, x, k)
            k[] += 1
            :($t($x[$(k[]-1)]))
        end
    end
end

end # module


text:
using Test

tests = [
    "tsr",
]

@testset "DataGen" begin
    @info("Running tests:")

    for test ∈ tests
        @info("\t* $test ...")
        include("$test.jl")
    end
end


text:
using Rankers.Storages # store_with
using LOLTools: MatchV4, http_action
using LOLToolsDataDragon.DataDragon: en_US, ko_KR

api_key = get(ENV, "RIOT_TOKEN", "")
region = "kr"
matchId = 3999591966 # 3999564765
match = MatchV4.match_by_id(api_key, region, matchId; action=store_with(region, http_action, debug=true));
@info match.participantIdentities[7].player.summonerName
stats = match.participants[7].stats
f(d, k) = iszero(k) ? nothing : getindex(d, k).name
# (println ∘ f).(Ref(en_US.Items), [stats.item0, stats.item1, stats.item2, stats.item3, stats.item4, stats.item5, stats.item6])
(println ∘ f).(Ref(ko_KR.Items), [stats.item0, stats.item1, stats.item2, stats.item3, stats.item4, stats.item5, stats.item6])
using Calendars.Times
@info :gameDuration Time(Nanosecond(Second(match.gameDuration)))
@info :gameCreation millisecond2datetime(match.gameCreation)


text:
using SDFReader
using Test
using Serialization

@testset "SDFReader.jl" begin
    include("sdf.jl")
end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -16.424759279035918 <= q <= 20.23452089871623)

                     add_NL_constraint(m, :(sigmoid(0.8420310557117703 + 0.34198141401583015*sigmoid(-0.5389217244927917 + 0.5798025520034438*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + 0.4428632765655749*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + 0.9033224385695577*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4]))) + -0.1887519896617813*sigmoid(-0.23921485861490321 + 0.5030319708107833*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + -0.2081481501163327*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + -0.2939216839670169*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4]))) + -0.2882490310085508*sigmoid(-0.6545847926105979 + 0.9774628615021363*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + -0.21180887804098658*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + -0.7013494059895882*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4])))) + sigmoid(-0.35188345522271236 + -0.11415509244934263*sigmoid(-0.5389217244927917 + 0.5798025520034438*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + 0.4428632765655749*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + 0.9033224385695577*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4]))) + -0.8088728193024375*sigmoid(-0.23921485861490321 + 0.5030319708107833*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + -0.2081481501163327*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + -0.2939216839670169*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4]))) + 0.8271426964720727*sigmoid(-0.6545847926105979 + 0.9774628615021363*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + -0.21180887804098658*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + -0.7013494059895882*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4])))) + sigmoid(0.4798469380896484 + 0.8134725364327324*sigmoid(-0.5389217244927917 + 0.5798025520034438*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + 0.4428632765655749*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + 0.9033224385695577*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4]))) + -0.12585836968848696*sigmoid(-0.23921485861490321 + 0.5030319708107833*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + -0.2081481501163327*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + -0.2939216839670169*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4]))) + 0.613011911779838*sigmoid(-0.6545847926105979 + 0.9774628615021363*sigmoid(0.8931287257144485 + 0.7044247236000336*$(x[1]) + 0.9746565459142618*$(x[2]) + -0.37297510015038515*$(x[3]) + 0.2762326932574686*$(x[4])) + -0.21180887804098658*sigmoid(-0.08155088834170021 + 0.2657824238897022*$(x[1]) + -0.7808635499515542*$(x[2]) + -0.8071446449659851*$(x[3]) + 0.29051627947702796*$(x[4])) + -0.7013494059895882*sigmoid(0.5572058350137663 + -0.028961919981738138*$(x[1]) + -0.957730339736961*$(x[2]) + -0.8718568123342201*$(x[3]) + -0.17639541567112582*$(x[4])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
using NeRCA
using Test

const DETX = joinpath(@__DIR__, "data", "detx_v3.detx")
const ONLINEFILE = joinpath(@__DIR__, "data", "km3net_online.root")
const OFFLINEFILE = joinpath(@__DIR__, "data", "km3net_offline.root")
const IO_EVT = joinpath(@__DIR__, "data", "IO_EVT.dat")

@testset "calibration" begin
    calib = Calibration(DETX)

    @test 23 == calib.det_id
    @test 6 == length(values(calib.pos))
    @test 6 == length(values(calib.dir))
    @test 6 == length(values(calib.t0))
    @test 6 == length(values(calib.du))
    @test 6 == length(values(calib.floor))
    @test 6 == length(values(calib.omkeys))
    @test 23.9 ≈ calib.max_z
    @test 2 == calib.n_dus
    @test 3.6 ≈ calib.pos[3][2].z

    @test 4 == omkey2domid(calib, 2, 1)
    @test 4 == omkey2domid(calib, OMKey(2, 1))
end

@testset "km3net online files" begin
    f = NeRCA.OnlineFile(ONLINEFILE)
    hits = NeRCA.read_snapshot_hits(f)
    @test 3 == length(hits)  # grouped by event
    @test 96 == length(hits[1])
    @test [806451572, 806451572, 806455814] == [h.dom_id for h in hits[1][1:3]]
    @test [809524432, 809524432, 809544061] == [h.dom_id for h in hits[1][end-2:end]]
    @test [30733918, 30733916, 30733256] == [h.t for h in hits[1][1:3]]
    @test [30733864, 30734686, 30735112] == [h.t for h in hits[1][end-2:end]]
    @test [10, 13, 0, 3, 1] == [h.channel_id for h in hits[1][1:5]]
    @test [10, 10, 22, 24, 17] == [h.channel_id for h in hits[1][end-4:end]]
    @test [26, 19, 25, 22, 28] == [h.tot for h in hits[1][1:5]]
    @test [6, 10, 29, 28, 27] == [h.tot for h in hits[1][end-4:end]]
    @test 124 == length(hits[2])
    @test [806455814, 806483369, 806483369] == [h.dom_id for h in hits[2][1:3]]
    @test [809521500, 809526097, 809526097, 809544058, 809544061] == [h.dom_id for h in hits[2][end-4:end]]
    @test [58728018, 58728107, 58729094] == [h.t for h in hits[2][1:3]]
    @test [58729410, 58729741, 58729262] == [h.t for h in hits[2][end-2:end]]
    @test [15, 5, 14, 23, 9] == [h.channel_id for h in hits[2][1:5]]
    @test [17,  5, 18, 24,  8] == [h.channel_id for h in hits[2][end-4:end]]
    @test [27, 24, 21, 17, 22] == [h.tot for h in hits[2][1:5]]
    @test [21, 23, 25, 27, 27] == [h.tot for h in hits[2][end-4:end]]
    @test 78 == length(hits[3])
    @test [806451572, 806483369, 806483369] == [h.dom_id for h in hits[3][1:3]]
    @test [809526097, 809526097, 809526097, 809544058, 809544061] == [h.dom_id for h in hits[3][end-4:end]]
    @test [63512204, 63511134, 63512493] == [h.t for h in hits[3][1:3]]
    @test [63511894, 63511798, 63512892] == [h.t for h in hits[3][end-2:end]]
    @test [4, 9, 5, 17, 20] == [h.channel_id for h in hits[3][1:5]]
    @test [5,  7, 24, 23, 10] == [h.channel_id for h in hits[3][end-4:end]]
    @test [26, 29, 30, 23, 30] == [h.tot for h in hits[3][1:5]]
    @test [28, 11, 27, 24, 23] == [h.tot for h in hits[3][end-4:end]]

    thits = NeRCA.read_triggered_hits(f)
    @test 3 == length(thits)
    @test 18 == length(thits[1])
    @test 53 == length(thits[2])
    @test 9 == length(thits[3])

    headers = NeRCA.read_headers(f)
    @test length(headers) == 3
    for header in headers
        @test header.run == 6633
        @test header.detector_id == 44
        @test header.UTC_seconds == 0x5dc6018c
    end
    @test headers[1].frame_index == 127
    @test headers[2].frame_index == 127
    @test headers[3].frame_index == 129
    @test headers[1].UTC_16nanosecondcycles == 0x029b9270
    @test headers[2].UTC_16nanosecondcycles == 0x029b9270
    @test headers[3].UTC_16nanosecondcycles == 0x035a4e90
    @test headers[1].trigger_counter == 0
    @test headers[2].trigger_counter == 1
    @test headers[3].trigger_counter == 0
    @test headers[1].trigger_mask == 22
    @test headers[2].trigger_mask == 22
    @test headers[3].trigger_mask == 4
    @test headers[1].overlays == 6
    @test headers[2].overlays == 21
    @test headers[3].overlays == 0
end

@testset "DAQ readout" begin
    f = open(IO_EVT)
    daqevent = read(f, NeRCA.DAQEvent; legacy=true)
    @test 49 == daqevent.det_id
    @test 8142 == daqevent.run_id
    @test 107050 == daqevent.timeslice_id
    @test 1591693105 == daqevent.timestamp
    @test 0 == daqevent.ticks
    @test 9789 == daqevent.trigger_counter
    @test 2 == daqevent.trigger_mask
    @test 0 == daqevent.overlays
    @test 6 == daqevent.n_triggered_hits
    @test 6 == length(daqevent.triggered_hits)
    @test 106 == daqevent.n_hits
    @test 106 == length(daqevent.hits)
    @test NeRCA.is3dshower(daqevent)
    @test !NeRCA.ismxshower(daqevent)
    @test !NeRCA.is3dmuon(daqevent)
    @test !NeRCA.isnb(daqevent)
    close(f)
end


text:
# Autogenerated wrapper script for catch22_jll for x86_64-apple-darwin
export ccatch22

JLLWrappers.@generate_wrapper_header("catch22")
JLLWrappers.@declare_library_product(ccatch22, "@rpath/libcatch22.dylib")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        ccatch22,
        "lib/libcatch22.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
"""
Calculate the resistance of a full turn of wire, given the nominal diameter `d`.

    resistance(d::Unitful.Length, awg::AWG)
    resistance(d::Unitful.Length, Ω::Unitful.ElectricalResistance)

"""
resistance(d::Unitful.Length, awg::AWG) = resistance(d,awg.Ω)
resistance(d::Unitful.Length, Ω::ResistanceLength) = d*pi*Ω |> u"Ω"

"""
isvalidcoil(a::CoilGeometry)

Test whether the coil has valid geometry.
"""
isvalidcoil(a::CoilGeometry) = a.od > a.id

"""
Calculate the maximum number of turns of wire for a given `CoilGeometry`, coil resistance `Ω`, and `fillfactor`.
Optional keyword `AWG_Chart::AbstractDict{AWG}` to specify wire options.

    optimalcoil(coil::CoilGeometry,Ω::Unitful.ElectricalResistance,fillfactor::Real=1.0; AWG_Chart=AWG_Chart)

"""
function optimalcoil(coil::CoilGeometry,Ω::Unitful.ElectricalResistance,fillfactor::Real=1.0; AWG_Chart=AWG_Chart)
    @assert 0 < fillfactor <= 1
    y=[(gauge=gauge,estimatetruefill(coil,awg,Ω,fillfactor)...,fillfactor=fillfactor) for (gauge,awg) in AWG_Chart]
    filter!(x->x.Ω >= Ω,y)
    ind=last(findmax(map(x->x.turns,y)))
    y[ind]
end

"""
Calculate the maximum number of turns of wire for a given `CoilGeometry`, coil voltage `v`, power `p`, and `fillfactor`.
Optional keyword `AWG_Chart::AbstractDict{AWG}` to specify wire options.

    optimalcoil(coil::CoilGeometry,Ω::Unitful.ElectricalResistance,fillfactor::Real=1.0; AWG_Chart=AWG_Chart)

"""
optimalcoil(coil::CoilGeometry,v::Unitful.Voltage,p::Unitful.Power,fillfactor::Real=1.0; AWG_Chart=AWG_Chart) =
    optimalcoil(coil,u"Ω"(v^2/p),fillfactor;AWG_Chart=AWG_Chart)

estimatetruefill(coil::CoilGeometry, awg::AWG, Ω::Unitful.ElectricalResistance,fillfactor::Number=1.0;
    maxturns::Number=Inf) = estimatetruefill(coil,awg;Ω=Ω, fillfactor=fillfactor, maxturns=maxturns)

"""
Estimate the resulting winding given a `coil` and `gauge` of wire, and `fillfactor`.
Optionally limit resistance resistance `Ω` and turn count `maxturns`.

    estimatetruefill(coil::CoilGeometry, awg::AWG; Ω::Unitful.ElectricalResistance=Inf*u"Ω",fillfactor::Number=1.0, maxturns::Number=Inf)
"""
function estimatetruefill(coil::CoilGeometry, awg::AWG; Ω::Unitful.ElectricalResistance=Inf*u"Ω",
    fillfactor::Number=1.0, maxturns::Number=Inf)

    f=ideal_fill(coil,awg,fillfactor)
    if f.turns==0
        # return (turns=0,Ω=zero(Ω),winding=f,coil=coil)
    end
    maxturns=min(maxturns,f.turns) |> Int
    wire_diam = awg.d_insul / sqrt(fillfactor)
    d0 = coil.id + wire_diam
    dstep = √3*wire_diam
    state = (layers=0, turns=0, Ω=0.0u"Ω")
    for (layer,d) in Iterators.take(enumerate(Iterators.countfrom(d0, dstep)), f.layers)
        layerturns = isodd(layer) ? f.odd_layer : f.even_layer
        layerturns = Int(min(maxturns-state.turns,layerturns)) # turns due to turn limit)

        turnΩ = resistance(d,awg)
        Ωturns = isinf(Ω) ? layerturns : (Ω-state.Ω)/turnΩ |> ceil |> Int # turns due to resistance limit

        layerturns = min(layerturns, Ωturns)

        #update State
        state = (layers=layer,
            turns=state.turns+layerturns,
            Ω=state.Ω+layerturns*turnΩ)

        if (state.turns >= maxturns) || (state.Ω ≳ Ω)
            break
        end
    end
    dfinal = coil.id+begin
    #TODO# Replace with enclosewinding function
        if state.layers <= 1
            2wire_diam
        else
            2wire_diam + sqrt(3)*wire_diam*(state.layers-1)
        end
    end
    coil_section=(coil.od-coil.id)/2*coil.len
    relativefill= (state.turns * π*(wire_diam/2)^2)/coil_section |> upreferred
    truefill= (state.turns * π*(awg.d_insul/2)^2)/coil_section |> upreferred
    (turns=state.turns,Ω=state.Ω,winding=f,
        coil=CoilGeometry(coil.id,unit(coil.od)(dfinal),coil.len),
        relativefill=relativefill,truefill=truefill,
        coil_section=coil_section,
        wire_cross=pi*(wire_diam/2)^2,)
end

function optimalgeometry(coil::CoilGeometry, Ω::Unitful.ElectricalResistance, fillfactor::Real=1.0;
    rng=0u"mm":0.01u"mm":2u"mm")

    f = insul -> let
        newcoil = CoilGeometry(coil; od=coil.od-2insul)
        len=gaugebreakpoint(newcoil,Ω,newcoil.len,fillfactor)
        err = 2abs(coil.len-len) + insul
        (insul=insul,len=len,error=err)
    end
    y=tmap(f,rng)
    y[last(findmin(map(x->x.error,y)))]
end