text:
# test functionality and identities
using Polylogarithms
using SpecialFunctions
using Test
using DataFrames, CSV
import Base.MathConstants: π, pi, ℯ, e, γ, eulergamma, catalan, φ, golden
include("test_defs.jl")
Q = Polylogarithms.Q

@testset "Polylogarithm polylog identies and special values" begin

    # check throws errors
    @testset "    throws errors" begin
        # @test_throws DomainError polylog(1, 0)              # should work for all inputs
        # @test_throws MethodError polylog( 1, Float32(0.0))  # should work for all numbers
        @test_throws MethodError polylog( 1, "1.0")  # should work for all numbers
        @test_throws MethodError polylog( "1.0", 1)  # should work for all numbers
        @test_throws MethodError polylog( 1 )
        # @test_throws ArgumentError polylog( 1 )
    end

    # check output types
    @testset "    output types" begin
        @test typeof( polylog(1, 0) ) == Float64
        @test typeof( polylog(complex(1), 0) ) == Complex{Float64}
        # we need to do some more testing here and consider carefully how consistent this should be
    end

    # check keywords work
    @testset "    output types" begin
        @test polylog(-1, 0; level=1, accuracy=1.0e-5,  min_iterations=1, max_iterations=100) ≈ 0.0
    end

    # check subsidary routine errors, just to make coverage cleaner
    @testset "    unexported function errors" begin
        @test_throws DomainError Polylogarithms.polylog_series_1(1.0, 2.0)
        @test_throws DomainError Polylogarithms.polylog_series_1(1.0, 0.75)
        @test_throws DomainError Polylogarithms.polylog_series_2(1.0, 0.0)
        @test_throws DomainError Polylogarithms.polylog_series_3(1.0, 0.0)
        @test_throws DomainError Polylogarithms.polylog_series_3(-1.0, 0.6)
        @test_throws DomainError Polylogarithms.g_crandall(-1)
        @test_throws DomainError Polylogarithms.g_crandall(15)
        τ = 0.00001
        @test_throws DomainError Polylogarithms.Q_closed(0, τ, 0; n_terms=0)
        @test_throws DomainError Polylogarithms.Q_closed(0, τ, 0; n_terms=4)
        @test_throws DomainError Polylogarithms.Q(0, τ, 0; n_terms=0)
        @test_throws DomainError Polylogarithms.Q(0, τ, 0; n_terms=8)
     end

    @testset "    unexported function values" begin
        ℒ = 0.1
        d2 = SpecialFunctions.digamma(1) - ℒ
        @test Polylogarithms.c_closed(0, 0, ℒ) ≈ harmonic(0) - ℒ
        @test Polylogarithms.c_closed(0, 1, ℒ) ≈ -stieltjes(1)   - d2^2/2      - π^2/6 + SpecialFunctions.polygamma(1,1)/2
        @test Polylogarithms.c_closed(0, 2, ℒ) ≈  stieltjes(2)/2 + d2^3/6 + d2*( π^2/6 - SpecialFunctions.polygamma(1,1)/2 ) +
                                                       SpecialFunctions.polygamma(2,1)/6
        for n=0:3
            @test Polylogarithms.c_closed(n, 0, ℒ) ≈ Polylogarithms.c_crandall(n, 0, ℒ)
            @test Polylogarithms.c_closed(n, 1, ℒ) ≈ Polylogarithms.c_crandall(n, 1, ℒ)
            @test Polylogarithms.c_closed(n, 2, ℒ) ≈ Polylogarithms.c_crandall(n, 2, ℒ)
        end
        for k = 1:5
            @test Polylogarithms.f_crandall(k,1) == -harmonic(k)
        end
        @test Polylogarithms.g_crandall(1) ≈ -γ

        τ = 0.00001
        n = 1
        @test Polylogarithms.Q_closed(n, τ, ℒ; n_terms=3) == Polylogarithms.Q(n, τ, ℒ; n_terms=3)
    end

    # check it can handle all types of inputs
    @testset "    input formats" begin
        @test polylog(-1.0, 0.0) ≈ 0.0
        @test polylog(-1,   0.0) ≈ 0.0
        @test polylog(-1.0, Complex(0.0)) ≈ 0.0
        @test polylog(Complex(-1.0), 0.0) ≈ 0.0
        @test polylog(Complex(-1.0), 0.0) ≈ 0.0
        @test polylog(1, 0) ≈ 0.0
        @test polylog(1.0, 0) ≈ 0.0
        @test polylog(1, 0.0) ≈ 0.0
        @test polylog(1.0, 0.0) ≈ 0.0
        @test polylog(1 // 2, 0) ≈ 0.0
        @test polylog(1 // 2, 1 // 2) ≈ 0.8061267230428526
        @test polylog( π, π ) ≈ 3.920844767506471 - 1.8340224080910017im
        x = collect(0.0:0.1:0.9)
        @test all([polylog.(1, x)[i] ≈ polylog(1, x[i]) for i=1:length(x)])
        @test polylog(Complex(-1.0), Complex(0.3)) ≈ polylog(-1.0, 0.3)
        @test polylog(Complex(-1.0), Complex(0.3)) ≈ polylog(-1.0, 0.3)
    end

    # check specific known values
    @testset "     s = n (a real integer)" begin
        # simple cases
        @test polylog(1, 0.5) ≈ log(2)
        @test !isfinite( polylog(0.99999, 1.0) ) # the singularity
        @test !isfinite( polylog(-1, 1.0) )      # the singularity
        @test !isfinite( polylog(0.5 + 0.5im, 1.0) )     # the singularity
        @test polylog(1, 2) ≈ -pi*im # comes from -log(complex(1-2))

        @testset "     dilogarithm for real z" begin
            @test polylog(2,-1.0)    ≈ -pi^2/12.0
            @test polylog(2, 0.0)    ≈ 0.0
            @test polylog(2, 0.5)    ≈ pi^2/12 - 0.5*log(2)^2
            @test polylog(2, 1.0)    ≈ pi^2/6.0
            @test polylog(2, 2.0)    ≈ pi^2/4 - im*pi*log(2)
            @test polylog(2, -φ)     ≈ -pi^2/10 - log(φ)^2
            @test polylog(2, -1/φ)   ≈ -pi^2/15 + log(φ)^2/2
            @test polylog(2, 1/φ^2)  ≈  pi^2/15 - log(φ)^2
            @test polylog(2, 1/φ)    ≈  pi^2/10 - log(φ)^2
            @test polylog(2, φ)      ≈  11*pi^2/15 + log(Complex(-1/φ))^2/2 # wiki has this one, but no ref
            @test polylog(2, φ^2)    ≈ -11*pi^2/15 - log(Complex(-φ))^2

            # identities for dilogarithm
            Z = [3.0 + 0.4im, -3.0 + 0.4im, 3.0 - 0.4im, -3.0 + -0.4im]
            for i=1:length(Z)
                z = Z[i]
                @test polylog(2, z) + polylog(2, 1/z) ≈ -pi^2/6.0 - log(Complex(-z))^2/2.0
            end
        end

        @testset "     trilogarithm for real z" begin
            # https://mathworld.wolfram.com/Trilogarithm.html
            @test polylog(3,-1.0)             ≈ -3*zeta(3)/4
            @test polylog(3, 0.0)             ≈ 0.0
            @test polylog(3, 0.5)             ≈ log(2)^3/6.0 - pi^2*log(2)/12.0 + (7.0/8.0)*zeta(3)
            @test polylog(3, 1.0)             ≈ zeta(3)
            @test polylog(3, Float64(φ)^(-2)) ≈ 4*zeta(3)/5 + 2*log(φ)^3/3 - 2*pi^2*log(φ)/15
        end

        @testset "     general case for real z" begin
            X = collect(-3.0:0.1:3.0)
            for i=1:length(X)
                x = X[i]
                # println("x = $x")
                @test polylog(1, x) ≈ -log(Complex(1-x))
                @test polylog(0, x) ≈ x ./ (1-x)
                @test polylog(-1, x) ≈ x ./ (1-x).^2
                @test polylog(-2, x) ≈ x .* (1+x) ./ (1-x).^3
                @test polylog(-3, x) ≈ x .* (1+4*x+x.^2) ./ (1-x).^4
                @test polylog(-4, x) ≈ x .* (1+x) .* (1+10*x+x.^2) ./ (1-x).^5
            end
        end

        @testset "     general case for complex z" begin
            X = collect(-3.0:0.5:3.0)
            Y = [-1.3, -0.4, 0.4, 1.5]
            for i=1:length(X)
                for j=1:length(Y)
                    z = Complex(X[i], Y[j])
                    # println("z = $z")
                    @test polylog(1, z) ≈ -log(Complex(1-z))
                    @test polylog(0, z) ≈ z ./ (1-z)
                    @test polylog(-1, z) ≈ z ./ (1-z).^2
                    @test polylog(-2, z) ≈ z .* (1+z) ./ (1-z).^3
                    @test polylog(-3, z) ≈ z .* (1+4*z+z.^2) ./ (1-z).^4
                    @test polylog(-4, z) ≈ z .* (1+z) .* (1+10*z+z.^2) ./ (1-z).^5
                end
            end
        end
    end

    @testset "    particular values |z| == 1" begin
        S_r = [2.1 2.5 3.0]
        S_i = [-1.3, -1.0, -0.5, 0.0, 0.5, 1.0, 1.3]
        for i=1:length(S_r)
            for j=1:length(S_i)
                s = Complex(S_r[i], S_i[j])
                @test polylog(s, -1.0) ≈ -eta(s)
                @test polylog(s,  1.0) ≈  zeta(s)
                @test polylog(s,  im) ≈  - 2.0^(-s)*eta(s) + im*dirichlet_beta(s)
            end
        end
    end

    @testset "    additional Identities" begin
        z = 0.5
        for n=1:5
            @test polylog(-n,z) + (-1)^n * polylog(-n, 1/z) ≈ 0.0
        end

        # for real s, and real z<1, polylog should be real
        S = [-1, 0.1, 2]
        Z = [-2, -1.0, 0.1, 0.95]
        for i=1:length(S)
            for j=1:length(Z)
                s = S[i]
                z = Z[j]
                # println("s = $s; z = $z")
                @test abs( imag( polylog(s,  z) ) ) < 1.0e-12 # actually they usually do better than this
            end
        end

        # for real s, and real z>=1, the imaginary part is given
        S = [-1.5, 0.1, 2]
        Z = [1.05, 3.0]
        for i=1:length(S)
            for j=1:length(Z)
                s = S[i]
                z = Z[j]
                # println("s = $s; z = $z")
                μ = log(z)
                @test imag( polylog(s,  z) ) ≈ -pi*μ^(s-1)/gamma(s)
            end
        end

        # duplication formula (or square formula)
        S_r = [2.1 2.5 3.0]
        S_i = [-1.3, -1.0, -0.5, 0.0, 0.5, 1.0, 1.3]
        z = 3.0 - 0.1im
        for i=1:length(S_r)
            for j=1:length(S_i)
                s = Complex(S_r[i], S_i[j])
                # println("s = $s; z = $z")
                @test polylog(s,z) + polylog(s,-z) ≈ complex(2)^(1-s) * polylog(s, z^2)
            end
        end
    end
end


text:
#####
##### Utilities for distributed game simulation
#####

# This code factorizes out all the complexity of:
#   1. Batching inference requests across game simulations
#   2. Distributing simulations across multiple machines (optional)

#####
##### Inference server
#####

# """
#    fill_and_evaluate(net, batch; batch_size, fill_batches)

# Evaluate a neural network on a batch of inputs.

# If `fill_batches=true`, the batch is padded with dummy inputs until
# it has size `batch_size` before it is sent to the network.

# This function is typically called by inference servers.
# """
function fill_and_evaluate(net, batch; batch_size, fill_batches)
  n = length(batch)
  @assert n > 0
  if !fill_batches
    return Network.evaluate_batch(net, batch)
  else
    nmissing = batch_size - n
    @assert nmissing >= 0
    if nmissing == 0
      return Network.evaluate_batch(net, batch)
    else
      batch = vcat(batch, [batch[1] for _ in 1:nmissing])
      return Network.evaluate_batch(net, batch)[1:n]
    end
  end
end

# """
#     launch_inference_server(net; num_workers, fill_batches)

# Start a server that processes inference requests for a single neural network.

# Return a request channel (see [`Batchifier.launch_server`](@ref)).
# """
function launch_inference_server(
    net::AbstractNetwork;
    num_workers,
    batch_size,
    fill_batches)

  return Batchifier.launch_server(;num_workers, batch_size) do batch
    fill_and_evaluate(net, batch; batch_size, fill_batches)
  end
end

# """
#     launch_inference_server(net1, net2, num_workers; fill_batches)
#
# Start a server that processes inference requests for TWO networks.
# This is needed when pitting two players together that rely on distinct networks for
# example. The reason we do not spawn two separate servers in this case is that each
# server would not know how many queries to wait for.
#
# Such a server expects modified queries that are named tuples with field `query` and
# `netid`. The latter must be an integer in {1, 2} indicating whether `net1` or `net2`
# is queried.
# """
function launch_inference_server(
    net1::AbstractNetwork,
    net2::AbstractNetwork;
    num_workers,
    batch_size,
    fill_batches)

  return Batchifier.launch_server(;num_workers, batch_size) do batch
    n = length(batch)
    mask1 = findall(b -> b.netid == 1, batch)
    mask2 = findall(b -> b.netid == 2, batch)
    @assert length(mask1) + length(mask2) == n
    state(x) = x.state # TODO: this is ugly
    batch1 = state.(batch[mask1])
    batch2 = state.(batch[mask2])
    if isempty(mask2) # there are only queries from net1
      return fill_and_evaluate(net1, batch1; batch_size=n, fill_batches)
    elseif isempty(mask1) # there are only queries from net2
      return fill_and_evaluate(net2, batch2; batch_size=n, fill_batches)
    else # both networks sent queries
      res1 = fill_and_evaluate(net1, batch1; batch_size=n, fill_batches)
      res2 = fill_and_evaluate(net2, batch2; batch_size=n, fill_batches)
      @assert typeof(res1) == typeof(res2)
      res = similar(res1, n)
      res[mask1] = res1
      res[mask2] = res2
      return res
    end
  end
end

# """
#     batchify_oracles(oracles; num_workers, batch_size, fill_batches)

# Take some oracles, launch a server to call them by batches if needed and return a pair
# of functions to be called by each concurrent worker:

#   - A `spawn_oracles()` function that expects no argument and returns `BatchedOracle`s
#     to be used by the worker.
#   - A `done!()` function to be called when the worker is finished and won't make anymore
#     query.

# The `oracles` argument can be either an oracle or a pair of oracles. If one of the two
# oracles at least is a neural network, an inference server will be launched using
# [`launch_inference_server`](@ref).
# """
function batchify_oracles() end

ret_oracle(x) = () -> x
do_nothing!() = nothing
send_done!(reqc) = () -> Batchifier.client_done!(reqc)
zipthunk(f1, f2) = () -> (f1(), f2())

# Two neural network oracles
function batchify_oracles(os::Tuple{AbstractNetwork, AbstractNetwork}; kwargs...)
  reqc = launch_inference_server(os[1], os[2]; kwargs...)
  make1() = Batchifier.BatchedOracle(reqc, st -> (state=st, netid=1))
  make2() = Batchifier.BatchedOracle(reqc, st -> (state=st, netid=2))
  return zipthunk(make1, make2), send_done!(reqc)
end

function batchify_oracles(os::Tuple{<:Any, AbstractNetwork}; kwargs...)
  reqc = launch_inference_server(os[2]; kwargs...)
  make2() = Batchifier.BatchedOracle(reqc)
  return zipthunk(ret_oracle(os[1]), make2), send_done!(reqc)
end

function batchify_oracles(os::Tuple{AbstractNetwork, <:Any}; kwargs...)
  reqc = launch_inference_server(os[1]; kwargs...)
  make1() = Batchifier.BatchedOracle(reqc)
  return zipthunk(make1, ret_oracle(os[2])), send_done!(reqc)
end

function batchify_oracles(os::Tuple{<:Any, <:Any}; kwargs...)
  return zipthunk(ret_oracle(os[1]), ret_oracle(os[2])), do_nothing!
end

function batchify_oracles(o::AbstractNetwork; kwargs...)
  reqc = launch_inference_server(o; kwargs...)
  make() = Batchifier.BatchedOracle(reqc)
  return make, send_done!(reqc)
end

function batchify_oracles(o::Any; kwargs...)
  return ret_oracle(o), do_nothing!
end

#####
##### Distributed simulator
#####

"""
    Simulator(make_player, make_oracles, measure)

A distributed simulator that encapsulates the details of running simulations
across multiple threads and multiple machines.

# Arguments

  - `make_oracles`: a function that takes no argument and returns
      the oracles used by the player, which can be either `nothing`,
      a single oracle or a pair of oracles.
  - `make_player`: a function that takes as an argument the result of `make_oracles`
      and builds a player from it. In practice, an oracle returned by `make_oracles`
      may be replaced by a `BatchedOracle` before it is passed to `make_player`, which
      is why these two functions are specified separately.
  - `measure(trace, colors_flipped, player)`: the function that is used to
      take measurements after each game simulation.
"""
struct Simulator{MakePlayer, Oracles, Measure}
  make_player :: MakePlayer
  make_oracles :: Oracles
  measure :: Measure
end

"""
    record_trace

A measurement function to be passed to a [`Simulator`](@ref) that produces
named tuples with two fields: `trace::Trace` and `colors_flipped::Bool`.
"""
record_trace(t, cf, p) = (trace=t, colors_flipped=cf)

"""
    simulate(::Simulator, ::AbstractGameSpec; ::SimParams; <kwargs>)

Play a series of games using a given [`Simulator`](@ref).

# Keyword Arguments

  - `game_simulated` is called every time a game simulation is completed
    (with no arguments)

# Return

Return a vector of objects computed by `simulator.measure`.
"""
function simulate(
    simulator::Simulator,
    gspec::AbstractGameSpec,
    p::SimParams;
    game_simulated)

  oracles = simulator.make_oracles()
  spawn_oracles, done =
    batchify_oracles(oracles; p.num_workers, p.batch_size, p.fill_batches)
  return Util.mapreduce(1:p.num_games, p.num_workers, vcat, []) do
    oracles = spawn_oracles()
    player = simulator.make_player(oracles)
    worker_sim_id = 0
    # For each worker
    function simulate_game(sim_id)
      worker_sim_id += 1
      # Switch players' colors if necessary: "_pf" stands for "possibly flipped"
      if isa(player, TwoPlayers) && p.alternate_colors
        colors_flipped = sim_id % 2 == 1
        player_pf = colors_flipped ? flipped_colors(player) : player
      else
        colors_flipped = false
        player_pf = player
      end
      # Play the game and generate a report
      trace = play_game(gspec, player_pf, flip_probability=p.flip_probability)
      report = simulator.measure(trace, colors_flipped, player)
      # Reset the player periodically
      if !isnothing(p.reset_every) && worker_sim_id % p.reset_every == 0
        reset_player!(player)
      end
      # Signal that a game has been simulated
      game_simulated()
      return report
    end
    return (process=simulate_game, terminate=done)
  end
end

"""
    simulate_distributed(::Simulator, ::AbstractGameSpec, ::SimParams; <kwargs>)

Identical to [`simulate`](@ref) but splits the work across all available distributed
workers, whose number is given by `Distributed.nworkers()`.
"""
function simulate_distributed(
    simulator::Simulator,
    gspec::AbstractGameSpec,
    p::SimParams;
    game_simulated)

  # Spawning a task to keep count of completed simulations
  chan = Distributed.RemoteChannel(()->Channel{Nothing}(1))
  Util.@tspawn_main begin
    for i in 1:p.num_games
      take!(chan)
      game_simulated()
    end
  end
  remote_game_simulated() = put!(chan, nothing)
  # Distributing the simulations across workers
  num_each, rem = divrem(p.num_games, Distributed.nworkers())
  @assert num_each >= 1
  workers = Distributed.workers()
  tasks = map(workers) do w
    Distributed.@spawnat w begin
      Util.@printing_errors begin
        simulate(
          simulator,
          gspec,
          SimParams(p; num_games=(w == workers[1] ? num_each + rem : num_each)),
          game_simulated=remote_game_simulated)
        end
    end
  end
  results = fetch.(tasks)
  # If one of the worker raised an exception, we print it
  for r in results
    if isa(r, Distributed.RemoteException)
      showerror(stderr, r, catch_backtrace())
    end
  end
  return reduce(vcat, results)
end

function compute_redundancy(states)
  unique = Set(states)
  return 1. - length(unique) / length(states)
end

# From a vector of measures, extract all rewards w.r.t. white and compute redundancy
# `samples` is a vector of named tuples with fields `trace` and `colors_flipped`
function rewards_and_redundancy(samples; gamma)
  rewards = map(samples) do s
    wr = total_reward(s.trace, gamma)
    return s.colors_flipped ? -wr : wr
  end
  states = [st for s in samples for st in s.trace.states]
  redundancy = compute_redundancy(states)
  return rewards, redundancy
end

ext: jl


text:
using StatsBase

# Distance function -------------------------------------------------

"""
    euclidean(a::AbstractVector, b::AbstractVector)

Calculate euclidean distance from two vectors. √∑(aᵢ - bᵢ)²
"""
function euclidean(a::AbstractVector{T},
                   b::AbstractVector{T}) where {T<:AbstractFloat}
    @assert length(a) == length(b)

    # euclidean(a, b) = √∑(aᵢ- bᵢ)²
    s = zero(T)
    @simd for i in eachindex(a)
        @inbounds s += (a[i] - b[i])^2
    end
    return √s
end

# KnnClassifier -----------------------------------------------------
struct KnnClassifier{T<:Real, S}
    X::AbstractMatrix{T}
    y::AbstractVector{S}
    K::Int

    # Internal constructor
    function KnnClassifier(X::AbstractMatrix{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S}
        @assert size(X, 1) == size(y, 1)
        return new{T, S}(X, y, K)
    end
end

# External constructor
function KnnClassifier(Xv::AbstractVector{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S}
    X = reshape(Xv, length(Xv), 1)
    return KnnClassifier(X, y, K)
end

function predict(model::KnnClassifier{T, S}, input::AbstractVector{T}) where {T<:Real, S}
    nl, nc = size(model.X, 1), size(model.X, 2)

    @assert size(input, 1) == nc

    # calculate all the distances
    dists = Vector{T}(undef, nl)
    for i in axes(model.X, 1)
        dists[i] = euclidean(input, model.X[i, :])
    end

    # sort index of distances and select the k nearest indexes
    idx = sortperm(dists)[1:model.K]

    # select the KNearestNeighbours
    kneighbours = model.y[idx]

    # get prediction
    return mode(kneighbours)
end

function predict(model::KnnClassifier{T, S}, input::AbstractMatrix{T}) where {T<:Real, S}
    nl, nc = size(model.X, 1), size(model.X, 2)
    nli = size(input, 1)

    @assert size(input, 2) == nc

    results = Vector{S}(undef, nli)
    for j in axes(input, 1)

        # calculate all the distances
        dists = Vector{T}(undef, nl)
        for i in axes(model.X, 1)
            dists[i] = euclidean(input[j, :], model.X[i, :])
        end

        # sort index of distances and select the k nearest indexes
        idx = sortperm(dists)[1:model.K]

        # select the KNearestNeighbours
        kneighbours = model.y[idx]

        # get prediction
        results[j] = mode(kneighbours)
    end

    return results
end

# KnnRegression -----------------------------------------------------
struct KnnRegression{T<:Real, S<:Real}
    X::AbstractMatrix{T}
    y::AbstractVector{S}
    K::Int

    # Internal constructor
    function KnnRegression(X::AbstractMatrix{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S<:Real}
        @assert size(X, 1) == size(y, 1)
        return new{T, S}(X, y, K)
    end
end

# External constructor
function KnnRegression(Xv::AbstractVector{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S<:Real}
    X = reshape(Xv, length(Xv), 1)
    return new{T, S}(X, y, K)
end

function predict(model::KnnRegression{T, S}, input::AbstractVector{T}, w::Bool=false) where {T<:Real, S<:Real}
    nl, nc = size(model.X, 1), size(model.X, 2)

    @assert size(input, 1) == nc

    # calculate all the distances
    dists = Vector{T}(undef, nl)
    for i in axes(model.X, 1)
        dists[i] = euclidean(input, model.X[i, :])
    end

    # sort index of distances and select the k nearest indexes
    idx = sortperm(dists)[1:model.K]

    # select the KNearestNeighbours
    y_kneighbours = model.y[idx]

    # get prediction by mean or weighted mean
    if w
        k_min_dists = sort(dists)[1:model.K]
        return mean(y_kneighbours, weights(k_min_dists))
    else
        return mean(y_kneighbours)
    end

end

function predict(model::KnnRegression{T, S}, input::AbstractMatrix{T}, w::Bool=false) where {T<:Real, S<:Real}
    nl, nc = size(model.X, 1), size(model.X, 2)
    nli = size(input, 1)

    @assert size(input, 2) == nc

    results = Vector{S}(undef, nli)
    for j in axes(input, 1)

        # calculate all the distances
        dists = Vector{T}(undef, nl)
        for i in axes(model.X, 1)
            dists[i] = euclidean(input[j, :], model.X[i, :])
        end

        # sort index of distances and select the k nearest indexes
        idx = sortperm(dists)[1:model.K]

        # select the KNearestNeighbours
        y_kneighbours = model.y[idx]

        # get prediction by mean or weighted mean
        if w
            k_min_dists = sort(dists)[1:model.K]
            results[j] = mean(y_kneighbours, weights(k_min_dists))
        else
            results[j] = mean(y_kneighbours)
        end
    end

    return results
end

max_stars_repo_name: AugustoCL/KNN-from-scratch-

max_issues_repo_name: AugustoCL/KNN-from-scratch-

max_forks_repo_name: AugustoCL/KNN-from-scratch-


text:
"""
Construct and evaluate ``Ψ(s)`` for a triditional SFM.

Uses newtons method to solve the Ricatti equation
``D⁺⁻(s) + Ψ(s)D⁻⁺(s)Ψ(s) + Ψ(s)D⁻⁻(s) + D⁺⁺(s)Ψ(s) = 0.``

    PsiFun( model::SFFM.Model; s = 0, MaxIters = 1000, err = 1e-8)

# Arguments
- `model`: a Model object
- `s::Real`: a value to evaluate the LST at
- `MaxIters::Int`: the maximum number of iterations of newtons method
- `err::Float64`: an error tolerance for terminating newtons method. Terminates
    when `max(Ψ_{n} - Ψ{n-1}) .< eps`.

# Output
- `Ψ(s)::Array{Float64,2}` the matrix ``Ψ``
"""
function PsiFunX( model::SFFM.Model; s = 0, MaxIters = 1000, err = 1e-8)
    SDict, TDict = modelDicts(model)

    T00inv = inv(TDict["00"] - s * LinearAlgebra.I)
    # construct the generator Q(s)
    Q =
        (1 ./ abs.(model.C[SDict["bullet"]])) .* (
            TDict["bulletbullet"] - s * LinearAlgebra.I -
            TDict["bullet0"] * T00inv * TDict["0bullet"]
        )

    QDict = Dict{String,Array}("Q" => Q)
    for ℓ in ["+" "-"], m in ["+" "-"]
        QDict[ℓ*m] = Q[SDict[ℓ], SDict[m]]
    end

    Ψ = zeros(Float64, length(SDict["+"]), length(SDict["-"]))
    A = QDict["++"]
    B = QDict["--"]
    D = QDict["+-"]
    # use netwons method to solve the Ricatti equation
    for n in 1:MaxIters
        Ψ = LinearAlgebra.sylvester(A,B,D)
        if maximum(abs.(sum(Ψ,dims=2).-1)) < err
            break
        end
        A = QDict["++"] + Ψ * QDict["-+"]
        B = QDict["--"] + QDict["-+"] * Ψ
        D = QDict["+-"] - Ψ * QDict["-+"] * Ψ
    end

    return Ψ
end

"""
Construct the vector ``ξ`` containing the distribution of the phase at the time
when ``X(t)`` first hits `0`.

    MakeXiX( model::SFFM.Model, Ψ::Array)

# Arguments
- `model`: a Model object
- `Ψ`: an array as output from `PsiFunX`

# Output
- the vector `ξ`
"""
function MakeXiX( model::SFFM.Model, Ψ::Array)
    # the system to solve is [ξ 0](-[B₋₋ B₋₀; B₀₋ B₀₀])⁻¹[B₋₊; B₀₊]Ψ = ξ
    # writing this out and using block inversion (as described on wikipedia)
    # we can solve this in the following way
    SDict, TDict = modelDicts(model)

    T00inv = inv(TDict["00"])
    invT₋₋ =
        inv(TDict["--"] - TDict["-0"] * T00inv * TDict["0-"])
    invT₋₀ = -invT₋₋ * TDict["-0"] * T00inv

    A =
        -(
            invT₋₋ * TDict["-+"] * Ψ + invT₋₀ * TDict["0+"] * Ψ +
            LinearAlgebra.I
        )
    b = zeros(1, size(TDict["--"], 1))
    A[:, 1] .= 1.0 # normalisation conditions
    b[1] = 1.0 # normalisation conditions

    ξ = b / A

    return ξ
end

"""
Construct the stationary distribution of the SFM

    StationaryDistributionX( model::SFFM.Model, Ψ::Array, ξ::Array)

# Arguments
- `model`: a Model object
- `Ψ`: an array as output from `PsiFunX`
- `ξ`: an array as returned from `MakeXiX`

# Output
- `pₓ::Array{Float64,2}`: the point masses of the SFM
- `πₓ(x)` a function with two methods
    - `πₓ(x::Real)`: for scalar inputs, returns the stationary density evaluated
        at `x` in all phases.
    - `πₓ(x::Array)`: for array inputs, returns an array with the same shape
        as is output by Coeff2Dist.
- `K::Array{Float64,2}`: the matrix in the exponential of the density.
"""
function StationaryDistributionX( model::SFFM.Model, Ψ::Array, ξ::Array)
    # using the same block inversion trick as in MakeXiX
    SDict, TDict = modelDicts(model)

    T00inv = inv(TDict["00"])
    invT₋₋ =
        inv(TDict["--"] - TDict["-0"] * T00inv * TDict["0-"])
    invT₋₀ = -invT₋₋ * TDict["-0"] * T00inv

    Q =
        (1 ./ abs.(model.C[SDict["bullet"]])) .* (
            TDict["bulletbullet"] -
            TDict["bullet0"] * T00inv * TDict["0bullet"]
        )

    QDict = Dict{String,Array}("Q" => Q)
    for ℓ in ["+" "-"], m in ["+" "-"]
        QDict[ℓ*m] = Q[SDict[ℓ], SDict[m]]
    end

    K = QDict["++"] + Ψ * QDict["-+"]

    A = -[invT₋₋ invT₋₀]

    # unnormalised values
    αpₓ = ξ * A

    απₓ = αpₓ *
        [TDict["-+"]; TDict["0+"]] *
        -inv(K) *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(model.C[SDict["bullet"]]))

    απₓ0 = -απₓ * [TDict["+0"];TDict["-0"]] * T00inv

    # normalising constant
    α = sum(αpₓ) + sum(απₓ) + sum(απₓ0)

    # normalised values
    # point masses
    pₓ = αpₓ/α
    # density method for scalar x-values
    function πₓ(x::Real)
        pₓ *
        [TDict["-+"]; TDict["0+"]] *
        exp(K*x) *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(model.C[SDict["bullet"]])) *
        [LinearAlgebra.I(sum(model.C .!= 0)) [TDict["+0"];TDict["-0"]] * T00inv]
    end
    # density method for arrays so that πₓ returns an array with the same shape
    # as is output by Coeff2Dist
    function πₓ(x::Array)
        temp = πₓ.(x)
        Evalπₓ = zeros(Float64, size(x,1), size(x,2), NPhases(model))
        for cell in 1:size(x,2)
            for basis in 1:size(x,1)
                Evalπₓ[basis,cell,:] = temp[basis,cell]
            end
        end
        return Evalπₓ
    end

    # CDF method for scalar x-values
    function Πₓ(x::Real)
        [pₓ zeros(1,sum(model.C.>0))] .+
        pₓ *
        [TDict["-+"]; TDict["0+"]] *
        (exp(K*x) - LinearAlgebra.I) / K *
        [LinearAlgebra.I(length(SDict["+"])) Ψ] *
        LinearAlgebra.diagm(1 ./ abs.(model.C[SDict["bullet"]])) *
        [LinearAlgebra.I(sum(model.C .!= 0)) [TDict["+0"];TDict["-0"]] * T00inv]
    end
    # CDF method for arrays so that Πₓ returns an array with the same shape
    # as is output by Coeff2Dist
    function Πₓ(x::Array)
        temp = Πₓ.(x)
        Evalπₓ = zeros(Float64, size(x,1), size(x,2), NPhases(model))
        for cell in 1:size(x,2)
            for basis in 1:size(x,1)
                Evalπₓ[basis,cell,:] = temp[basis,cell]
            end
        end
        return Evalπₓ
    end

    return pₓ, πₓ, Πₓ, K
end

max_stars_repo_name: angus-lewis/SFFM

max_issues_repo_name: angus-lewis/SFFM

max_forks_repo_name: angus-lewis/SFFM


text:
### File for testing the plotting functions.

#Tests plottings for normal one dimensional bifurcation diagrams.
for rn in reaction_networks_standard
    p = 10*rand(length(rn.params))
    p_bif = rand(rn.params)
    p_range = (0.4:0.4:2.0)*p[1]
    bif = bifurcation_diagram(rn,p,p_bif,p_range)
    plot(bif)
    plot!(bif)
    plot(bif,rand(1:length(rn.syms)))
    plot!(bif,rand(1:length(rn.syms)))
    plot(bif,rand(rn.syms))
    plot!(bif,rand(rn.syms))
    plot_bifurcation_diagram(rand(1:length(rn.syms)),rn,p,p_bif,p_range)
    plot_bifurcation_diagram!(rand(1:length(rn.syms)),rn,p,p_bif,p_range)
    plot_bifurcation_diagram(rand(rn.syms),rn,p,p_bif,p_range)
    plot_bifurcation_diagram!(rand(rn.syms),rn,p,p_bif,p_range)
end

#Tests plottings fortwo dimensional bifurcation diagrams.
for rn in reaction_networks_standard
    p = 10*rand(length(rn.params))
    p_bif_1 = rand(rn.params)
    p_range_1 = (1.0:2.0)*p[1]
    p_bif_2 = rand(setdiff(rn.params,Set([p_bif_1])))
    p_range_2 = (1.0:2.0)*p[1]
    bif_2d = bifurcation_diagram_2d(rn,p,p_bif_1,p_bif_2,p_range_1,p_range_2)
    plot(bif_2d)
    plot!(bif_2d)
    plot(bif_2d,rand(1:length(rn.syms)))
    plot!(bif_2d,rand(1:length(rn.syms)))
    plot(bif_2d,rand(rn.syms))
    plot!(bif_2d,rand(rn.syms))
    plot_bifurcation_diagram_2d(rand(1:length(rn.syms)),rn,p,p_bif_1,p_bif_2,p_range_1,p_range_2)
    plot_bifurcation_diagram_2d!(rand(1:length(rn.syms)),rn,p,p_bif_1,p_bif_2,p_range_1,p_range_2)
    plot_bifurcation_diagram_2d(rand(rn.syms),rn,p,p_bif_1,p_bif_2,p_range_1,p_range_2)
    plot_bifurcation_diagram_2d!(rand(rn.syms),rn,p,p_bif_1,p_bif_2,p_range_1,p_range_2)
end

max_stars_repo_name: TorkelE/BRNEquilibrate

max_issues_repo_name: TorkelE/BRNEquilibrate

max_forks_repo_name: TorkelE/BRNEquilibrate


text:
module ShaleDrillingLikelihood_DynamicDrillModel_Test_LearningVersions

DOBTIME = false
PRINTSTUFF = false
DOPROFILE = false

using ShaleDrillingLikelihood
using Test
using SparseArrays
using BenchmarkTools
using Random
# using ProfileView
using InteractiveUtils

using Base.Iterators: product, OneTo

using ShaleDrillingLikelihood: DCDPEmax,
    DCDPTmpVars,
    solve_vf_all_timing!,
    EV,
    solve_vf_all!,
    PerfectInfo,
    MaxLearning,
    _nSexp,
    _ρ,
    end_lrn,
    end_ex0,
    ObservationDrill,
    SimulationDraw,
    flow!,
    vw_revenue,
    learn,
    revenue,
    _Dgt0,
    Eexpψ

println("print to keep from blowing up")

@testset "Versions of Dynamic Drilling Models" begin

    Random.seed!(1234)

    nψ, dmx, nz =  13, 3, 11

    # psi
    ψs = range(-4.5; stop=4.5, length=nψ)

    # z transition
    zs = ( range(-1.5; stop=1.5, length=nz), )
    nz = prod(length.(zs))
    ztrans = spdiagm(-1 => fill(1.0/3, nz-1), 0 => fill(1.0/3, nz), 1 => fill(1.0/3, nz-1) )
    ztrans[1,1] = ztrans[1,2] = ztrans[end,end-1] = ztrans[end,end] = 0.5

    # wp = LeasedProblemContsDrill(dmx,4,5,3,2)
    wp = LeasedProblem(dmx,4,5,3,2)

    # ichars
    ichar = (2.0, 0.25,)

    @testset "Perfect Info" begin
        f1 = DrillReward(DrillingRevenue(Constrained(), NoTrend(), NoTaxes(), Learn(),       WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())
        f2 = DrillReward(DrillingRevenue(Constrained(), NoTrend(), NoTaxes(), PerfectInfo(), WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())

        theta1 = randn(_nparm(f1))
        theta2 = copy(theta1)
        theta1[end] = 20
        theta2[end] = 0

        ddm1 = DynamicDrillModel(f1, 0.9, wp, zs, ztrans, ψs, false)
        ddm2 = DynamicDrillModel(f2, 0.9, wp, zs, ztrans, ψs, false)

        ev1 = ValueFunctionArrayOnly(ddm1)
        ev2 = ValueFunctionArrayOnly(ddm2)
        tmpv = DCDPTmpVars(ddm1)

        fill!(ev1, 0)
        fill!(ev2, 0)

        solve_vf_all!(tmpv, ddm1, theta1, ichar, false)
        solve_vf_all!(tmpv, ddm2, theta2, ichar, false)

        @test EV(ev1) ≈ EV(ev2)
        @test EV(ev1) != EV(ev2)
    end # perfect info

    @testset "Max Learning" begin
        f1 = DrillReward(DrillingRevenue(Constrained(), NoTrend(), NoTaxes(), Learn(),       WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())
        f2 = DrillReward(DrillingRevenue(Constrained(), NoTrend(), NoTaxes(), MaxLearning(), WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())

        theta1 = randn(_nparm(f1))
        theta2 = copy(theta1)
        theta1[end] = -20
        theta2[end] = 0

        ddm1 = DynamicDrillModel(f1, 0.9, wp, zs, ztrans, ψs, false)
        ddm2 = DynamicDrillModel(f2, 0.9, wp, zs, ztrans, ψs, false)

        ev1 = ValueFunctionArrayOnly(ddm1)
        ev2 = ValueFunctionArrayOnly(ddm2)
        tmpv = DCDPTmpVars(ddm1)

        fill!(ev1, 0)
        fill!(ev2, 0)

        solve_vf_all!(tmpv, ddm1, theta1, ichar, false)
        solve_vf_all!(tmpv, ddm2, theta2, ichar, false)

        @test EV(ev1) ≈ EV(ev2)
        @test EV(ev1) != EV(ev2)
    end # max learning

    @testset "NoRoyalty" begin
        f1 = DrillReward(DrillingRevenue(Constrained(), NoTrend(), NoTaxes(), Learn(), WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())
        f2 = DrillReward(DrillingRevenue(Constrained(), NoTrend(), NoTaxes(), Learn(), NoRoyalty()),   DrillingCost_constant(), ExtensionCost_Constant())

        theta1 = randn(_nparm(f1))
        theta2 = copy(theta1)

        ddm1 = DynamicDrillModel(f1, 0.9, wp, zs, ztrans, ψs, false)
        ddm2 = DynamicDrillModel(f2, 0.9, wp, zs, ztrans, ψs, false)

        ev1 = ValueFunctionArrayOnly(ddm1)
        ev2 = ValueFunctionArrayOnly(ddm2)
        tmpv = DCDPTmpVars(ddm1)

        fill!(ev1, 0)
        fill!(ev2, 0)

        solve_vf_all!(tmpv, ddm1, theta1, (2.0, 0.0), false)
        solve_vf_all!(tmpv, ddm2, theta2, (2.0, 0.25), false)

        @test EV(ev1) == EV(ev2)
    end

    @testset "NoLearn" begin

        α0 = -3.0
        αψ0 = 0.35
        sig0 = 0.7
        theta0 = [-5.45746, -0.9, α0, sig0] # α₀, α_cost, α_extend
        ρ0 = _ρ(sig0)
        @test ρ0 < sig0

        α1 = α0 + 0.5*αψ0^2*(1-ρ0^2)
        αψ1 = αψ0*ρ0
        sig1 = 60.0
        theta1 = vcat(theta0[1:2], α1, sig1)

        f0 = DrillReward(DrillingRevenue(Constrained(log_ogip=0.64, α_ψ=αψ0, α_t=0.025), NoTrend(), NoTaxes(), NoLearn(),     WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())
        f1 = DrillReward(DrillingRevenue(Constrained(log_ogip=0.64, α_ψ=αψ1, α_t=0.025), NoTrend(), NoTaxes(), PerfectInfo(), WithRoyalty()), DrillingCost_constant(), ExtensionCost_Constant())

        @test learn(revenue(f0)) == NoLearn()

        @test vw_revenue(f0, theta0) == theta0[3:4]
        @test vw_revenue(f1, theta1) == theta1[3:4]

        ddm0 = DynamicDrillModel(f0, 0.9, wp, zs, ztrans, ψs, false)
        ddm1 = DynamicDrillModel(f1, 0.9, wp, zs, ztrans, ψs, false)

        let z = (2.5, 2.0, 2010),
            grad = similar(theta0),
            d = 1,
            i = length(statespace(ddm0))-1,
            obs0 = ObservationDrill(ddm0, ichar, z, d, i),
            obs1 = ObservationDrill(ddm1, ichar, z, d, i),
            u = randn()
            @test _Dgt0(obs0)
            @test _Dgt0(obs1)

            # @show @which Eexpψ(revenue(f0), d, obs0, theta0, u)
            # @show @which Eexpψ(revenue(f1), d, obs1, theta1, u)

            ff0 = flow!(grad, f0, d, obs0, theta0, u, false)
            ff1 = flow!(grad, f1, d, obs1, theta1, u, false)

            @test ff0 ≈ ff1
        end
        ev0 = ValueFunctionArrayOnly(ddm0)
        ev1 = ValueFunctionArrayOnly(ddm1)
        tmpv = DCDPTmpVars(ddm0)

        fill!(ev0, 0)
        fill!(ev1,0)

        solve_vf_all!(tmpv, ddm0, theta0, ichar, false)
        @test all(isfinite.(EV(ev0)))

        solve_vf_all!(tmpv, ddm1, theta1, ichar, false)
        @test all(isfinite.(EV(ev1)))

        @test EV(ev0) ≈ EV(ev1)
    end

end # testset
end # module


text:
module ClassicalSpinMonteCarlo

using StaticArra
using QuantumLattices

using LinearAlgebra::norm,cross

end # module


text:
function make_public_forecast(forecast::ForecastInternal, d::TimeSeriesData)
    error("$(typeof(forecast)) must implement make_public_forecast")
end

function make_internal_forecast(forecast::Forecast)
    error("$(typeof(forecast)) must implement make_internal_forecast")
end

"""
    Deterministic(label::String,
                  resolution::Dates.Period,
                  initial_time::Dates.DateTime,
                  time_steps::Int)

Constructs Deterministic after constructing a TimeArray from initial_time and time_steps.
"""
function Deterministic(label::String,
                       resolution::Dates.Period,
                       initial_time::Dates.DateTime,
                       time_steps::Int)
    data = TimeSeries.TimeArray(
        initial_time : Dates.Hour(1) : initial_time + resolution * (time_steps-1),
        ones(time_steps)
    )
    return Deterministic(label, data)
end

function Deterministic(forecasts::Vector{Deterministic})
    @assert !isempty(forecasts)
    timestamps = collect(Iterators.flatten((TimeSeries.timestamp(get_data(x))
                                            for x in forecasts)))
    data = collect(Iterators.flatten((TimeSeries.values(get_data(x)) for x in forecasts)))
    ta = TimeSeries.TimeArray(timestamps, data)

    forecast = Deterministic(get_label(forecasts[1]), ta)
    @debug "concatenated forecasts" forecast
    return forecast
end

# TODO: need to make concatenation constructors for Probabilistic and ScenarioBased.

function make_public_forecast(forecast::DeterministicInternal, data::TimeSeries.TimeArray)
    return Deterministic(get_label(forecast), data)
end

function make_internal_forecast(forecast::Deterministic, ts_data::TimeSeriesData)
    return DeterministicInternal(get_label(forecast), ts_data)
end

function DeterministicInternal(label::AbstractString, data::TimeSeriesData)
    return DeterministicInternal(label,
                                 get_resolution(data),
                                 get_initial_time(data),
                                 get_uuid(data),
                                 get_horizon(data))
end

"""
    Probabilistic(
                  label::String,
                  resolution::Dates.Period,
                  initial_time::Dates.DateTime,
                  percentiles::Vector{Float64},
                  time_steps::Int,
                 )
Constructs Probabilistic after constructing a TimeArray from initial_time and time_steps.
"""
function Probabilistic(
                       label::String,
                       resolution::Dates.Period,
                       initial_time::Dates.DateTime,
                       percentiles::Vector{Float64},
                       time_steps::Int,
                      )
    data = TimeSeries.TimeArray(
        initial_time : Dates.Hour(1) : initial_time + resolution * (time_steps-1),
        ones(time_steps, length(percentiles))
    )

    return Probabilistic(label, percentiles, data)
end

"""
    Probabilistic(
                  label::String,
                  percentiles::Vector{Float64},  # percentiles for the probabilistic forecast
                  data::TimeSeries.TimeArray,
                 )
Constructs Probabilistic Forecast after constructing a TimeArray from initial_time and time_steps.
"""
# TODO: do we need this check still?
#function Probabilistic(
#                       label::String,
#                       percentiles::Vector{Float64},  # percentiles for the probabilistic forecast
#                       data::TimeSeries.TimeArray,
#                      )
#    if !(length(TimeSeries.colnames(data)) == length(percentiles))
#        throw(DataFormatError(
#            "The size of the provided percentiles and data columns is inconsistent"))
#    end
#    initial_time = TimeSeries.timestamp(data)[1]
#    resolution = get_resolution(data)
#
#    return Probabilistic(label, percentiles, data)
#end

function Probabilistic(label::String,
                       resolution::Dates.Period,
                       initial_time::Dates.DateTime,
                       percentiles::Vector{Float64},  # percentiles for the probabilistic forecast
                       data::TimeSeries.TimeArray)
    return Probabilistic(label, percentiles, data)
end

function make_public_forecast(forecast::ProbabilisticInternal, data::TimeSeries.TimeArray)
    return Probabilistic(get_label(forecast), get_percentiles(forecast), data)
end

function make_internal_forecast(forecast::Probabilistic, ts_data::TimeSeriesData)
    return ProbabilisticInternal(get_label(forecast),
                                 get_resolution(forecast),
                                 get_initial_time(forecast),
                                 get_percentiles(forecast),
                                 get_uuid(ts_data),
                                 get_horizon(forecast),
                                )
end

function ScenarioBased(label::String, data::TimeSeries.TimeArray)
    initial_time = TimeSeries.timestamp(data)[1]
    resolution = get_resolution(data)
    scenario_count = length(TimeSeries.colnames(data))
    return ScenarioBased(label, scenario_count, data)
end

"""
Constructs ScenarioBased Forecast after constructing a TimeArray from initial_time and
time_steps.
"""
function ScenarioBased(label::String,
                       resolution::Dates.Period,
                       initial_time::Dates.DateTime,
                       scenario_count::Int,
                       time_steps::Int)
    data = TimeSeries.TimeArray(
        initial_time : Dates.Hour(1) : initial_time + resolution * (time_steps-1),
        ones(time_steps, scenario_count)
    )

    return ScenarioBased(label, data)
end

function make_public_forecast(forecast::ScenarioBasedInternal, data::TimeSeries.TimeArray)
    return ScenarioBased(get_label(forecast), data)
end

function make_internal_forecast(forecast::ScenarioBased, ts_data::TimeSeriesData)
    return ScenarioBasedInternal(get_label(forecast),
                                 get_resolution(forecast),
                                 get_initial_time(forecast),
                                 get_scenario_count(forecast),
                                 get_uuid(ts_data),
                                 get_horizon(forecast),
                                )
end

function forecast_external_to_internal(::Type{T}) where T <: Forecast
    if T <: Deterministic
        forecast_type = DeterministicInternal
    elseif T <: Probabilistic
        forecast_type = ProbabilisticInternal
    elseif T <: ScenarioBased
        forecast_type = ScenarioBasedInternal
    else
        @assert false
    end

    return forecast_type
end

function forecast_internal_to_external(::Type{T}) where T <: ForecastInternal
    if T <: DeterministicInternal
        forecast_type = Deterministic
    elseif T <: ProbabilisticInternal
        forecast_type = Probabilistic
    elseif T <: ScenarioBasedInternal
        forecast_type = ScenarioBased
    else
        @assert false
    end

    return forecast_type
end


text:
test1a=JuMP.Model()
ux(i)=[100;100][i]
lx(i)=[-100;-100][i]
JuMP.@variable(test1a,x[i=1:2], upper_bound=ux(i), lower_bound=lx(i),start=1.0)
JuMP.@NLobjective(test1a,Min,x[1]-x[2])
JuMP.@constraint(test1a,1-x[2]>=0)

testa = MathOptNLPModel(test1a) #MathProgNLPModel(test1a)
mpcca = MPCCAmpl(testa, [0], 0)

@test mpcca.meta.ncc == 0
@test mpcca.meta.nvar == 2
@test mpcca.cvar == [0]
@test mpcca.meta.lvar == [-100.;-100.]
@test mpcca.meta.uvar == [ 100.; 100.]
@test mpcca.meta.ncon == 1

@test obj(mpcca.mp, mpcca.meta.x0)  == 0.0
@test grad(mpcca.mp, mpcca.meta.x0) == [1.0, -1.0]
@test cons(mpcca.mp, mpcca.meta.x0) == [-1.0]
@test jac(mpcca.mp, mpcca.meta.x0)  == [0.0 -1.0]

@test hessG(mpcca, mpcca.meta.x0, zeros(0)) == zeros(2,2)
@test hGprod(mpcca, mpcca.meta.x0, zeros(0), ones(2)) == zeros(2)
@test issparse(jacG(mpcca, mpcca.meta.x0))
@test consG(mpcca, mpcca.meta.x0) == zeros(0)

@test hessH(mpcca, mpcca.meta.x0, zeros(0)) == zeros(2,2)
@test hHprod(mpcca, mpcca.meta.x0, zeros(0), ones(2)) == zeros(2)
@test sum(jacH(mpcca, mpcca.meta.x0)) == 0
@test consH(mpcca, mpcca.meta.x0) == zeros(0)

@test jprod(mpcca, mpcca.meta.x0, zeros(2)) == zeros(1)
@test jtprod(mpcca, mpcca.meta.x0, zeros(1)) == zeros(2)

@test size(hess(mpcca, mpcca.meta.x0)) == (2,2)
@test hess(mpcca, mpcca.meta.x0, zeros(1), obj_weight = 0.0) == zeros(2,2)
@test hprod(mpcca, mpcca.meta.x0, zeros(2)) == zeros(2)
@test hprod(mpcca, mpcca.meta.x0, zeros(1), ones(2), obj_weight = 0.0) == zeros(2)

@test hessH_structure(mpcca) == hessG_structure(mpcca)
@test length(hessH_coord(mpcca, mpcca.meta.x0, zeros(0))) == mpcca.meta.nvar * (mpcca.meta.nvar + 1) / 2
@test length(hessG_coord(mpcca, mpcca.meta.x0, zeros(0))) == mpcca.meta.nvar * (mpcca.meta.nvar + 1) / 2


text:
using BHAtp, Statistics

ProjDir = @__DIR__
!isdir(joinpath(ProjDir, "plots")) && mkdir(joinpath(ProjDir, "plots"))

ProjName = split(ProjDir, "/")[end]
bhaj = BHAtp.BHAJ(ProjName, ProjDir)
bhaj.ratio = 2.5

segs = [
	# Element type,  Material,    Length,     ID,   OD,    fc
	(:bit,           :steel,        0.00,   2.000, 12.250,  0.0),
  (:collar,        :steel,        1.00,   2.000,  7.875,  0.0),
  (:collar,        :steel,        1.85,   2.750,  7.875,  0.0),
	(:stabilizer,    :steel,        0.00,   2.750, 12.250,  0.0),
  (:collar,        :steel,        3.90,   2.750,  7.875,  0.0),
	(:collar,        :monel,       28.51,   2.750,  8.000,  0.0),
  (:collar,        :monel,        1.77,   3.000,  8.000,  0.0),
  (:stabilizer,    :monel,        0.00,   3.000, 12.250,  0.0),
  (:collar,        :monel,        3.78,   3.000,  8.000,  0.0),
	(:collar,        :monel,        2.50,   3.000,  8.875,  0.0),
  (:collar,        :monel,       25.00,   5.500,  8.875,  0.0),
  (:stabilizer,    :monel,        0.00,   5.500, 12.250,  0.0),
  (:collar,        :monel,       14.23,   5.500,  8.875,  0.0),
  (:collar,        :monel,        2.46,   3.000,  8.875,  0.0),
  (:collar,        :monel,       29.68,   2.750,  8.313,  0.0),
  (:collar,        :monel,        1.45,   3.000,  8.000,  0.0),
  (:stabilizer,    :monel,        0.00,   3.000, 12.250,  0.0),
  (:collar,        :monel,        3.75,   3.000,  8.000,  0.0),
  (:collar,        :monel,       29.62,   2.750,  8.313,  0.0),
  (:collar,        :steel,      218.39,   2.750,  8.000,  0.0),
  (:collar,        :steel,       31.60,   2.750,  8.000,  0.0)
]

traj = [
	#Heading,  Diameter
	( 60.0,      12.250)
]

incls = 45:5:55
wobs = 20:20:60

@time bhaj(segs, traj, wobs, incls)

display("Fetch tp=false, wob=40, incl@bit=50 solution:")
df,df_tp = show_solution(ProjDir, 40, 50, show=false, tp=false)
df[1:25,[2; 5:6; 9:12]] |> display
println()

println("Fetch tp=true, wob=40, incl@bit=50 solution:")
df,df_tp = show_solution(ProjDir, 40, 50, show=false);
df[1:25,[2; 5:6; 9:12]] |> display
println()

tpf = create_final_tp_df(ProjDir, wobs, incls; ofu=false)
tpf |> display
println()

# Collect computed side forces on NBS @ 45ft

sfs = []
for wob in wobs
  for incl in incls
    df, df_tp = show_solution(ProjDir, wob, incl, show=false);
    sf = @linq df |> where(:dtb .== 2.85)
    append!(sfs, sf.fx)
  end
end

# Actual tps

actual = DataFrame(
  incls = [49.75, 50.75, 51.50, 52.80, 54, 55],
  atp = [1.05, 1.07, 1.20, 1.15, 1.60, 1.22]
)

name = "figure_07_1"

figure_07_1 = @pgf GroupPlot({
    group_style = {
      group_size = "1 by 1",
      vertical_sep = "0pt",
      xticklabels_at = "edge bottom"
    }
  },
  Axis({
      height = "6cm", width = " 8cm",
      xmin = 42, xmax = 56,
      ymin = 0.0, ymax = 2.5,
      grid = "major",
      xlabel = "Inclination [°]",
      ylabel = "TP [°/100ft]",
      title = "Figure 7:Actual vs. Isotropic TP",
      legend_pos = "south west",
    },
    Plot({ color="blue", mark = "." },
      Table(tpf.incl[1:3], tpf.tp[1:3])),
    LegendEntry("TP wob=20"),
    Plot({ color="gray", mark = "." },
      Table(tpf.incl[4:6], tpf.tp[4:6])),
    LegendEntry("TP wob=40"),
    Plot({ color="green", mark = "." },
      Table(tpf.incl[7:9], tpf.tp[7:9])),
    LegendEntry("TP wob=60"),
    Plot({ thick, color="black", mark = "+", },
      Table(actual.incls, actual.atp)),
    LegendEntry("TP, actual")))

savefig(joinpath(ProjDir, "plots", name), figure_07_1, ProjDir)


text:
# # Practical error bounds for the forces
#
# This is a simple example showing how to compute error estimates for the forces
# on a ``{\rm TiO}_2`` molecule, from which we can either compute asymptotically
# valid error bounds or increase the precision on the computation of the forces.
#
# The strategy we follow is described with more details in [^CDKL2021] and we
# will use in comments the density matrices framework. We will also needs
# operators and functions from
# [`src/scf/newton.jl`](https://dftk.org/blob/master/src/scf/newton.jl).
#
# [^CDKL2021]:
#     E. Cancès, G. Dusson, G. Kemlin, and A. Levitt
#     *Practical error bounds for properties in plane-wave electronic structure
#     calculations* Preprint, 2021. [arXiv](https://arxiv.org/abs/2111.01470)
using DFTK
using Printf
using LinearAlgebra
using ForwardDiff
using LinearMaps
using IterativeSolvers

# ## Setup
# We setup manually the ``{\rm TiO}_2`` configuration from
# [Materials Project](https://materialsproject.org/materials/mp-2657/).
Ti = ElementPsp(:Ti, psp=load_psp("hgh/lda/ti-q4.hgh"))
O  = ElementPsp(:O, psp=load_psp("hgh/lda/o-q6.hgh"))
atoms     = [Ti, Ti, O, O, O, O]
positions = [[0.5,     0.5,     0.5],  # Ti
             [0.0,     0.0,     0.0],  # Ti
             [0.19542, 0.80458, 0.5],  # O
             [0.80458, 0.19542, 0.5],  # O
             [0.30458, 0.30458, 0.0],  # O
             [0.69542, 0.69542, 0.0]]  # O
lattice   = [[8.79341  0.0      0.0];
             [0.0      8.79341  0.0];
             [0.0      0.0      5.61098]];

# We apply a small displacement to one of the ``\rm Ti`` atoms to get nonzero
# forces.
positions[1] .+= [-0.022, 0.028, 0.035]

# We build a model with one k-point only, not too high `Ecut_ref` and small
# tolerance to limit computational time. These parameters can be increased for
# more precise results.
model = model_LDA(lattice, atoms, positions)
kgrid = [1, 1, 1]
Ecut_ref = 35
basis_ref = PlaneWaveBasis(model; Ecut=Ecut_ref, kgrid)
tol = 1e-8;

# ## Computations

# We compute the reference solution ``P_*`` from which we will compute the
# references forces.
scfres_ref = self_consistent_field(basis_ref; tol, callback=info->nothing)
ψ_ref, _ = DFTK.select_occupied_orbitals(basis_ref, scfres_ref.ψ,
                                         scfres_ref.occupation);

# We compute a variational approximation of the reference solution with
# smaller `Ecut`. `ψr`, `ρr` and `Er` are the quantities computed with `Ecut`
# and then extended to the reference grid.
#
# !!! note "Choice of convergence parameters"
#     Be careful to choose `Ecut` not too close to `Ecut_ref`.
#     Note also that the current choice `Ecut_ref = 35` is such that the
#     reference solution is not converged and `Ecut = 15` is such that the
#     asymptotic regime (crucial to validate the approach) is barely established.
Ecut = 15
basis = PlaneWaveBasis(model; Ecut=Ecut, kgrid)
scfres = self_consistent_field(basis; tol, callback=info->nothing)
ψr = DFTK.transfer_blochwave(scfres.ψ, basis, basis_ref)
ρr = compute_density(basis_ref, ψr, scfres.occupation)
Er, hamr = energy_hamiltonian(basis_ref, ψr, scfres.occupation; ρ=ρr);

# We then compute several quantities that we need to evaluate the error bounds.

# - Compute the residual ``R(P)``, and remove the virtual orbitals, as required
#   in [`src/scf/newton.jl`](https://github.com/JuliaMolSim/DFTK.jl/blob/fedc720dab2d194b30d468501acd0f04bd4dd3d6/src/scf/newton.jl#L121).
res = DFTK.compute_projected_gradient(basis_ref, ψr, scfres.occupation)
res, occ = DFTK.select_occupied_orbitals(basis_ref, res, scfres.occupation)
ψr, _ = DFTK.select_occupied_orbitals(basis_ref, ψr, scfres.occupation);

# - Compute the error ``P-P_*`` on the associated orbitals ``ϕ-ψ`` after aligning
#   them: this is done by solving ``\min |ϕ - ψU|`` for ``U`` unitary matrix of
#   size ``N\times N`` (``N`` being the number of electrons) whose solution is
#   ``U = S(S^*S)^{-1/2}`` where ``S`` is the overlap matrix ``ψ^*ϕ``.
function compute_error(basis, ϕ, ψ)
    map(zip(ϕ, ψ)) do (ϕk, ψk)
        S = ψk'ϕk
        U = S*(S'S)^(-1/2)
        ϕk - ψk*U
    end
end
err = compute_error(basis_ref, ψr, ψ_ref);

# - Compute ``{\bm M}^{-1}R(P)`` with ``{\bm M}^{-1}`` defined in [^CDKL2021]:
P = [PreconditionerTPA(basis_ref, kpt) for kpt in basis_ref.kpoints]
map(zip(P, ψr)) do (Pk, ψk)
    DFTK.precondprep!(Pk, ψk)
end
function apply_M(φk, Pk, δφnk, n)
    DFTK.proj_tangent_kpt!(δφnk, φk)
    δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk
    DFTK.proj_tangent_kpt!(δφnk, φk)
    δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk
    DFTK.proj_tangent_kpt!(δφnk, φk)
end
function apply_inv_M(φk, Pk, δφnk, n)
    DFTK.proj_tangent_kpt!(δφnk, φk)
    op(x) = apply_M(φk, Pk, x, n)
    function f_ldiv!(x, y)
        x .= DFTK.proj_tangent_kpt(y, φk)
        x ./= (Pk.mean_kin[n] .+ Pk.kin)
        DFTK.proj_tangent_kpt!(x, φk)
    end
    J = LinearMap{eltype(φk)}(op, size(δφnk, 1))
    δφnk = cg(J, δφnk, Pl=DFTK.FunctionPreconditioner(f_ldiv!),
              verbose=false, reltol=0, abstol=1e-15)
    DFTK.proj_tangent_kpt!(δφnk, φk)
end
function apply_metric(φ, P, δφ, A::Function)
    map(enumerate(δφ)) do (ik, δφk)
        Aδφk = similar(δφk)
        φk = φ[ik]
        for n = 1:size(δφk,2)
            Aδφk[:,n] = A(φk, P[ik], δφk[:,n], n)
        end
        Aδφk
    end
end
Mres = apply_metric(ψr, P, res, apply_inv_M);

# We can now compute the modified residual ``R_{\rm Schur}(P)`` using a Schur
# complement to approximate the error on low-frequencies[^CDKL2021]:
#
# ```math
# \begin{bmatrix}
# (\bm \Omega + \bm K)_{11} & (\bm \Omega + \bm K)_{12} \\
# 0 & {\bm M}_{22}
# \end{bmatrix}
# \begin{bmatrix}
# P_{1} - P_{*1} \\ P_{2}-P_{*2}
# \end{bmatrix} =
# \begin{bmatrix}
# R_{1} \\ R_{2}
# \end{bmatrix}.
# ```

# - Compute the projection of the residual onto the high and low frequencies:
resLF = DFTK.transfer_blochwave(res, basis_ref, basis)
resHF = res - DFTK.transfer_blochwave(resLF, basis, basis_ref);

# - Compute ``{\bm M}^{-1}_{22}R_2(P)``:
e2 = apply_metric(ψr, P, resHF, apply_inv_M);

# - Compute the right hand side of the Schur system:
## Rayleigh coefficients needed for `apply_Ω`
Λ = map(enumerate(ψr)) do (ik, ψk)
    Hk = hamr.blocks[ik]
    Hψk = Hk * ψk
    ψk'Hψk
end
ΩpKe2 = DFTK.apply_Ω(e2, ψr, hamr, Λ) .+ DFTK.apply_K(basis_ref, e2, ψr, ρr, occ)
ΩpKe2 = DFTK.transfer_blochwave(ΩpKe2, basis_ref, basis)
rhs = resLF - ΩpKe2;

# - Solve the Schur system to compute ``R_{\rm Schur}(P)``: this is the most
#   costly step, but inverting ``\bm{\Omega} + \bm{K}`` on the small space has
#   the same cost than the full SCF cycle on the small grid.
ψ, _ = DFTK.select_occupied_orbitals(basis, scfres.ψ, scfres.occupation)
e1 = DFTK.solve_ΩplusK(basis, ψ, rhs, occ; tol).δψ
e1 = DFTK.transfer_blochwave(e1, basis, basis_ref)
res_schur = e1 + Mres;

# ## Error estimates

# We start with different estimations of the forces:
# - Force from the reference solution
f_ref = compute_forces(scfres_ref)
forces   = Dict("F(P_*)" => f_ref)
relerror = Dict("F(P_*)" => 0.0)
compute_relerror(f) = norm(f - f_ref) / norm(f_ref)

# - Force from the variational solution and relative error without
#   any post-processing:
f = compute_forces(scfres)
forces["F(P)"]   = f
relerror["F(P)"] = compute_relerror(f)

# We then try to improve ``F(P)`` using the first order linearization:
#
# ```math
# F(P) = F(P_*) + {\rm d}F(P)\cdot(P-P_*).
# ```

# To this end, we use the `ForwardDiff.jl` package to compute ``{\rm d}F(P)``
# using automatic differentiation.
function df(basis, occupation, ψ, δψ, ρ)
    δρ = DFTK.compute_δρ(basis, ψ, δψ, occupation)
    ForwardDiff.derivative(ε -> compute_forces(basis, ψ.+ε.*δψ, occupation; ρ=ρ+ε.*δρ), 0)
end;

# - Computation of the forces by a linearization argument if we have access to
#   the actual error ``P-P_*``. Usually this is of course not the case, but this
#   is the "best" improvement we can hope for with a linearisation, so we are
#   aiming for this precision.
df_err = df(basis_ref, occ, ψr, DFTK.proj_tangent(err, ψr), ρr)
forces["F(P) - df(P)⋅(P-P_*)"]   = f - df_err
relerror["F(P) - df(P)⋅(P-P_*)"] = compute_relerror(f - df_err)

# - Computation of the forces by a linearization argument when replacing the
#   error ``P-P_*`` by the modified residual ``R_{\rm Schur}(P)``. The latter
#   quantity is computable in practice.
df_schur = df(basis_ref, occ, ψr, res_schur, ρr)
forces["F(P) - df(P)⋅Rschur(P)"]   = f - df_schur
relerror["F(P) - df(P)⋅Rschur(P)"] = compute_relerror(f - df_schur)

# Summary of all forces on the first atom (Ti)
for (key, value) in pairs(forces)
    @printf("%30s = [%7.5f, %7.5f, %7.5f]   (rel. error: %7.5f)\n",
            key, (value[1])..., relerror[key])
end

# Notice how close the computable expression ``F(P) - {\rm d}F(P)⋅R_{\rm Schur}(P)``
# is to the best linearization ansatz ``F(P) - {\rm d}F(P)⋅(P-P_*)``.


text:
using Geodesy, SeisIO
export rms, dist, normalize, normalize!

"""
    rms(A::AbstractArray, B::AbstractArray)

Compute the root mean squared error between two arrays
"""
function rms(A::AbstractArray, B::AbstractArray)
    return sqrt(sum((B.-A).^2)/length(A))
end

"""
    dist(loc1::GeoLoc, loc2::GeoLoc)

Compute the distance (m) between two points (lat, lon, elev)
This assumes the distance is not large, such that the curvature of the earth is negligible.
"""
function dist(loc1::GeoLoc, loc2::GeoLoc)
    loc1_lla = LLA(loc1.lat, loc1.lon, loc1.el)
    loc2_lla = LLA(loc2.lat, loc2.lon, loc2.el)

    return distance(loc1_lla, loc2_lla)
end

"""
    window_read(len::Int64, max_len::Int64; startInd::Int64=2)

Define subarrays of a given length for data reading.
"""
function window_read(len::Int64, max_len::Int64)
    fullArr = collect(1:len)
    # overflow
    rem = length(fullArr) % max_len
    # number of full sized subarrays
    nWin = div(length(fullArr), max_len)

    # define subarrays
    windows = [fullArr[1+(i*max_len):(i+1)*max_len] for i=0:nWin-1]
    if rem != 0 push!(windows, collect((max_len*nWin)+1:len)) end

    return windows
end

"""
    normalize!(x::AbstractArray; mag::Float64=1.0)

Normalize an array and scale to a given magnitude.
"""
function normalize!(x::AbstractArray; mag::Float64=1.0)
    x ./= maximum(abs.(x)) .* mag
    return nothing
end
normalize(x::AbstractArray; mag::Float64=1.0) = (U = deepcopy(x); normalize!(U, mag=mag); return U)

"""
    Base.copy(x::T)

Universal copy function used to copy composite types such as structs.
"""
Base.copy(x::T) where T = T([deepcopy(getfield(x, k)) for k ∈ fieldnames(T)]...)

max_stars_repo_name: kura-okubo/SeisXcorrelation

max_issues_repo_name: kura-okubo/SeisXcorrelation

max_forks_repo_name: kura-okubo/SeisXcorrelation


text:
@testset "$(rpad("Stochastic equations",80))" begin

    ################################################################################
    # Test SDE: Stochastic Differential Equation
    ################################################################################

    function v_sde(λ, t, q, v)
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function u_sde(μ, t, q, u)
        u[1] = μ*q[1]
        u[2] = μ*q[2]
    end

    λ  = 2.
    μ  = 1.

    v_sde_params = (t, q, v) -> v_sde(λ, t, q, v)
    u_sde_params = (t, q, v) -> u_sde(μ, t, q, v)

    sde  = SDE{eltype(x₀), typeof(t₀), typeof(v_sde_params), typeof(u_sde_params), 1}(2, 1, 1, 1, v_sde_params, u_sde_params, t₀, x₀)
    sde1 = SDE(1, 1, v_sde_params, u_sde_params, t₀, x₀)
    sde2 = SDE(1, 1, v_sde_params, u_sde_params, x₀)

    @test sde == sde1
    @test sde == sde2

end


text:
function rref(mat::Matrix)
    # make a copy of the original matrix
    ret = float(copy(mat))
    # # ret = copy(mat)
    # initialize an array to memorize the pivots
    pivot = []
    # it's assumed that the matrix has $m$ rows and $n$ columns
    m, n = size(ret)
    # $i$ and $j$ are used to represent iterators of row and column,
    # respectively
    i = j = 1
    # part of `forward elimination'
    while i <= m && j <= n
        # find the entry with largest absolute value
        imax, vmax = 0.0, 0.0
        for itr1 in i : m
            if vmax < abs(ret[itr1, j])
                imax, vmax = itr1, abs(ret[itr1, j])
            end
        end
        # but this one seems to work pretty fine?
        (vmax > 0) ? (k = imax) : (j += 1; continue)
        # # if vmax > 0
        # #     k = imax
        # # else
        # #     j += 1
        # #     continue
        # # end
        # add the pivot position to the array
        push!(pivot, (i, j))
        # exchange the rows
        for itr1 in i : n
            ret[i, itr1], ret[k, itr1] = ret[k, itr1], ret[i, itr1]
        end
        # `forward elimination'
        for itr1 in (i + 1) : m
            scaler = ret[itr1, j] / ret[i, j]
            for itr2 in (j + 1) : n
                ret[itr1, itr2] -= ret[i, itr2] * scaler
            end
            # manually set this entry to 0, in case of using
            # LU factorization
            ret[itr1, j] = 0
        end
        i += 1
        j += 1
    end
    # part of `back subsitution'
    for itr1 in length(pivot) : -1 : 1
        # get the pivot
        i, j = pivot[itr1]
        # `back subsitution`
        for itr2 in 1 : (i - 1)
            scaler = ret[itr2, j] / ret[i, j]
            for itr3 in j : n
                ret[itr2, itr3] -= ret[i, itr3] * scaler
            end
        end
        # make entry on the pivot be 1
        # and a little trick in order not to create a new variable
        for itr2 in n : -1 : j
            ret[i, itr2] /= ret[i, j]
        end
    end
    return ret
end

max_stars_repo_name: Shimushushushu/Etude-Rhapsody

max_issues_repo_name: Shimushushushu/Etude-Rhapsody

max_forks_repo_name: Shimushushushu/Etude-Rhapsody


text:
#!/usr/bin/env julia

module GetRate

using CSV, DataFrames, Serialization, HDF5, YAML
using ArgParse, Printf

using CarrierCapture

# read arguments and input file
s = ArgParseSettings()
@add_arg_table s begin
    "--input", "-i"
        help = "input file in a yaml format"
        default = "input.yaml"
        arg_type = String
    "--wave", "-w"
        help = "wave function file serialized"
        default = "wave.jld"
        arg_type = String
    "--plot", "-p"
        help = "Plot potentials"
        action = :store_true
    "--verbose", "-v"
        help = "write verbose capture coefficient"
        action = :store_true
end

args = parse_args(ARGS, s)
input = YAML.load(open(args["input"]))
wave_path = args["wave"]
is_verbose = args["verbose"]

println(raw"
      _____                _            _____            _
     / ____|              (_)          / ____|          | |
    | |     __ _ _ __ _ __ _  ___ _ __| |     __ _ _ __ | |_ _   _ _ __ ___
    | |    / _` | '__| '__| |/ _ \ '__| |    / _` | '_ \| __| | | | '__/ _ \
    | |___| (_| | |  | |  | |  __/ |  | |___| (_| | |_) | |_| |_| | | |  __/
     \_____\__,_|_|  |_|  |_|\___|_|   \_____\__,_| .__/ \__|\__,_|_|  \___|
                                                  | |
                                                  |_| v 0.1
      ____      _   ____       _
     / ___| ___| |_|  _ \ __ _| |_ ___
    | |  _ / _ \ __| |_) / _` | __/ _ \
    | |_| |  __/ |_|  _ < (_| | ||  __/
     \____|\___|\__|_| \_\__,_|\__\___|
")

# capture variables
input_capt = input["captures"]
V = input_capt["Volume"]
Tmin, Tmax, NT = input_capt["Tmin"], input_capt["Tmax"], input_capt["NT"]
cut_off = input_capt["cut_off"]
σ = input_capt["σ"]
temperature = range(Tmin, stop=Tmax, length=NT)

println("==============================")
println("Read wave functions from <$(wave_path)>")
file = open(wave_path, "r"); pots = deserialize(file); close(file)
println(keys(pots))
println("==============================")

# TODO: raise no eigenvalue error!
#       after dryrun of GetPotential
# calculate capture coefficient
ccs = Array{conf_coord,1}()
for cc in input_capt["ccs"]
    cc_cfg = cc["cc"]
    pot_i = pots[cc_cfg["initial"]]
    pot_f = pots[cc_cfg["final"]]
    cc = cc_from_dict(pot_i, pot_f, cc_cfg)
    calc_overlap!(cc, cut_off = cut_off, σ = σ)
    calc_capt_coeff!(cc, V, temperature)
    append!(ccs, [cc])
end

# plot
if args["plot"] == true
    println("====Ploting capture rate====\n")
    plot_cfg = get(input, "plot", Dict())
    plot_ccs(ccs, plot_cfg)
end

# write capture coefficient into cvs files
for (i, cc) in enumerate(ccs)
    CSV.write("capt_coeff_$(i).csv", DataFrame([cc.temperature, cc.capt_coeff], [:T, :C]))
end

# write verbose ouput
if is_verbose
    for (i, cc) in enumerate(ccs)
        h5open("overlap_$(i).hdf5", "w") do file
            write(file, "overlap", collect(cc.overlap_matrix))
            write(file, "delta_Eij", collect(cc.δ_matrix))
        end
        h5open("partial_capt_coeff_$(i).hdf5","w") do file
            write(file, "partial_capt_coeff", collect(cc.partial_capt_coeff))
        end
    end
end

end # module


text:
# Unstable Bickley jet

using Printf

ENV["GKSwstype"] = "nul"
using Plots
using Revise
using CUDA
using ClimateMachine

ClimateMachine.init()

using ClimateMachine.Ocean
using ClimateMachine.Ocean.Domains
using ClimateMachine.Ocean.Fields

using ClimateMachine.Mesh.Grids: DiscontinuousSpectralElementGrid
using ClimateMachine.GenericCallbacks: EveryXSimulationTime
using ClimateMachine.GenericCallbacks: EveryXSimulationSteps
using ClimateMachine.Ocean: current_step, Δt, current_time
using ClimateMachine.Ocean: JLD2Writer, OutputTimeSeries, write!
using CLIMAParameters: AbstractEarthParameterSet, Planet

struct NonDimensionalParameters <: AbstractEarthParameterSet end
Planet.grav(::NonDimensionalParameters) = 10
c = sqrt(Planet.grav(NonDimensionalParameters())) # gravity wave speed for unit depth

include("Bickley.jl")
using .Bickley

# Low-p assumption:
effective_node_spacing(Ne, Np, Lx=4π) = Lx / (Ne * (Np + 1))

function run(;
             Ne = 4,
             Np = 4,
             ν = 0,
             time_step = 0.1 * effective_node_spacing(Ne, Np) / c,
             array_type = Array,
             output_time_interval = 2,
             stabilizing_dissipation = nothing,
             stop_time = 200)

    name = @sprintf("climate_machine_unstable_bickley_jet_Ne%d_Np%d_ν%.1e_no_rotation", Ne, Np, ν)

    ClimateMachine.Settings.array_type = array_type

    # Domain

    domain = RectangularDomain(Ne = (Ne, Ne, 1), Np = Np,
                               x = (-2π, 2π), y = (-2π, 2π), z = (0, 1),
                               periodicity = (true, true, false))

    # Physical parameters:
    g = Planet.grav(NonDimensionalParameters())

    # Non-dimensional parameters
    ϵ = 0.1 # Perturbation amplitude
    ℓ = 0.5 # Perturbation width
    k = 0.5 # Perturbation wavenumber

    # Initial conditions: Jet/tracer + perturbations
    uᵢ(x, y, z) = Bickley.U(y, domain.L.y) + ϵ * Bickley.ũ(x, y, ℓ, k)
    vᵢ(x, y, z) = ϵ * Bickley.ṽ(x, y, ℓ, k)
    θᵢ(x, y, z) = Bickley.C(y, domain.L.y)

    initial_conditions = InitialConditions(u=uᵢ, v=vᵢ, θ=θᵢ)

    model = Ocean.HydrostaticBoussinesqSuperModel(
        domain = domain,
        time_step = time_step,
        initial_conditions = initial_conditions,
        parameters = NonDimensionalParameters(),
        turbulence_closure = (νʰ = ν, κʰ = ν, νᶻ = ν, κᶻ = ν),
        rusanov_wave_speeds = (cʰ = sqrt(g * domain.L.z), cᶻ = 1e-2),
        stabilizing_dissipation = stabilizing_dissipation,
        coriolis = (f₀ = 0, β = 0),
        buoyancy = (αᵀ = 0,),
        boundary_tags = ((0, 0), (1, 1), (1, 2)),
        boundary_conditions = (OceanBC(Impenetrable(FreeSlip()), Insulating()),
                               OceanBC(Penetrable(FreeSlip()), Insulating()))
    )

    name = @sprintf("climate_machine_unstable_bickley_jet_Ne%d_Np%d_ν%.1e_no_rotation", Ne, Np, ν)

    ClimateMachine.Settings.array_type = array_type

    # Domain

    domain = RectangularDomain(Ne = (Ne, Ne, 1), Np = Np,
                               x = (-2π, 2π), y = (-2π, 2π), z = (0, 1),
                               periodicity = (true, true, false))

    # Physical parameters:
    g = Planet.grav(NonDimensionalParameters())

    # Non-dimensional parameters
    ϵ = 0.1 # Perturbation amplitude
    ℓ = 0.5 # Perturbation width
    k = 0.5 # Perturbation wavenumber

    # Initial conditions: Jet/tracer + perturbations
    uᵢ(x, y, z) = Bickley.U(y, domain.L.y) + ϵ * Bickley.ũ(x, y, ℓ, k)
    vᵢ(x, y, z) = ϵ * Bickley.ṽ(x, y, ℓ, k)
    θᵢ(x, y, z) = Bickley.C(y, domain.L.y)

    initial_conditions = InitialConditions(u=uᵢ, v=vᵢ, θ=θᵢ)

    model = Ocean.HydrostaticBoussinesqSuperModel(
        domain = domain,
        time_step = time_step,
        initial_conditions = initial_conditions,
        parameters = NonDimensionalParameters(),
        turbulence_closure = (νʰ = ν, κʰ = ν, νᶻ = ν, κᶻ = ν),
        rusanov_wave_speeds = (cʰ = sqrt(g * domain.L.z), cᶻ = 1e-2),
        stabilizing_dissipation = stabilizing_dissipation,
        coriolis = (f₀ = 0, β = 0),
        buoyancy = (αᵀ = 0,),
        boundary_tags = ((0, 0), (1, 1), (1, 2)),
        boundary_conditions = (OceanBC(Impenetrable(FreeSlip()), Insulating()),
                               OceanBC(Penetrable(FreeSlip()), Insulating()))
    )    # We prepare a callback that periodically fetches the horizontal velocity and
    # tracer concentration for later animation,

    writer = JLD2Writer(model, filepath = name * ".jld2", overwrite_existing = true)
    cpu_grid = DiscontinuousSpectralElementGrid(domain, array_type=Array)

    start_time = time_ns()

    data_fetcher = EveryXSimulationTime(output_time_interval) do
        write!(writer)

        cpu_data = convert(Array, model.state.realdata)
        u = SpectralElementField(domain, cpu_grid, view(cpu_data, :, 1, :))

        # Print a helpful message
        step = @sprintf("Step: %d", current_step(model))
        time = @sprintf("time: %.2f", current_time(model))
        max_u = @sprintf("max|u|: %.6f", maximum(abs, u))

        elapsed = (time_ns() - start_time) * 1e-9
        wall_time = @sprintf("elapsed wall time: %.2f min", elapsed / 60)

        isnan(maximum(abs, u)) && error("NaNs.")

        @info "$step, $time, $max_u, $wall_time"
    end

    # and then run the simulation.

    model.solver_configuration.timeend = stop_time

    total_steps = ceil(Int, stop_time / time_step)
    @info @sprintf("Running a simulation of the instability of the Bickley jet (Δt=%.2e, steps=%d)", time_step, total_steps)

    try
        result = ClimateMachine.invoke!(model.solver_configuration;
                                        user_callbacks = [data_fetcher])
    catch err
        @warn "Simulation ended prematurely because $(sprint(showerror, err))"
    end

    return name
end

function visualize(name, contours=false)

    filepath = name * ".jld2"

    u_timeseries = OutputTimeSeries(:u, filepath)
    v_timeseries = OutputTimeSeries(:v, filepath)
    η_timeseries = OutputTimeSeries(:η, filepath)
    c_timeseries = OutputTimeSeries(:θ, filepath)

    u₀ = u_timeseries[1]
    domain = u₀.domain
    assembled_u₀ = assemble(u₀)
    x = assembled_u₀.x[:, 1, 1]
    y = assembled_u₀.y[1, :, 1]

    times = u_timeseries.times

    animation = @animate for i = 1:length(u_timeseries)

        @info "Plotting frame $i of $(length(u_timeseries))..."

        kwargs = (xlim = domain.x, ylim = domain.y, linewidth = 0, aspectratio = 1)

        u = assemble(u_timeseries[i]).data[:, :, 1]
        v = assemble(v_timeseries[i]).data[:, :, 1]
        η = assemble(η_timeseries[i]).data[:, :, 1]
        c = assemble(c_timeseries[i]).data[:, :, 1]

        if ~isnan(maximum(abs, u))

            s = @. sqrt(u^2 + v^2)

            ηmin = minimum(η)
            ηmax = maximum(η)
            ηlevels = range(ηmin, ηmax, length=31)

            kwargs = Dict(:xlabel => "x",
                          :ylabel => "y",
                          :aspectratio => 1,
                          :linewidth => 0,
                          :colorbar => true,
                          :xlims => (-domain.L.x/2, domain.L.x/2),
                          :ylims => (-domain.L.y/2, domain.L.y/2))

            plotter = contours ? contourf : heatmap

            kwargs[:clims] = (0, 1)
            contours && (kwargs[:levels] = range(0, 1, length=31))

            s_plot = plotter(x, y, clamp.(s, 0, 1)'; color = :thermal, kwargs...)

            kwargs[:clims] = (-1, 1)
            contours && (kwargs[:levels] = range(-1, 1, length=31))

            u_plot = plotter(x, y, clamp.(u, -1, 1)'; color = :balance, kwargs...)
            c_plot = plotter(x, y, clamp.(c, -1, 1)'; color = :thermal, kwargs...)
            #η_plot = plotter(x, y, clamp.(c, -1, 1)'; color = :thermal, kwargs...)

            u_title = @sprintf("u at t = %.2f", times[i])
            s_title = @sprintf("speed at t = %.2f", times[i])
            c_title = @sprintf("c at t = %.2f", times[i])

            plot(u_plot, s_plot, c_plot,
                 title = [u_title s_title c_title],
                 layout = (1, 3),
                 size = (1600, 400))

        end
    end

    gif(animation, name * ".gif", fps = 8)

    return nothing
end

include("StabilizingDissipations.jl")
using .StabilizingDissipations: StabilizingDissipation

Ne = 8
Np = 3

time_step = 0.1 * effective_node_spacing(Ne, Np) / c

test_dissipation = StabilizingDissipation(minimum_node_spacing = effective_node_spacing(Ne, Np),
                                          time_step = time_step,
                                          Δu = 1e-3,
                                          Δθ = 1e-3)

name = run(Ne=Ne, Np=Np, stabilizing_dissipation=test_dissipation)
visualize(name)

#=
for DOF in (32, 64, 128, 256, 512)
    for Np in (2, 3, 4, 5, 6)
        Ne = round(Int, DOF / (Np+1))
        name = run(Ne=16, Np=3, safety=0.1)
        visualize(name)
    end
end

for DOF in (512,)
    for Np in (2, 3, 4, 5, 6)
        Ne = round(Int, DOF / (Np+1))
        name = run(Ne=16, Np=3, safety=0.1, ν=1e-4)
        visualize(name)
    end
end

for DOF in (1024,)
    for Np in (2, 3, 4, 5, 6)
        Ne = round(Int, DOF / (Np+1))
        name = run(Ne=16, Np=3, safety=0.1, ν=1e-5)
        visualize(name)
    end
end
=#


text:
"""
    CheckSolution

Type that wraps the solution of a property checking problem.

### Fields

- `satisfied` -- is the property satisfied?
- `violation` -- step at which the property is violated (-1 otherwise)
- `options`   -- the dictionary of options

### Notes

This type contains the answer if the property is satisfied, and if not, it
contains the index at which the property might be violated for the first time.
"""
struct CheckSolution <: AbstractSolution
    satisfied::Bool
    violation::Int
    options::AbstractOptions
end

# constructor with no options
CheckSolution(satisfied::Bool, violation::Int) =
    CheckSolution(satisfied, violation, Options())


text:
"""
Base struct for all DiffEqEnvironments.jl observations.
"""
struct ObservationFunction
    f
end

(of::ObservationFunction)(s,a) = of.f(s, a)

"""
General observation of form ``o=of(s,a)``.
Passed function `of` requires a case `of(s, a::Nothing)`
    for ini
"""
CustomObservation(of) = ObservationFunction(of)

"""
General observation of form ``o=of(s)``.
"""
CustomStateObservation(of) = ObservationFunction((s, a) -> of(s))

"""
Full observation of state: ``o=s``.
"""
FullStateObservation() = ObservationFunction((s, a) -> s)

"""
Linear relationship of form ``o=\\mathbf{C}s`` between state and observation.
"""
LinearStateObservation(C) = ObservationFunction((s, a) -> C * s)

function LinearObservation(C, D)
    of(s, a::Nothing) = C * s # used for step 0
    of(s, a) = C * s + D * a
    return ObservationFunction(of)
end


text:
using Base, OrthogonalPolynomialsQuasi, ContinuumArrays, QuasiArrays, FillArrays,
        LazyArrays, BandedMatrices, LinearAlgebra, FastTransforms, IntervalSets,
        InfiniteLinearAlgebra, Test
using ForwardDiff, SemiseparableMatrices, SpecialFunctions, LazyBandedMatrices
import ContinuumArrays: BasisLayout, MappedBasisLayout
import OrthogonalPolynomialsQuasi: jacobimatrix, ∞, ChebyshevInterval, LegendreWeight,
            Clenshaw, forwardrecurrence!, singularities
import LazyArrays: ApplyStyle, colsupport, MemoryLayout, arguments
import SemiseparableMatrices: VcatAlmostBandedLayout
import QuasiArrays: MulQuasiMatrix
import Base: OneTo
import InfiniteLinearAlgebra: KronTrav, Block
import FastTransforms: clenshaw!

@testset "singularities" begin
    x = Inclusion(ChebyshevInterval())
    @test singularities(x) == singularities(exp.(x)) == singularities(x.^2) ==
        singularities(x .+ 1) == singularities(1 .+ x) == singularities(x .+ x) ==
        LegendreWeight()
    @test singularities(exp.(x) .* JacobiWeight(0.1,0.2)) ==
        singularities(JacobiWeight(0.1,0.2) .* exp.(x)) ==
        JacobiWeight(0.1,0.2)

    x = Inclusion(0..1)
    @test singularities(x) == singularities(exp.(x)) == singularities(x.^2) ==
        singularities(x .+ 1) == singularities(1 .+ x) == singularities(x .+ x) ==
        legendreweight(0..1)
end

include("test_chebyshev.jl")
include("test_legendre.jl")
include("test_ultraspherical.jl")
include("test_jacobi.jl")
include("test_fourier.jl")
include("test_odes.jl")
include("test_normalized.jl")
include("test_lanczos.jl")

@testset "Auto-diff" begin
    U = Ultraspherical(1)
    C = Ultraspherical(2)

    f = x -> ChebyshevT{eltype(x)}()[x,5]
    @test ForwardDiff.derivative(f,0.1) ≈ 4*U[0.1,4]
    f = x -> ChebyshevT{eltype(x)}()[x,5][1]
    @test ForwardDiff.gradient(f,[0.1]) ≈ [4*U[0.1,4]]
    @test ForwardDiff.hessian(f,[0.1]) ≈ [8*C[0.1,3]]

    f = x -> ChebyshevT{eltype(x)}()[x,1:5]
    @test ForwardDiff.derivative(f,0.1) ≈ [0;(1:4).*U[0.1,1:4]]
end


text:
module Chemostat_InSilico

    import ProjAssistant
    const PjAss = ProjAssistant
    PjAss.@gen_top_proj

    include("Dynamic/Dynamic.jl")

end


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

using Markdown
using InteractiveUtils

# ╔═╡ 2d73bd62-3739-11eb-35c0-09836bdd3d24
file = open(f->read(f, String), "day_2.txt")

# ╔═╡ 431ed85e-3739-11eb-0f7d-bd408abb654f
begin
	#separing each line
	input = eachmatch(r"(\d+)-(\d+) (\w): (\w+)\n",file)

	#getting first match of regex and converting to int
	input = map(x-> x.captures, input)
end

# ╔═╡ db441a22-3739-11eb-3657-319f8fbfbbe2
[x[1] for x in input]

# ╔═╡ ebaf6628-3739-11eb-3ec3-3120c2937e13
function valid_line(min ,max ,password, letter::AbstractString)
	min = parse(Int64, min)
	max = parse(Int64, max)

	matchs = collect( eachmatch( Regex(letter) , password) )
	return size(matchs)[1] ≥ min && size(matchs)[1] ≤ max
end

# ╔═╡ a9690290-373b-11eb-0312-671f6f134c36
resp1 = sum([valid_line(x[1],x[2],x[4],x[3]) for x in input])

# ╔═╡ ba4b5356-373b-11eb-2dc9-65a12a13c5d2
md"awnser for problem1 $(resp1)"

# ╔═╡ 65c76852-373c-11eb-3b42-8bdfacf8c6b3
md"# Problem 2"

# ╔═╡ 769c461e-373c-11eb-050a-1d496a98c2ab
function valid_line2(pos1 ,pos2 ,password, letter::AbstractString)
	pos1 = parse(Int64, pos1)
	pos2 = parse(Int64, pos2)

	matchs =  collect(eachmatch( Regex(letter) , password))
	if size(matchs)[1] == 0
		return false
	end
	apperance= map(x -> x.offset, matchs)
	return xor((pos1 in apperance) ,(pos2 in apperance))
end

# ╔═╡ b71c9022-373c-11eb-1168-599806156629
resp2 = sum([valid_line2(x[1],x[2],x[4],x[3]) for x in input])

# ╔═╡ Cell order:
# ╠═2d73bd62-3739-11eb-35c0-09836bdd3d24
# ╠═431ed85e-3739-11eb-0f7d-bd408abb654f
# ╠═db441a22-3739-11eb-3657-319f8fbfbbe2
# ╠═ebaf6628-3739-11eb-3ec3-3120c2937e13
# ╠═a9690290-373b-11eb-0312-671f6f134c36
# ╟─ba4b5356-373b-11eb-2dc9-65a12a13c5d2
# ╠═65c76852-373c-11eb-3b42-8bdfacf8c6b3
# ╠═769c461e-373c-11eb-050a-1d496a98c2ab
# ╠═b71c9022-373c-11eb-1168-599806156629


text:
abstract type CurveProblem end
abstract type CurveModifier end
abstract type WhatJump <: CurveModifier end
abstract type WhatDynamics <: CurveModifier end

struct JumpStart{F <: AbstractFloat} <: WhatJump
    jumpsize::F
end
struct ZeroStart <: WhatJump end

function make_spans(c::CurveProblem, span)
    return make_spans(c::CurveProblem, span, isjumped(c))
end

struct MDCDynamics <: WhatDynamics end

dynamics(c::CurveProblem) = dynamics(c::CurveProblem, whatdynamics(c))
build_cond(c::CurveProblem, r, tol) = build_cond(c, r, tol, whatdynamics(c))
dHdu_residual(c::CurveProblem, u, t, p) = dHdu_residual(c, u, t, p, whatdynamics(c))
build_affect(c::CurveProblem, affect) = build_affect(c, affect, whatdynamics(c))

"""
For callbacks to tune MD Curve
"""
abstract type ConditionType end
struct ResidualCondition <: ConditionType end
struct CostCondition <: ConditionType end

abstract type CallbackCallable end
abstract type AdjustmentCallback <: CallbackCallable end

"""
    MomentumReadjustment(tol::AbstractFloat, verbose::Bool)
Ideally, dHdu = 0 throughout curve evolution, where H is the Hamiltonian/momentum, and u is the curve velocity in parameter space. Numerical error integrates and prevents this. This struct readjusts momentum when `abs(dHdu) > tol`, so that `dHdu = 0` is recovered.
"""
struct MomentumReadjustment{T <: AbstractFloat} <: AdjustmentCallback
    tol::T
    verbose::Bool
end

"""
    Terminates curve evolution when the cost exceeds the momentum
"""
struct TerminalCond <: AdjustmentCallback end
MomentumReadjustment(a; verbose=false) = MomentumReadjustment(a, verbose)

"""
Experimental. See documentation for MomentumReadjustment. This acts the  same, but instead of modifying the momentum, it modifies the state of the curve (i.e. current parameters) itself, by doing gradient descent to minimise the cost function, subject to the constraint that the distance from the initial parameters does not decrease.
"""
struct StateReadjustment{T <: AbstractFloat} <: AdjustmentCallback
    tol::T
verbose::Bool
end
StateReadjustment(a; verbose=false) = StateReadjustment(a, verbose)

abstract type AffectType end
struct StateAffect <: AffectType end
struct CostateAffect <: AffectType end


text:
using PowerModels
using .PP2PM
import JuMP

function increase_rate_a(pm)
    for (i, branch) in pm["branch"]
            branch["rate_a"] = 1.01 * branch["rate_a"]
    end
    return pm
end

function run_powermodels(json_path)
    # function to run optimal transmission switching (OTS) optimization from powermodels.jl
    pm = PP2PM.load_pm_from_json(json_path)
    model = PP2PM.get_model(pm["pm_model"])

    solver = PP2PM.get_solver(pm["pm_solver"], pm["pm_nl_solver"], pm["pm_mip_solver"],
    pm["pm_log_level"], pm["pm_time_limit"], pm["pm_nl_time_limit"], pm["pm_mip_time_limit"])

    result = run_ots(pm, model, solver,
                        setting = Dict("output" => Dict("branch_flows" => true)))
    i = 0
    max_iteration = 6
    while i < max_iteration && !(result["termination_status"] == JuMP.MOI.LOCALLY_SOLVED || result["termination_status"] == JuMP.MOI.OPTIMAL)
        print("\n\ntrying to increase rate A by 1%: ", i, "/", max_iteration, "\n\n")
        pm = increase_rate_a(pm)
        result = run_ots(pm, model, solver,
                        setting = Dict("output" => Dict("branch_flows" => true)))
        i = i + 1
        print_summary(result)
    end

    return result
end


text:
using AtomicLevels
using HalfIntegers
using Test

@testset "Intermediate terms" begin
    @testset "LS coupling" begin
        @test !AtomicLevels.istermvalid(T"2S", Seniority(2))

        @test string(IntermediateTerm(T"2D", Seniority(3))) == "₃²D"
        @test string(IntermediateTerm(T"2D", 3)) == "₍₃₎²D"

        for t in [T"2D", T"2Do"], ν in [1, Seniority(2), 3]
            it = IntermediateTerm(t, ν)
            @test length(propertynames(it)) == 3
            @test hasproperty(it, :term)
            @test hasproperty(it, :ν)
            @test hasproperty(it, :nu)
            @test it.term == t
            @test it.ν == it.nu
            @test it.nu == ν
            @test !hasproperty(it, :foo)
            @test_throws ErrorException it.foo
        end

        @testset "Seniority" begin
            # Table taken from p. 25 of
            #
            # - Froese Fischer, C., Brage, T., & Jönsson, P. (1997). Computational
            #   Atomic Structure : An MCHF Approach. Bristol, UK Philadelphia, Penn:
            #   Institute of Physics Publ.
            subshell_terms = [
                o"1s" => (1 => "2S1",
                          2 => "1S0"),
                o"2p" => (1 => "2P1",
                          2 => ("1S0", "1D2", "3P2"),
                          3 => ("2P1", "2D3", "4S3")),
                o"3d" => (1 => "2D1",
                          2 => ("1S0", "1D2", "1G2", "3P2", "3F2"),
                          3 => ("2D1", "2P3", "2D3", "2F3", "2G3", "2H3", "4P3", "4F3"),
                          4 => ("1S0", "1D2", "1G2", "3P2", "3F2", "1S4", "1D4", "1F4", "1G4", "1I4", "3P4", "3D4", "3F4", "3G4", "3H4", "5D4"),
                          5 => ("2D1", "2P3", "2D3", "2F3", "2G3", "2H3", "4P3", "4F3", "2S5", "2D5", "2F5", "2G5", "2I5", "4D5", "4G5", "6S5")),
                o"4f" => (1 => "2F1",
                          2 => ("1S0", "1D2", "1G2", "1I2", "3P2", "3F2", "3H2"))
            ]

            foreach(subshell_terms) do (orb, data)
                foreach(data) do (occ, expected_its)
                    p = parity(orb)^occ
                    expected_its = map(expected_its isa String ? (expected_its,) : expected_its) do ei
                        t_str = "$(ei[1:2])$(isodd(p) ? "o" : "")"
                        IntermediateTerm(parse(Term, t_str), Seniority(parse(Int, ei[end])))
                    end
                    its = intermediate_terms(orb, occ)
                    @test its == [expected_its...]
                end
            end
        end
    end

    @testset "jj coupling" begin
        # Taken from Table A.10 of
        #
        # - Grant, I. P. (2007). Relativistic Quantum Theory of Atoms and
        #   Molecules : Theory and Computation. New York: Springer.
        #
        # except for the last row, which seems to have a typo since J=19/2 is
        # missing. Cf. Table 4-5 of
        #
        # - Cowan, R. (1981). The Theory of Atomic Structure and
        #   Spectra. Berkeley: University of California Press.
        ref_table = [[rc"1s2", rc"2p-2"] => [0 => [0]],
                     [rc"1s", rc"2p-"] => [1 => [1/2]],
                     [rc"2p4", rc"3d-4"] => [0 => [0]],
                     [rc"2p", rc"2p3", rc"3d-", rc"3d-3"] => [1 => [3/2]],
                     [rc"2p2", rc"3d-2"] => [0 => [0], 2 => [2]],
                     [rc"3d6", rc"4f-6"] => [0 => [0]],
                     [rc"3d", rc"3d5", rc"4f-", rc"4f-5"] => [1 => [5/2]],
                     [rc"3d2", rc"3d4", rc"4f-2", rc"4f-4"] => [0 => [0], 2 => [2,4]],
                     [rc"3d3", rc"4f-3"] => [1 => [5/2], 3 => [3/2, 9/2]],
                     [rc"4f8", rc"5g-8"] => [0 => [0]],
                     [rc"4f", rc"4f7", rc"5g-", rc"5g-7"] => [1 => [7/2]],
                     [rc"4f2", rc"4f6", rc"5g-2", rc"5g-6"] => [0 => [0], 2 => [2,4,6]],
                     [rc"4f3", rc"4f5", rc"5g-3", rc"5g-5"] => [1 => [7/2], 3 => [3/2, 5/2, 9/2, 11/2, 15/2]],
                     [rc"4f4", rc"5g-4"] => [0 => [0], 2 => [2,4,6], 4 => [2,4,5,8]],
                     [rc"5g10", rc"6h-10"] => [0 => [0]],
                     [rc"5g", rc"5g9", rc"6h-", rc"6h-9"] => [1 => [9/2]],
                     [rc"5g2", rc"5g8", rc"6h-2", rc"6h-8"] => [0 => [0], 2 => [2,4,6,8]],
                     [rc"5g3", rc"5g7", rc"6h-3", rc"6h-7"] => [1 => [9/2],
                                                                3 => vcat((3:2:17), 21)./2],
                     [rc"5g4", rc"5g6", rc"6h-4", rc"6h-6"] => [0 => 0, 2 => [2,4,6,8],
                                                                4 => [0,2,3,4,4,5,6,6,7,8,9,10,12]],
                     [rc"5g5", rc"6h-5"] => [1 => [9/2], 3 => vcat((3:2:17), 21)./2,
                                             5 => vcat(1, (5:2:19), 25)./2]]

        for (cfgs,cases) in ref_table
            ref = [reduce(vcat, [[IntermediateTerm(convert(HalfInt, J), Seniority(ν))
                                  for J in Js]
                                 for (ν,Js) in cases])]
            for cfg in cfgs
                @test intermediate_terms(cfg) == ref
            end
        end
    end
end


text:
function add_kwargs!(ex::Expr; kwargs...)
    ex.head === :call || throw(ArgumentError("expression is not a function call"))
    isempty(ex.args) && throw(ArgumentError("expression body is empty"))
    if !isempty(kwargs)
        params = Expr(:parameters)
        for (name, value) in kwargs
            push!(params.args, Expr(:kw, name, value))
        end
        # Parameters need to come after the function name and before positional arguments
        if length(ex.args) == 1
            push!(ex.args, params)
        else
            insert!(ex.args, 2, params)
        end
    end
    ex
end

function get_union_call(foo::Symbol, type_tuple::Expr; kwargs...)
    # Get types from tuple and create a collection of symbols for use in the call.
    types = get_types(get_body(type_tuple))
    arg_names = [Symbol("x$j") for j in 1:length(types)]

    # Generate the call.
    typed_args = map((name, typ) -> :($name::$(unionise_type(typ))), arg_names, types)
    call = add_kwargs!(Expr(:call, foo, typed_args...); kwargs...)

    return replace_body(type_tuple, call), arg_names
end

"""
    unionise_type(tp::Union{Symbol, Expr})

Returns an expression for the type union of `tp` and `Node{<:tp}`. e.g.
`unionise_type(:Real)` returns `:(Union{Real, Node{<:Real}})`.
"""
function unionise_type(tp::Union{Symbol, Expr})
    (_tp, _info) = remove_vararg(tp)
    tp_clean = (isa(_tp, Expr) && _tp.head == Symbol("<:")) ? _tp.args[1] : _tp
    return replace_vararg(:(Union{$_tp, Node{<:$tp_clean}}), (_tp, _info))
end

"""
    replace_body(unionall::Union{Symbol, Expr}, replacement::Union{Symbol, Expr})

Replace the body of an expression representing a `UnionAll`. e.g.
replace_body(:(Tuple{T, T} where T), :foo) returns the (nonsensical) expression
:(foo where T). If `unionall` is a `Symbol`, then `replacement` is returned.
"""
replace_body(unionall::Expr, replacement::Union{Symbol, Expr}) =
    unionall.head == :where ?
        Expr(:where, replace_body(unionall.args[1], replacement), unionall.args[2:end]...) :
        replacement
replace_body(::Symbol, replacement::Union{Symbol, Expr}) = replacement

"""
    get_body(x::Union{Symbol, Expr})

Get the body from an expression representing a `UnionAll`. e.g. :(Tuple{T, T} where T)
returns :(Tuple{T, T}). If `x` is a `Symbol`, then `x` is returned unaltered.
"""
get_body(unionall::Expr) = unionall.head == :where ? get_body(unionall.args[1]) : unionall
get_body(body::Symbol) = body

"""
    get_types(type_tuple::Expr)

Return the types from a type-tuple expression. e.g. :(Tuple{Float64, Real}) returns a
vector [:Float64, :Real]
"""
get_types(type_tuple::Expr) = type_tuple.args[2:end]

"""
    isa_vararg(symbol_or_expr)

Returns a bool indicating whether `symbol_or_expr` is a `Vararg`.
"""
isa_vararg(sym::Symbol) = (sym == :Vararg)
isa_vararg(expr::Expr) =
    (expr.head == :curly && expr.args[1] == :Vararg) ?
        true :
        expr.head == :where ? isa_vararg(expr.args[1]) : false

"""
    remove_vararg(typ::Union{Symbol, Expr})

Return the type contained by the `typ` if it's a `Vararg` and the `N` parameter if provided.
"""
remove_vararg(typ::Symbol) = isa_vararg(typ) ? (:Any, :Vararg) : (typ, :nothing)
function remove_vararg(typ::Expr)
    if isa_vararg(typ)
        body = get_body(typ)
        new_typ = replace_body(typ, body.args[2])
        vararg_info = length(body.args) == 3 ? body.args[3] : :Vararg
        return new_typ, vararg_info
    else
        return (typ, :nothing)
    end
end

"""
    replace_vararg(typ::SymOrExpr, vararg_info::Tuple)

Convert `typ` to the `Vararg` containing elements of type `typ` specified by
`vararg_info`, which should be a `Tuple` returned from `remove_vararg`.
"""
replace_vararg(typ::SymOrExpr, vararg_info::Tuple) =
    vararg_info[2] == :nothing ?
        typ :
        vararg_info[2] == :no_N || vararg_info[2] == :Vararg ?
            replace_body(typ, :(Vararg{$(get_body(typ))})) :
            replace_body(typ, :(Vararg{$(get_body(typ)), $(vararg_info[2])}))

"""
    parse_kwargs(nt_expr) -> NamedTuple

Accepts an expression containing a `NamedTuple` literal and parses it into a `NamedTuple`
with expressions as values.
"""
function parse_kwargs(nt_expr)
    if isempty(nt_expr.args) || nt_expr == Expr(:call, :NamedTuple)
        return NamedTuple()
    end

    first_arg = first(nt_expr.args)
    if first_arg isa Expr && first_arg.head == :parameters
        return parse_kwargs_parameters(nt_expr)
    elseif first_arg isa Expr && first_arg.head == :(=)
        return parse_kwargs_tuple(nt_expr)
    else
        throw(ArgumentError("Unsupported expression $nt_expr for kwargs;"
            * " they must be passed as a NamedTuple literal"))
    end
end

function parse_kwargs_tuple(tup_expr)
    nt_names = Tuple(first(ex.args) for ex in tup_expr.args)
    return NamedTuple{nt_names}(last(ex.args) for ex in tup_expr.args)
end

function parse_kwargs_parameters(param_tuple_expr)
    # code is the same even though the inner args are also different expression types
    # (:kw vs :())
    return parse_kwargs_tuple(param_tuple_expr.args[1])
end

"""
    parse_is_node(bool_array_expr) -> Vector{Bool}

Accepts an expression containing a `Vector{Bool}` literal and parses it into a
`Vector{Bool}`.
"""
function parse_is_node(bool_array_expr)
    if bool_array_expr.head != :vect
        throw(ArgumentError("Unsupported expression $bool_array_expr for is_node; "
            * "it must be passed as a `Vector{Bool}` literal (e.g., `[true, false]`)"))
    end

    return collect(Bool, bool_array_expr.args)
end


text:
using BrainFlow

# specify logging library to use
BrainFlow.enable_dev_logger(BrainFlow.BOARD_CONTROLLER)

params = BrainFlowInputParams()
board_shim = BrainFlow.BoardShim(BrainFlow.SYNTHETIC_BOARD, params)
board_descr = BrainFlow.get_board_descr(BrainFlow.SYNTHETIC_BOARD)
sampling_rate = board_descr["sampling_rate"]
nfft = BrainFlow.get_nearest_power_of_two(sampling_rate)

BrainFlow.prepare_session(board_shim)
BrainFlow.start_stream(board_shim)
sleep(5)
BrainFlow.stop_stream(board_shim)
data = BrainFlow.get_board_data(board_shim)
BrainFlow.release_session(board_shim)

eeg_channels = board_descr["eeg_channels"]
# second channel of synthetic board is sine wave at 10 Hz, should see huge 'alpha'
data_second_channel = data[eeg_channels[2], :]

# optional: detrend
BrainFlow.detrend(data_second_channel, BrainFlow.LINEAR)
# psd is a tuple of ampls and freqs
psd = BrainFlow.get_psd_welch(data_second_channel, nfft, Integer(nfft / 2), sampling_rate, BrainFlow.BLACKMAN_HARRIS)
band_power_alpha = BrainFlow.get_band_power(psd, 7.0, 13.0)
band_power_beta = BrainFlow.get_band_power(psd, 14.0, 30.0)
println(band_power_alpha / band_power_beta)


text:
# Create example data
x = LinRange(-2, 2, 40)
y = x.^3 .+ x.^2 .+ x
# Draw a plot with panoramic ratio (16:9)
plot(x, y)
aspectratio(16/9)