text:
using Distributions, ZVSimulator

function C(n::Int, c::Float64)
  X = eye(n)
  [(j <= n-i) ? X[i+j, i] = X[i, i+j] = c^j : nothing for i = 1:(n-1), j = 1:(n-1)]
  X
end

df, npars, a = 5., 4, 0.25
c = ((df-2)/df)*C(npars, a)

d = MvTDist(df, zeros(npars), c)

results = psim_rand_vrf(d, nsets=5)

# writedlm_vrf(results, "./", ' ')
# writedlm_reformat("_means.txt", "./_means_pprint.txt", fill("%.3e", 3))


text:
#=

The robust control problem for a monopolist with adjustment costs.  The
inverse demand curve is:

  p_t = a_0 - a_1 y_t + d_t

where d_{t+1} = \rho d_t + \sigma_d w_{t+1} for w_t ~ N(0,1) and iid.
The period return function for the monopolist is

  r_t =  p_t y_t - gam (y_{t+1} - y_t)^2 / 2 - c y_t

The objective of the firm is E_t \sum_{t=0}^\infty \beta^t r_t

For the linear regulator, we take the state and control to be

    x_t = (1, y_t, d_t) and u_t = y_{t+1} - y_t

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-07-05

References
----------

Simple port of the file examples/robust_monopolist.py

http://quant-econ.net/robustness.html#application

=#
using QuantEcon
using Plots
pyplot()
using Grid

# model parameters
a_0     = 100
a_1     = 0.5
rho     = 0.9
sigma_d = 0.05
bet    = 0.95
c       = 2
gam   = 50.0
theta = 0.002
ac    = (a_0 - c) / 2.0

# Define LQ matrices
R = [0 ac    0
     ac -a_1 0.5
     0. 0.5  0]
R = -R  # For minimization
Q = [gam / 2.0]'
A = [1. 0. 0.
     0. 1. 0.
     0. 0. rho]
B = [0. 1. 0.]'
C = [0. 0. sigma_d]'

## Functions

function evaluate_policy(theta, F)
    rlq = RBLQ(Q, R, A, B, C, bet, theta)
    K_F, P_F, d_F, O_F, o_F = evaluate_F(rlq, F)
    x0 = [1.0 0.0 0.0]'
    value = - x0'*P_F*x0 - d_F
    entropy = x0'*O_F*x0 + o_F
    return value[1], entropy[1]  # return scalars
end

function value_and_entropy(emax, F, bw, grid_size=1000)
    if lowercase(bw) == "worst"
        thetas = 1 ./ linspace(1e-8, 1000, grid_size)
    else
        thetas = -1 ./ linspace(1e-8, 1000, grid_size)
    end

    data = Array{Float64}(grid_size, 2)

    for (i, theta) in enumerate(thetas)
        data[i, :] = collect(evaluate_policy(theta, F))
        if data[i, 2] >= emax  # stop at this entropy level
            data = data[1:i, :]
            break
        end
    end
    return data
end

## Main

# compute optimal rule
optimal_lq = LQ(Q, R, A, B, C, zero(B'A), bet=bet)
Po, Fo, Do = stationary_values(optimal_lq)

# compute robust rule for our theta
baseline_robust = RBLQ(Q, R, A, B, C, bet, theta)
Fb, Kb, Pb = robust_rule(baseline_robust)

# Check the positive definiteness of worst-case covariance matrix to
# ensure that theta exceeds the breakdown point
test_matrix = eye(size(Pb, 1)) - (C' * Pb * C ./ theta)[1]
eigenvals, eigenvecs = eig(test_matrix)
@assert all(eigenvals .>= 0)

emax = 1.6e6

# compute values and entropies
optimal_best_case = value_and_entropy(emax, Fo, "best")
robust_best_case = value_and_entropy(emax, Fb, "best")
optimal_worst_case = value_and_entropy(emax, Fo, "worst")
robust_worst_case = value_and_entropy(emax, Fb, "worst")

# we reverse order of "worst_case"s so values are ascending
data_pairs = ((optimal_best_case, optimal_worst_case),
              (robust_best_case, robust_worst_case))

egrid = linspace(0, emax, 100)
egrid_data = Array{Float64}[]
for data_pair in data_pairs
    for data in data_pair
        x, y = data[:, 2], data[:, 1]
        curve(z) = InterpIrregular(x, y, BCnearest, InterpLinear)[z]
        push!(egrid_data, curve(egrid))
    end
end
plot(egrid, egrid_data, color=[:red :red :blue :blue])
plot!(egrid, egrid_data[1], fillrange=egrid_data[2],
      fillcolor=:red, fillalpha=0.1, color=:red, legend=:none)
plot!(egrid, egrid_data[3], fillrange=egrid_data[4],
      fillcolor=:blue, fillalpha=0.1, color=:blue, legend=:none)
plot!(xlabel="Entropy", ylabel="Value")

ext: jl


text:
#=
    Author: Jonathan Payne
    This is simple illustration of the method lines technique for solving
    differential equations.
    cd("/Users/Juddy/Dropbox/Studies/Classes/2014A/Financial Economics/6_Replication/2_Code")
=#

# Runge-Kutta Method
function rk4mono(a, b, h, yinit, fdash)
    # This code calculates ODE using Runge-Kutta 4th order method
    # Solves the problem:
    # dy(x)/dx = fdash(x,y)
    x = a:h:b
    y = zeros(1, length(x))
    y[1] = yinit;
    for i = 1:(length(x)-1)
        k_1 = fdash(x[i], y[i]);
        k_2 = fdash(x[i] + 0.5*h, y[i] + 0.5*h*k_1)
        k_3 = fdash(x[i] + 0.5*h, y[i] + 0.5*h*k_2)
        k_4 = fdash(x[i] + h, y[i] + k_3*h)
        y[i+1] = y[i] + (1/6)*(k_1 + 2*k_2 + 2*k_3 + k_4)*h
    end
    return y
end
a, b, h, yinit = 0, 10, 0.1, 48
fdash(x,y) = 9.8 - 0.196*y
ymono = rk4mono(a, b, h, yinit, fdash)

# Multi-dimensional Runge-Kutta Method
function rk4old(y, dydx, n, x, h, yout, derivs)
    # Given values for n variables y[1,...,n] and their derviatives
    # dydx[1,...,n] known at x, this algorithm uses the fourth-order
    # Runge-Kutta method to advance the solution over an interval h
    # and return the incremented variables as yout[1,...n], which need
    # not be a distinct array from y. The user supplies the routine
    # derivs(x, y, dydx) which returns derivatives dydx at x

    dym = zeros(n)
    dyt = zeros(n)
    yt = zeros(n)
    hh = h*0.5
    h6 = h/6.0
    xh = x+hm

end

function rk4(fdash, t0, x0, tf, N)
    # Approximate the solution of the initial value problem:
    # x'(t) = fdash(t, x),
    # x(t0) = x0
    # using the classical fourth order Runge-Kutta method.
    # Inputs:
    #   * fdash: user defined function
    #   * t0: initial value of the independent variable
    #   * x0: initial value of the dependent variable(s) (column vector)
    #   * tf: final value of the independent variable
    #   * N: number of universally sized time steps to be taken to advance the
    #        the solution from t = t0 to t = tf
    # Outputs:
    #   * wi: vector / matrix containing the values of the approximate solution
    #         to the ode
    #   * ti: number of uniformly sized time steps to be taken to advance t
    neqn = length(x0)
    ti = linspace(t0, tf, N+1)
    wi = zeros(neqn, N+1)
    wi[1:neqn, 1] = x0
    h = (tf - t0) / N
    for i = 1:N
        k1 = h*fdash(t0, x0)
        k2 = h*fdash(t0 + h/2, x0 + k1/2)
        k3 = h*fdash(t0 + h/2, x0 + k2/2)
        k4 = h*fdash(t0 + h, x0 + k3)
        x0 = x0 + (k1 + 2*k2 + 2k3 + k4) / 6
        t0 = t0 + h
        wi[1:neqn, i+1] = x0
    end
    return wi, ti
end
t0, x0, tf, N = 0, 48, 10, 100
fdash(t, x) = 9.8 - 0.196*x
wi, ti = rk4(fdash, t0, x0, tf, N)

t0, x0, tf, N = 0, [0,-4], 10, 100
A = [[1 2], [3 2]]
fdash(t, x) = A*x
xtrue(t) = [ 8/5*exp(-t) - 8/5*exp(4*t), -8/5*exp(-t) - 12/5*exp(4*t) ]
wi, ti = rk4(fdash, t0, x0, tf, N)

function rk4step(fdash, t, x, h)
    # Calculates the the nth step in the classical fourth order Runga-Kutta
    # procedure for solving a differential equation of the form:
    #   x'(t) = fdash(t, x),
    # Inputs:
    #   * fdash: user defined function
    #   * t: previous value of independent variable
    #   * x: previous value of dependent variable
    #   * h: increment
    # Outputs:
    #   * xnew: next step
    #   * tnew: new time
    k1 = h*fdash(t, x)
    k2 = h*fdash(t + h/2, x + k1/2)
    k3 = h*fdash(t + h/2, x + k2/2)
    k4 = h*fdash(t + h, x + k3)
    xnew = x + (k1 + 2*k2 + 2k3 + k4) / 6
    tnew = t + h
    return xnew, tnew
end

function rk4uniform(fdash, t0, x0, tf, N)
    # Approximate the solution of the initial value problem using the
    # classical fourth order Runge-Kutta method with uniform step sizes:
    #   * x'(t) = fdash(t, x),
    #   * x(t0) = x0
    # Inputs:
    #   * fdash: user defined function
    #   * t0: initial value of the independent variable
    #   * x0: initial value of the dependent variable(s) (column vector)
    #   * tf: final value of the independent variable
    #   * N: number of universally sized time steps to be taken to advance the
    #        the solution from t = t0 to t = tf
    # Outputs:
    #   * wi: vector / matrix containing the values of the approximate solution
    #         to the ode
    #   * ti: number of uniformly sized time steps to be taken to advance t
    neqn = length(x0)
    ti = linspace(t0, tf, N+1)
    wi = zeros(neqn, N+1)
    wi[1:neqn, 1] = x0
    h = (tf - t0) / N
    for i = 1:N
        x0, t0 = rk4step(fdash, t0, x0, h)
        wi[1:neqn, i+1] = x0
    end
    return wi, ti
end

# Usage example 1:
t0, x0, tf, N = 0, [0,-4], 10, 100
A = [[1 2], [3 2]]
fdash(t, x) = A*x
xtrue(t) = [ 8/5*exp(-t) - 8/5*exp(4*t), -8/5*exp(-t) - 12/5*exp(4*t) ]
wi2, ti2 = rk4uniform(fdash, t0, x0, tf, N)

# Usage example 2:
function Fdash(t, x)
    xp = zeros(2,1)
    xp[1] = x[2]
    xp[2] = -t*x[1] + exp(t)*x[2] + 3*sin(2*t)
    return xp
end
t0, x0, tf, N = 0, [2,8], 4, 100
xi5, ti5 = rk4uniform(Fdash, t0, x0, tf, N)

# ----------------------------------------
# METHOD OF LINES
# ----------------------------------------
# Solves the heat (or diffusion equation) using the method of lines technique:
#   * u_t = kappa u_xx
#   * u(x, 0) = n(x)
#   * u(0, t) = g_0(t), t>0
#   * u(1, t) = g_1(t), t>0

xs, xf = 0, 5
M = 11 # number of mesh points
h = (xf - xs)/M # size of the space mesh
k = 0.2 # size of the time mesh
xgrid = linspace(xs, xf, M)
uxs(t) = 1
uxf(t) = 1
ut0(x) = x

# Set up the discretisation in the space dimension
# Need to discretise into the form:
#   * U'(t) = A U(t) + g(t)
A = (1/h^2)*(diagm(ones(M-1),-1)+diagm(-2*ones(M))+diagm(ones(M-1),1))
function gfunc(t)
    gvec = zeros(M);
    gvec[1] = uxs(t)
    gvec[end] = uxf(t)
    return gvec
end

# Run the ode solver
fdash(t, Uall) = A*Uall + gfunc(t)
t0, Unit, tf, N = 0, ut0(xgrid), 10, 100
wi2, ti2 = rk4uniform(fdash, t0, Unit, tf, N)

########################################################################
#                                                                      %
#    Example of ADI Method for 2D heat equation                        %
#                                                                      %
#          u_t = u_{xx} + u_{yy} + f(x,t)                              %
#                                                                      %
#    Test problme:                                                     %
#      Exact solution: u(t,x,y) = exp(-t) sin(pi*x) sin(pi*y)          %
#      Source term:  f(t,x,y) = exp(-t) sin(pi*x) sin(pi*y) (2pi^2-1)  %
#                                                                      %
#    Files needed for the test:                                        %
#                                                                      %
#     adi.m:      This file, the main calling code.                    %
#     f.m:        The file defines the f(t,x,y)                        %
#     uexact.m:    The exact solution.                                 %
#                                                                      %
#     Results:         n              e            ratio               %
#                     10           0.0041                              %
#     t_final=0.5     20           0.0010           4.1                %
#                     40           2.5192e-04       3.97               %
#                     80           6.3069e-05       3.9944             %
########################################################################

uexact(t,x,y) = exp(-t)*sin(pi*x)*sin(pi*y)
f(t,x,y) = exp(-t)*sin(pi*x)*sin(pi*y)*(2*pi^2-1)

a = 0; b = 1; c = 0; d = 1; n = 40; tfinal = 0.5
m = n
h = (b-a)/n; dt = h; h1 = h*h
x = a:h:b; y = c:h:d

# ----- Initial condition
t = 0
u1 = zeros(m+1, m+1)
for i=1:m+1
    for j=1:m+1
        u1[i,j] = uexact(t, x[i], y[j])
    end
end

# ----- Loop for time t
k_t = floor(tfinal/dt)
for k = 1:k_t
    t1 = t + dt
    t2 = t + dt/2
    # A. Sweep in the x-direction
    # -- Set the boundary condition
    u2 = zeros(m+1, n+1)
    for i = 1:m+1
        u2[i, 1] = uexact(t2, x[i], y[1])
        u2[i, n+1] = uexact(t2, x[i], y[n+1])
        u2[1, i] = uexact(t2, x[1], y[i])
        u2[m+1, i] = uexact(t2, x[m+1], y[i])
    end
    # -- Iterate
    for j = 2:n
        A = spzeros(m-1, m-1)
        b = zeros(m-1,1)
        for i = 2:m
            b[i-1] = ( (u1[i,j-1] - 2*u1[i,j] + u1[i,j+1])/h1
                      + f(t2, x[i], y[j]) + 2*u1[i,j]/dt )
            if i == 2
                b[i-1] = b[i-1] + uexact(t2, x[i-1], y[j])/h1
                A[i-1, i] = -1/h1
            else
                if i == m
                    b[i-1] = b[i-1] + uexact(t2, x[i+1], y[j])/h1
                    A[i-1, i-2] = -1/h1
                else
                    A[i-1, i] = -1/h1
                    A[i-1, i-2] = -1/h1
                end
            end
            A[i-1, i-1] = 2/dt + 2/h1
        end
        ut = A\b
        for i = 1:m-1
            u1[i+1, j] = ut[i]
        end
    end
    # B. Sweep in the y-direction
    # -- Set the boundary condition
    u1 = zeros(m+1, n+1)
    for i = 1:m+1
        u1[i,1] = uexact(t1, x[i], y[1])
        u1[i,n+1] = uexact(t1, x[i], y[m+1])
        u1[1,i] = uexact(t1, x[1], y[i])
        u1[m+1,i] = uexact(t1, x[m+1], y[i])
    end
    # -- Iterate
    for i = 2:m
        A = spzeros(m-1, m-1)
        b = zeros(m-1,1)
        for j = 2:n
            b[j-1] = ( (u2[i-1,j] - 2*u2[i,j] + u2[i+1,j])/h1
                      + f(t2, x[i], y[j]) + 2*u2[i,j]/dt )
            if j == 2
                b[j-1] = b[j-1] + uexact(t1, x[i], y[j-1])/h1
                A[j-1,j] = -1/h1
            else
                if j == n
                    b[j-1] = b[j-1] + uexact(t1, x[i], y[j+1])/h1
                    A[j-1,j-2] = -1/h1
                else
                    A[j-1,j] = -1/h1
                    A[j-1,j-2] = -1/h1
                end
            end
            A[j-1,j-1] = 2/dt + 2/h1
        end
        # -- Solve the system
        ut = A\b
        for j = 1:n-1
            u1[i,j+1] = ut[j]
        end
    end
t = t + dt
end

# ----- Data analysis
ue = zeros(m+1, m+1)
for i = 1:m+1
    for j = 1:n+1
        ue[i,j] = uexact(tfinal, x[i], y[i])
    end
end
err = maximum(abs(u1-ue))

max_stars_repo_name: jepayne/NYU

max_issues_repo_name: jepayne/NYU

max_forks_repo_name: jepayne/NYU


text:
module CTScan

Base.Experimental.@optlevel 0

export AbstractCTScanner, CTScanner, FBPScanner
export rename, rename!, transform_data, transform_data!
export compute_gray_scale
export make_phantom, create_image!
export project_image!, reconstruct_image!
export para2fan!, fan2para!
export compute_rmse, compute_χ2, fit_gray_scale
export load_image!, load_sinogram!, load_tomogram!
export sample_sinogram!

import Base: show, copy, similar, getproperty
using IntervalSets
import Statistics; const Stat = Statistics
import CurveFit; const Fit = CurveFit
import ..Utils: _atype
using ..Applicative
using ..Monads
using ..Geometry
import ..AbstractAlgorithms: reconstruct_image, project_image
using ..AbstractAlgorithms
using ..RadonAlgorithm
import ..CTImages: ctimage, ctsinogram, cttomogram
using ..CTImages
import ..FanBeam: para2fan, fan2para
import ..CalibrationBase:
    calibrate_image,
    calibrate_image!,
    calibrate_tomogram,
    calibrate_tomogram!
using ..Calibration
using ..TestImages
import ..CTIO: load_image, write_image
import ..CTIO: load_sinogram, write_sinogram
import ..CTIO: load_tomogram, write_tomogram

include("CTImageData.jl")
using .CTImageData: AbstractCTData, CTData

"""
    abstract type AbstractCTScanner{<:AbstractArray,<:AbstractGeometry} end

Abstract type for tests in this suite.
"""
abstract type AbstractCTScanner end

const CTScannerNameType = Union{Symbol,String}

function rename(gst::AbstractCTScanner, name::AbstractString)
    typeof(gst)(gst.geometry, gst.data; name)
end

function rename!(gst::AbstractCTScanner, name::AbstractString)
    gst.name = name
end

_atype(gst::AbstractCTScanner) = _atype(gst.data)

ctimage(gst::AbstractCTScanner) = ctimage(gst.data)
ctsinogram(gst::AbstractCTScanner) = ctsinogram(gst.data)
cttomogram(gst::AbstractCTScanner) = cttomogram(gst.data)

similar(gst::AbstractCTScanner, img::AbstractCTImage...) =
    similar(gst, CTData(img...))

function similar(
    gst::AbstractCTScanner,
    g::AbstractGeometry,
    img::AbstractCTImage...,
)
    similar(gst, g, CTData(img...))
end

for (nm, ctf) ∈ pairs(CTImages.ctfn)
    v = Val{nm}
    @eval begin
        t = CTImages.ctnames.$nm
        function transform_data(f::Function, gst::AbstractCTScanner, ::$v)
            mbind($ctf(gst)) do img
                newdata = similar(gst.data, map(f, something(img)))
                similar(gst, newdata)
            end
        end
        function transform_data!(f::Function, gst::AbstractCTScanner, ::$v)
            mbind($ctf(gst)) do img
                newdata = similar(gst.data, map(f, something(img)))
                gst.data = newdata
            end
        end
        transform_data(f::Function, gst::AbstractCTScanner, ::Type{t}) =
            transform_data(f, gst, $v())
        transform_data!(f::Function, gst::AbstractCTScanner, ::Type{t}) =
            transform_data!(f, gst, $v())
    end
end

transform_data(f::Function, gst::AbstractCTScanner, s::Symbol) =
    transform_data(f, gst, Val(s))
transform_data!(f::Function, gst::AbstractCTScanner, s::Symbol) =
    transform_data!(f, gst, Val(s))

struct CTScanner{A<:AlgorithmInfo,D<:AbstractCTData} <: AbstractCTScanner
    name::CTScannerNameType
    study_id::String
    ainfo::A
    data::D

    """
    CTScanner{A,D}(
        ainfo::A[, data::D];
        name::Optional{Union{Symbol,String}} = nothing,
        study_id::Optional{String} = "Unknown"
    ) where {A <: AlgorithmInfo, D <: AbstractCTData}

    Construct `CTScanner` from an algorithm.
    """
    function CTScanner{A,D}(
        ainfo::A,
        data::D = CTData();
        name::Optional{CTScannerNameType} = nothing,
        study_id::Optional{String} = nothing,
    ) where {A <: AlgorithmInfo,D <: AbstractCTData}
        new(
            maybe(alg_name(ainfo), name),
            maybe("Unknown", study_id),
            ainfo,
            data
        )
    end
end

@inline function getproperty(ct::CTScanner, s::Symbol)
    s ≡ :geometry && return ct.ainfo.geometry
    s ≡ :params && return ct.ainfo.params
    s ≡ :algorithm && return ct.ainfo.algorithm
    getfield(ct, s)
end

function CTScanner(alg::AlgorithmInfo, data::AbstractCTData; kwargs...)
    CTScanner{typeof(alg), typeof(data)}(alg, data; kwargs...)
end

function CTScanner(
    alg::AlgorithmInfo,
    imgs::T...;
    kwargs...
) where {T<:AbstractCTImage}
    CTScanner{typeof(alg), CTData{eltype(T),_atype(T)}}(
        alg, CTData(imgs...); kwargs...)
end

function CTScanner(alg::AlgorithmInfo, gst::AbstractCTScanner; kwargs...)
    CTScanner(alg, gst.data; kwargs...)
end

function similar(gst::CTScanner, data::Optional{AbstractCTData} = nothing)
    CTScanner(
        gst.ainfo,
        maybe(similar(gst.data), data);
        gst.name,
        gst.study_id,
    )
end

const FBPScanner{G,D} = CTScanner{FBPInfo{G},D}

"""
    FBPScanner(
        geometry::AbstractGeometry,
        data::AbstractCTData;
        name::Optional{String},
        study_id::Optional{String}
    )

Construct a `FBPScanner` object.
"""
function FBPScanner(
    geometry::G,
    data::D = CTData(eltype(G));
    name::Optional{CTScannerNameType} = nothing,
    study_id::Optional{String} = nothing,
) where {G <: AbstractGeometry, D <: AbstractCTData}
    FBPScanner{G,D}(FBPInfo(geometry), data; name, study_id)
end

"""
    FBPScanner(
        geometry::AbstractGeometry,
        image::AbstractTestImage;
        name::Optional{String},
        study_id::Optional{String}
    )

Construct a `FBPScanner` object.
"""
function FBPScanner(
    geometry::AbstractGeometry,
    gs::AbstractTestImage;
    name::Optional{CTScannerNameType} = nothing,
    study_id::Optional{String} = nothing,
)
    FBPScanner(geometry, gs.image; name, study_id)
end

"""
    FBPScanner(
        image::AbstractTestImage;
        name::Optional{String},
        study_id::Optional{String}
    )

Construct a `FBPScanner` object.
"""
function FBPScanner(
    gs::AbstractTestImage;
    name::Optional{CTScannerNameType} = nothing,
    study_id::Optional{String} = nothing,
)
    FBPScanner(ParallelBeamGeometry(gs), gs; name, study_id)
end

"""
    FBPScanner(gst::AbstractCTScanner; name::Optional{String}, study_id::Optional{String})

Construct a `FBPScanner` object from another `AbstractCTScanner` object. This allows to reuse data.
"""
function FBPScanner(
    gst::AbstractCTScanner;
    name::Optional{CTScannerNameType} = nothing,
    study_id::Optional{String} = nothing,
)
    data = CTData(ctimage(gst), ctsinogram(gst), nothing)
    name = maybe(gst.name, name)
    study_id = maybe(gst.study_id, study_id)
    FBPScanner(gst.ainfo, data; name, study_id)
end

function FBPScanner(g::AbstractGeometry, imgs::AbstractCTImage...; kwargs...)
    FBPScanner(g, CTData(imgs...); kwargs...)
end

function copy(gst::FBPScanner)
    FBPScanner(gst.geometry, copy(gst.data); gst.name, gst.study_id)
end

function similar(
    gst::FBPScanner,
    g::AbstractGeometry,
    data::Optional{CTData} = nothing,
)
    FBPScanner(
        g,
        maybe(gst.data, data);
        gst.name,
        gst.study_id,
    )
end

function show(io::IO, gst::CTScanner)
    print(
        io,
        """*** $(gst.name) Scanner ***
        $(gst.algorithm)""",
    )
end

"""
    create_image!(gst::AbstractCTScanner, par::ImageParams)

Create gray scale image for the test `gst`.
"""
function create_image!(gst::AbstractCTScanner, par::ImageParams)
    #similar(gst, similar(gst.data, create_image(par)))
    gst.data.image = create_image(par)
    gst
end

"""
    project_image(gst::AbstractCTScanner, alg::Optional{A}; <keyword arguments>)
        where {A <: AbstractProjectionAlgorithm}

Compute sinogram for the test `gst`.

Keyword arguments depend on the algorithm employed, please see the relative
documentation.
"""
function project_image(
    gst::AbstractCTScanner,
    alg::Optional{A} = nothing;
    kwargs...
) where {A <: AbstractProjectionAlgorithm}
    alg = isnothing(alg) ? Radon() : alg
    sinog = project_image(ctimage(gst), RadonInfo(gst.geometry, alg); kwargs...)
    similar(gst, similar(gst.data, sinog))
end

function project_image!(
    gst::AbstractCTScanner,
    alg::Optional{A} = nothing;
    kwargs...
) where {A <: AbstractProjectionAlgorithm}
    alg = isnothing(alg) ? Radon() : alg
    sinog = project_image(ctimage(gst), gst.geometry, alg; kwargs...)
    gst.data.sinog = sinog
    gst
end

function para2fan(gst::AbstractCTScanner; kwargs...)
    fbg, sinog_fan = ctsinogram(gst) ↣ para2fan(gst.geometry; kwargs...)
    similar(gst, fbg, similar(gst.data, sinog_fan))
end

function para2fan!(gst::AbstractCTScanner; kwargs...)
    fbg, sinog_fan = ctsinogram(gst) ↣ para2fan(gst.geometry; kwargs...)
    gst.data.sinog = sinog_fan
    gst
end

function para2fan(gst::AbstractCTScanner, fbg::FanBeamGeometry; kwargs...)
    fbg, sinog_fan = ctsinogram(gst) ↣ para2fan(fbg; kwargs...)
    similar(gst, fbg, similar(gst.data, sinog_fan))
end

function fan2para(gst::AbstractCTScanner; kwargs...)
    pbg, sinog_para = ctsinogram(gst) ↣ fan2para(gst.geometry; kwargs...)
    similar(gst, pbg, similar(gst.data, sinog_para))
end

function fan2para!(gst::AbstractCTScanner; kwargs...)
    pbg, sinog_para = ctsinogram(gst) ↣ fan2para(gst.geometry; kwargs...)
    gst.data.sinog = sinog_para
    gst
end

"""
    reconstruct_image(gst::AbstractCTScanner; <keyword arguments>)

Reconstruct CT image from test `gst`.

Keyword arguments depend on the specific algorithm used, please see the relative
documentation.
"""
function reconstruct_image(gst::AbstractCTScanner; kwargs...)
    tomog = reconstruct_image(
        ctsinogram(gst), gst.ainfo; kwargs...)
    similar(gst, similar(gst.data, tomog))
end

function reconstruct_image!(gst::AbstractCTScanner; kwargs...)
    tomog = reconstruct_image(
        ctsinogram(gst), gst.ainfo; kwargs...)
    gst.data.tomog = tomog
    gst
end

"""
    load_sinogram(io::IO, gst::AbstractCTScanner)

Load sinogram from stream `io`.
"""
function load_sinogram(io::IO, gst::AbstractCTScanner)
    sinog = load_sinogram(io; gst.geometry.nd, gst.geometry.nϕ)
    similar(gst, similar(gst.data, sinog))
end

function load_sinogram!(io::IO, gst::AbstractCTScanner)
    sinog = load_sinogram(io; gst.geometry.nd, gst.geometry.nϕ)
    gst.data.sinog = sinog
    gst
end

"""
    load_sinogram(f::AbstractString, gst::AbstractCTScanner)

Load sinogram from file `f`.
"""
function load_sinogram(f::AbstractString, gst::AbstractCTScanner)
    open(f) do s
        load_sinogram(s, gst)
    end
end

function load_sinogram!(f::AbstractString, gst::AbstractCTScanner)
    open(f) do s
        load_sinogram!(s, gst)
    end
end

"""
    load_image(io::IO, gst::AbstractCTScanner)

Load image from stream `io`.
"""
function load_image(io::IO, gst::AbstractCTScanner)
    image = load_image(io; gst.geometry.rows, gst.geometry.cols)
    similar(gst, similar(gst.data, image))
end

function load_image!(io::IO, gst::AbstractCTScanner)
    image = load_image(io; gst.geometry.rows, gst.geometry.cols)
    gst.data.image = image
    gst
end

"""
    load_image(f::AbstractString, gst::AbstractCTScanner)

Load image from file `f`.
"""
function load_image(f::AbstractString, gst::AbstractCTScanner)
    open(f) do s
        load_image(s, gst)
    end
end

function load_image!(f::AbstractString, gst::AbstractCTScanner)
    open(f) do s
        load_image!(s, gst)
    end
end

"""
    write_image(io::IO, gst::AbstractCTScanner)

Write input image to stream `io`.
"""
function write_image(io::IO, gst::AbstractCTScanner)
    mbind(ctimage(gst)) do x
        write_image(io, x; header = false, row_major = true)
    end
end

"""
    write_image(f::AbstractString, gst::AbstractCTScanner)

Write input image to file `f`.
"""
function write_image(f::AbstractString, gst::AbstractCTScanner)
    open(f, "w") do s
        write_image(s, gst)
    end
end

"""
    write_sinogram(io::IO, gst::AbstractCTScanner)

Write sinogram to stream `io`.
"""
function write_sinogram(io::IO, gst::AbstractCTScanner)
    mbind(ctsinogram(gst)) do x
        write_sinogram(io, x; header = false)
    end
end

"""
    write_sinogram(f::AbstractString, gst::AbstractCTScanner)

Write sinogram to file `f`.
"""
function write_sinogram(f::AbstractString, gst::AbstractCTScanner)
    open(f, "w") do s
        write_sinogram(s, gst)
    end
end

"""
    load_tomogram(io::IO, gst::AbstractCTScanner)

Load tomogram from stream `io`.
"""
function load_tomogram(io::IO, gst::AbstractCTScanner)
    tomog = load_tomogram(io; gst.geometry.rows, gst.geometry.cols)
    similar(gst, similar(gst.data, tomog))
end

function load_tomogram!(io::IO, gst::AbstractCTScanner)
    tomog = load_tomogram(io; gst.geometry.rows, gst.geometry.cols)
    gst.data.tomog = tomog
    gst
end

"""
    load_tomogram(f::AbstractString, gst::AbstractCTScanner)

Load tomogram from file `f`.
"""
function load_tomogram(f::AbstractString, gst::AbstractCTScanner)
    open(f) do s
        load_tomogram(s, gst)
    end
end

function load_tomogram!(f::AbstractString, gst::AbstractCTScanner)
    open(f) do s
        load_tomogram!(s, gst)
    end
end

"""
    write_tomogram(io::IO, gst::AbstractCTScanner)

Write reconstructed image to stream `io`.
"""
function write_tomogram(io::IO, gst::AbstractCTScanner)
    mbind(cttomogram(gst)) do x
        write_tomogram(io, x; header = false, row_major = true)
    end
end

"""
    write_tomogram(f::AbstractString, gst::AbstractCTScanner)

Write tomogram to file `f`.
"""
function write_tomogram(f::AbstractString, gst::AbstractCTScanner)
    open(f, "w") do s
        write_tomogram(s, gst)
    end
end

"""
    sample_sinogram(gst::AbstractCTScanner; <keyword arguments>)

Returns a new test object with the resampled sinogram.
"""
function sample_sinogram(gst::AbstractCTScanner; kwargs...)
    sinog = ctsinogram(gst) ↣ sample_sinogram(; kwargs...)
    similar(gst, similar(gst.data, sinog))
end

function sample_sinogram!(gst::AbstractCTScanner; kwargs...)
    sinog = ctsinogram(gst) ↣ sample_sinogram(; kwargs...)
    gst.data.sinog = sinog
    gst
end

"""
    calibrate_image(gst::AbstractCTScanner, imp::ImageParams; interval=nothing, window=nothing)

Perform calibration of input image using image parameters as reference.
"""
function calibrate_image(
    gst::AbstractCTScanner,
    imp::AbstractImageParams;
    interval::Optional{ClosedInterval{T}} = nothing,
    window::Optional{ClosedInterval{U}} = nothing,
) where {T<:Real,U<:Real}
    cimg = calibrate_image(
        ctimage(gst),
        imp;
        interval,
        window,
    )
    similar(gst, maybe(gst.data, cimg) do x
        similar(gst.data, x)
    end)
end

function calibrate_image!(
    gst::AbstractCTScanner,
    imp::AbstractImageParams;
    interval::Optional{ClosedInterval{T}} = nothing,
    window::Optional{ClosedInterval{U}} = nothing,
) where {T<:Real,U<:Real}
    cimg = calibrate_image(
        ctimage(gst),
        imp;
        interval,
        window,
    )
    gst.data.image = cimg
    gst
end

"""
    calibrate_image(gst::AbstractCTScanner; min_pos, max_pos, interval=0..1, window=nothing)

Perform calibration of input image with reference values.
"""
function calibrate_image(
    gst::AbstractCTScanner;
    min_pos,
    max_pos,
    interval::ClosedInterval = 0..1,
    window::Optional{ClosedInterval{T}} = nothing,
) where {T<:Real}
    cimg = calibrate_image(
        ctimage(gst);
        interval,
        min_pos,
        max_pos,
        window,
    )
    similar(gst, maybe(gst.data, cimg) do x
        similar(gst.data, x)
    end)
end

function calibrate_image!(
    gst::AbstractCTScanner;
    min_pos,
    max_pos,
    interval::ClosedInterval = 0..1,
    window::Optional{ClosedInterval{T}} = nothing,
) where {T<:Real}
    cimg = calibrate_image(
        ctimage(gst);
        interval,
        min_pos,
        max_pos,
        window,
    )
    gst.data.image = cimg
    gst
end

"""
    calibrate_tomogram(gst::AbstractCTScanner, imp::ImageParams; interval=nothing, window=nothing)

Perform calibration of reconstructed image using image parameters as reference.
"""
function calibrate_tomogram(
    gst::AbstractCTScanner,
    imp::AbstractImageParams;
    interval::Optional{ClosedInterval{T}} = nothing,
    window::Optional{ClosedInterval{U}} = nothing,
) where {T<:Real,U<:Real}
    cimg = calibrate_tomogram(
        cttomogram(gst),
        imp;
        interval,
        window,
    )
    similar(gst, maybe(gst.data, cimg) do x
        similar(gst.data, x)
    end)
end

function calibrate_tomogram!(
    gst::AbstractCTScanner,
    imp::AbstractImageParams;
    interval::Optional{ClosedInterval{T}} = nothing,
    window::Optional{ClosedInterval{U}} = nothing,
) where {T<:Real,U<:Real}
    cimg = calibrate_tomogram(
        cttomogram(gst),
        imp;
        interval,
        window,
    )
    gst.data.tomog = cimg
    gst
end

"""
    calibrate_tomogram(gst::AbstractCTScanner; min_pos, max_pos, interval=0..1, window=nothing)

Perform calibration of reconstructed image with reference values.
"""
function calibrate_tomogram(
    gst::AbstractCTScanner;
    min_pos,
    max_pos,
    interval::ClosedInterval = 0..1,
    window::Optional{ClosedInterval{T}} = nothing,
) where {T<:Real}
    cimg = calibrate_tomogram(
        cttomogram(gst);
        min_pos,
        max_pos,
        interval,
        window,
    )
    similar(gst, maybe(gst.data, cimg) do x
        similar(gst.data, x)
    end)
end

function calibrate_tomogram!(
    gst::AbstractCTScanner;
    min_pos,
    max_pos,
    interval::ClosedInterval = 0..1,
    window::Optional{ClosedInterval{T}} = nothing,
) where {T<:Real}
    cimg = calibrate_tomogram(
        cttomogram(gst);
        min_pos,
        max_pos,
        interval,
        window,
    )
    gst.data.tomog = cimg
    gst
end

"""
    compute_gray_scale(image::AbstractMatrix, imp::ImageParams; mean = false)

Compute the corresponding gray scale for the test image `image`.
"""
function compute_gray_scale(image::AbstractMatrix, imp::ImageParams; mean = false)
    mean && return Stat.mean(image[gray_scale_indices(imp)...], dims = 1) |> vec
    rind, cind = gray_scale_indices(imp)
    i = sum(extrema(rind)) ÷ 2
    image[i,cind]
end

# function compute_gray_scale(
#     image::CTImageOrTomog, params::ImageParams; kwargs...)
#     mbind(image) do x
#         compute_gray_scale(x, params; kwargs...)
#     end
# end

compute_gray_scale(grsc::AbstractGrayScale; kwargs...) =
    compute_gray_scale(grsc.image, grsc.params)

compute_gray_scale(imp::ImageParams; kwargs...) = x -> compute_gray_scale(x, imp; kwargs...)

"""
    compute_gray_scale(gst::AbstractCTScanner, imp::ImageParams; mean = false)

Compute the corresponding gray scale for the reconstructed image.
"""
function compute_gray_scale(gst::AbstractCTScanner, grsc::AbstractGrayScale; kwargs...)
    cttomogram(gst) ↣ compute_gray_scale(grsc.params; kwargs...)
end

function make_phantom(f::Function, g::AbstractParallelBeamGeometry)
    g |> f |> CTImage
end

function compute_rmse(
    grimg::AbstractGrayScale,
    expected::AbstractVector,
    measured::AbstractVector;
    relative = false,
    scale::Optional{ClosedInterval} = nothing,
)
    scale = maybe(grimg.gray_scale, scale)
    rmse = Stat.mean(abs2, measured - expected) |> sqrt
    relative && return rmse / width(scale)
    rmse
end

function _fit_line(ys)
    xs = eachindex(ys)
    a, b = Fit.linear_fit(xs, ys)
    @. a + b * xs
end

function fit_gray_scale(::Union{GrayScaleLine,WhiteRect}, data::AbstractVector)
    _fit_line(data)
end

function fit_gray_scale(grimg::GrayScalePyramid, data::AbstractVector)
    plateau_len = plateau_length(grimg)
    len = (grimg.width - plateau_len) ÷ 2
    len == 0 && return _fit_line(data)
    steps = accumulate((len, len + plateau_len, grimg.width), init = 0) do x, y
        (last(x) + 1):y
    end
    lines = _fit_line.(data[l] for l ∈ steps)
    reduce(vcat, lines)
end

function compute_rmse(grimg::Union{GrayScaleLine,WhiteRect}, data::AbstractVector; relative = false, scale = nothing)
    fitted_data = fit_gray_scale(grimg, data)
    compute_rmse(grimg, fitted_data, data; relative, scale)
end

function compute_rmse(grimg::GrayScalePyramid, data::AbstractVector; relative = false, scale = nothing)
    fitted_data = fit_gray_scale(grimg, data)
    compute_rmse(grimg, fitted_data, data; relative, scale)
end

function compute_χ2(grimg::AbstractGrayScale, expected::AbstractVector, measured::Vector, variances::Vector)
    ν = length(expected) - 2
    sum(@. (measured - expected) ^ 2 / variances) / ν
end

end # module


text:
using LabelledArrays, Test

AB = @SLVector (:a, :b)
x = AB(1, 2)
@test string(x) == "[a=1, b=2]"
@test string(Float16.(x)) == "Float16[a=1.0, b=2.0]"

y = @LArray [1 2; 3 4] (:a, :b, :c, :d)
@test string(y) == "[a=1 c=2; b=3 d=4]"
@test string(Float16.(y)) == "Float16[a=1.0 c=2.0; b=3.0 d=4.0]"


text:
# test sonar functions directly

# addprocs(2)

using KernelDensityEstimate
using TransformUtils
using IncrementalInference
using RoME
using Distributions

meas = LinearRangeBearingElevation((3.0,3e-4),(0.2,3e-4))

# Functor for efficient functional programming, avoids type_inference at each call
# fp! = WrapParam{reuseLBRA}(zeros(3), zeros(6), zeros(3), reuseLBRA(0))

N = 100
X, pts = 0.01*randn(6,N), zeros(3,N);
t = Array{Array{Float64,2},1}()
push!(t,X)
push!(t,pts)

fp! = GenericWrapParam{LinearRangeBearingElevation}(meas, t, 2, 1, (zeros(0,1),) , RoME.getSample)

# pre-emptively populate the measurements, kept separate since nlsolve calls fp(x, res) multiple times
fp!.measurement = fp!.samplerfnc(fp!.usrfnc!, N)
# fp!(x, res)
@time fp!(zeros(3), zeros(3))

@show zDim = size(fp!.measurement[1],1)
fr = FastRootGenericWrapParam{LinearRangeBearingElevation}(fp!.params[fp!.varidx], zDim, fp!)

@time for fp!.particleidx in 1:N
  numericRootGenericRandomizedFnc!( fr )
end

warn("still need to insert kld(..) test to ensure this is working")

p1 = kde!(pts);

println("Test back projection from ")

meas = LinearRangeBearingElevation((3.0,3e-4),(0.2,3e-4))

N = 200
pts, L = 0.01*randn(6,N), zeros(3,N);
L[1,:] += 3.0
L[2,:] += 0.65
t = Array{Array{Float64,2},1}()
push!(t,pts)
push!(t,L)

fp! = GenericWrapParam{LinearRangeBearingElevation}(meas, t, 1, 1, (zeros(0,1),) , RoME.getSample)

# pre-emptively populate the measurements, kept separate since nlsolve calls fp(x, res) multiple times
fp!.measurement = fp!.samplerfnc(fp!.usrfnc!, N)
# fp!(x, res)
@time fp!(zeros(3), zeros(6))

@show zDim = size(fp!.measurement[1],1)
fr = FastRootGenericWrapParam{LinearRangeBearingElevation}(fp!.params[fp!.varidx], zDim, fp!)

@time for fp!.particleidx in 1:N
  numericRootGenericRandomizedFnc!( fr )
end

p2 = kde!(pts);

# using Gadfly
#
# axis = [[1.5;3.5]';[-1.25;1.25]';[-1.0;1.0]']
# draw( PDF("/home/dehann/Desktop/test.pdf",30cm,20cm),
#       plotKDE( p1, dimLbls=["x";"y";"z"], axis=axis)  )
# #
#
# axis = [[-2.5;2.5]';[-2.5;2.5]';[-2.5;2.5]';[-2pi;2pi]';[-2pi;2pi]';[-2pi;2pi]']
# draw( PDF("/home/dehann/Desktop/test.pdf",30cm,20cm),
#       plotKDE( p2, dimLbls=["x";"y";"z"; "roll"; "pitch"; "yaw"], axis=axis)  )

# do test directly in factor graph
fg = initfg()
# @show fg.registeredModuleFunctions
N = 100
warn("Breaks if not set to 100")

initCov = eye(6)
[initCov[i,i] = 0.01 for i in 4:6];
odoCov = deepcopy(initCov)

println("Adding PriorPose3 to graph...")
X, pts = 0.01*randn(6,N), zeros(3,N);

v1 = addNode!(fg,:x1,  X,  N=N)
initPosePrior = PriorPose3( MvNormal(zeros(6), initCov) )
f1  = addFactor!(fg,[v1], initPosePrior)

println("Adding LinearRangeBearingElevation to graph...")
# implemented in SensorModels
meas = LinearRangeBearingElevation((3.0,3e-4),(0.2,3e-4))

@time X = getVal(v1)
@time pts = X + meas
p1 = kde!(pts); # visual checking

v2 = addNode!(fg, :l1, pts, N=N)
f2 = addFactor!(fg,[v1;v2],meas)

L1pts = evalFactor2(fg, f2, fg.IDs[:l1])
X1pts = evalFactor2(fg, f2, fg.IDs[:x1])

tree = wipeBuildNewTree!(fg)
inferOverTreeR!(fg, tree, N=N)
# inferOverTree!(fg, tree, N=N)

# using ProfileView
# Profile.clear()

# @time for i in 1:200
#   backprojectRandomized!(meas, L, pts, i, fp!)
# end
#
# # ProfileView.view()
#
#
# # work on speeding up and refactoring the numericroot operation
# # did not yet have the desired affect, lots of memory still being claimed
# println("Compute with FastGenericRoot memory structure")
# fgr = FastGenericRoot{WrapParam}(6, 3, fp!)
# @time for idx in 1:200
#   fp!.landmark[1:3] = L[1:3,idx]
#   fp!.pose[1:6] = pts[1:6,idx]
#   fp!.z[1:3] = getSample(meas)
#   copy!(fgr.X, pts[1:6,idx]) #initial guess x0
#   numericRootGenericRandomizedFnc!( fgr )
#   # backprojectRandomized!(fgr)
# 	pts[1:6,idx] = fgr.Y[1:6]
# end
#
#
# pts, L = 0.01*randn(6,200), zeros(3,200);
# L[1,:] += 3.0
# L[2,:] += 0.65
#
#
# # going faster
# fpA! = WrapParamArray(L, pts, zeros(3), 1, reuseLBRA(0))
#
# # @everywhere begin
# # (p::WrapParamArray{reuseLBRA})(x::Vector{Float64}, res::Vector{Float64}) =
# #     residualLRBE!(res, p.z, x, p.landmark[:,p.idx], p.reuse)
# # end
# # Profile.clear()
#
# println("Compute with FastGenericRoot memory structure")
# fgr = FastGenericRoot{WrapParamArray}(6, 3, fpA!)
#
# @time for idx in 1:200
#   fpA!.idx = idx
#   getSample!(fpA!.z,  meas)
#   # fpA!.z[1:3] = getSample(meas)
#   copy!(fgr.X, pts[1:6,idx]) #initial guess x0
#   numericRootGenericRandomizedFnc!( fgr )
#   # backprojectRandomized!( fgr )
# 	pts[1:6,idx] = fgr.Y[1:6]
# end
#

# ProfileView.view()

#  visualization


text:
export TabularPolicy

using Flux: onehot

Base.@kwdef struct TabularPolicy{S,A} <: AbstractPolicy
    table::Dict{S,A} = Dict{Int, Int}()
    n_action::Int
end

(p::TabularPolicy)(env::AbstractEnv) = p(state(env))
(p::TabularPolicy{S})(s::S) where S = p.table[s]

function RLBase.update!(p::TabularPolicy, target::Pair)
    p.table[first(target)] = last(target)
end

RLBase.prob(p::TabularPolicy, s, a) = p.table[s] == a ? 1.0 : 0.


text:
module Codex

using Requires: @require

using Reexport
@reexport using DataFrames
@reexport using Dates
@reexport using Statistics

const PROJECT_ROOT = string(pkgdir(Codex))::String

include("utils.jl")
export Output, apply, dirparent, has_duplicates, nt2dict
include("stats.jl")
include(joinpath("TransformExport", "transformexport.jl"))
include(joinpath("Questionnaires", "questionnaires.jl"))

function __init__()
    @require JSON2="2535ab7d-5cd8-5a07-80ac-9b1792aadce3" include("gitlab.jl")
end

end # module


text:
import JLD2

"""
Adds simplistic checkpointing to a DFTK self-consistent field calculation.
"""
function ScfSaveCheckpoints(filename="dftk_scf_checkpoint.jld2"; keep=false, overwrite=false)
    # TODO Save only every 30 minutes or so
    function callback(info)
        if info.n_iter == 1
            isfile(filename) && !overwrite && error(
                "Checkpoint $filename exists. Use 'overwrite=true' to force overwriting."
            )
        end
        if info.stage == :finalize
            if mpi_master() && !keep
                isfile(filename) && rm(filename)  # Cleanup checkpoint
            end
        else
            scfres = (; (k => v for (k, v) in pairs(info) if !startswith(string(k), "ρ"))...)
            scfres = merge(scfres, (ρ=info.ρout, ρspin=info.ρ_spin_out))
            save_scfres(filename, scfres)
        end
        info
    end
end

function save_scfres_master(file::AbstractString, scfres::NamedTuple, ::Val{:jld2})
    !mpi_master() && error(
        "This function should only be called on MPI master after the k-point data has " *
        "been gathered with `gather_kpts`."
    )

    JLD2.jldopen(file, "w") do jld
        jld["__propertynames"] = propertynames(scfres)
        jld["ρ_real"]          = scfres.ρ.real
        jld["ρspin_real"]      = isnothing(scfres.ρspin) ? nothing : scfres.ρspin.real
        jld["basis"]           = scfres.basis

        for symbol in propertynames(scfres)
            symbol in (:ham, :basis, :ρ, :ρspin, :energies) && continue  # special
            jld[string(symbol)] = getproperty(scfres, symbol)
        end

        jld
    end
end

function load_scfres(jld::JLD2.JLDFile)
    basis = jld["basis"]
    scfdict = Dict{Symbol, Any}(
        :ρ     => from_real(basis, jld["ρ_real"]),
        :ρspin => nothing,
        :basis => basis
    )
    if !isnothing(jld["ρspin_real"])
        scfdict[:ρspin] = from_real(basis, jld["ρspin_real"])
    end

    kpt_properties = (:ψ, :occupation, :eigenvalues)  # Need splitting over MPI processes
    for sym in kpt_properties
        scfdict[sym] = jld[string(sym)][basis.krange_thisproc]
    end
    for sym in jld["__propertynames"]
        sym in (:ham, :basis, :ρ, :ρspin, :energies) && continue  # special
        sym in kpt_properties && continue
        scfdict[sym] = jld[string(sym)]
    end

    energies, ham = energy_hamiltonian(basis, scfdict[:ψ], scfdict[:occupation];
                                       ρ=scfdict[:ρ], ρspin=scfdict[:ρspin],
                                       eigenvalues=scfdict[:eigenvalues],
                                       εF=scfdict[:εF])

    scfdict[:energies] = energies
    scfdict[:ham]      = ham
    (; (sym => scfdict[sym] for sym in jld["__propertynames"])...)
end
load_scfres(file::AbstractString) = JLD2.jldopen(load_scfres, file, "r")

#
# Custom serialisations
#
struct PlaneWaveBasisSerialisation{T <: Real}
    model::Model{T}
    Ecut::T
    variational::Bool
    kcoords::Vector{Vec3{T}}
    ksymops::Vector{Vector{SymOp}}
    kgrid::Union{Nothing,Vec3{Int}}
    kshift::Union{Nothing,Vec3{T}}
    fft_size::Tuple{Int, Int, Int}
    symmetries::Vector{SymOp}
end
JLD2.writeas(::Type{PlaneWaveBasis{T}}) where {T} = PlaneWaveBasisSerialisation{T}

function Base.convert(::Type{PlaneWaveBasisSerialisation{T}},
                      basis::PlaneWaveBasis{T}) where {T}
    # Notice: This function is only meaningful on a basis which has gathered
    # all k-Point information locally. So before using this in MPI-distributed calculations,
    # ensure to call `mpi_kgather` on the PlaneWaveBasis object.

    # Number of distinct k-Point coordinates
    n_kcoords = div(length(basis.kpoints), basis.model.n_spin_components)
    PlaneWaveBasisSerialisation{T}(
        basis.model,
        basis.Ecut,
        basis.variational,
        getproperty.(basis.kpoints[1:n_kcoords], :coordinate),
        basis.ksymops[1:n_kcoords],
        basis.kgrid,
        basis.kshift,
        basis.fft_size,
        basis.symmetries
    )
end

function Base.convert(::Type{PlaneWaveBasis{T}},
                      serial::PlaneWaveBasisSerialisation{T}) where {T}
    PlaneWaveBasis(serial.model, serial.Ecut, serial.kcoords,
                   serial.ksymops, serial.symmetries;
                   fft_size=serial.fft_size,
                   kgrid=serial.kgrid, kshift=serial.kshift,
                   variational=serial.variational)
end


text:
@testset "linalg" begin
    @testset "dot" begin
        @testset "Vector{$T}" for T in (Float64, ComplexF64)
            M = 3
            x, y = randn(T, M), randn(T, M)
            ẋ, ẏ = randn(T, M), randn(T, M)
            x̄, ȳ = randn(T, M), randn(T, M)
            frule_test(dot, (x, ẋ), (y, ẏ))
            rrule_test(dot, randn(T), (x, x̄), (y, ȳ))
        end
        @testset "Matrix{$T}" for T in (Float64, ComplexF64)
            M, N = 3, 4
            x, y = randn(T, M, N), randn(T, M, N)
            ẋ, ẏ = randn(T, M, N), randn(T, M, N)
            x̄, ȳ = randn(T, M, N), randn(T, M, N)
            frule_test(dot, (x, ẋ), (y, ẏ))
            rrule_test(dot, randn(T), (x, x̄), (y, ȳ))
        end
        @testset "Array{$T, 3}" for T in (Float64, ComplexF64)
            M, N, P = 3, 4, 5
            x, y = randn(T, M, N, P), randn(T, M, N, P)
            ẋ, ẏ = randn(T, M, N, P), randn(T, M, N, P)
            x̄, ȳ = randn(T, M, N, P), randn(T, M, N, P)
            frule_test(dot, (x, ẋ), (y, ẏ))
            rrule_test(dot, randn(T), (x, x̄), (y, ȳ))
        end
    end
    @testset "cross" begin
        @testset "frule" begin
            @testset "$T" for T in (Float64, ComplexF64)
                n = 3
                x, y = randn(T, n), randn(T, n)
                ẋ, ẏ = randn(T, n), randn(T, n)
                frule_test(cross, (x, ẋ), (y, ẏ))
            end
        end
        @testset "rrule" begin
            n = 3
            x, y = randn(n), randn(n)
            x̄, ȳ = randn(n), randn(n)
            ΔΩ = randn(n)
            rrule_test(cross, ΔΩ, (x, x̄), (y, ȳ))
        end
    end
    @testset "inv(::Matrix{$T})" for T in (Float64, ComplexF64)
        N = 3
        B = generate_well_conditioned_matrix(T, N)
        frule_test(inv, (B, randn(T, N, N)))
        rrule_test(inv, randn(T, N, N), (B, randn(T, N, N)))
    end
    @testset "pinv" begin
        @testset "$T" for T in (Float64, ComplexF64)
            test_scalar(pinv, randn(T))
            @test frule((Zero(), randn(T)), pinv, zero(T))[2] ≈ zero(T)
            @test rrule(pinv, zero(T))[2](randn(T))[2] ≈ zero(T)
        end
        @testset "Vector{$T}" for T in (Float64, ComplexF64)
            n = 3
            x, ẋ, x̄ = randn(T, n), randn(T, n), randn(T, n)
            tol, ṫol, t̄ol = 0.0, randn(), randn()
            Δy = copyto!(similar(pinv(x)), randn(T, n))
            frule_test(pinv, (x, ẋ), (tol, ṫol))
            @test frule((Zero(), ẋ), pinv, x)[2] isa typeof(pinv(x))
            rrule_test(pinv, Δy, (x, x̄), (tol, t̄ol))
            @test rrule(pinv, x)[2](Δy)[2] isa typeof(x)
        end
        @testset "$F{Vector{$T}}" for T in (Float64, ComplexF64), F in (Transpose, Adjoint)
            n = 3
            x, ẋ, x̄ = F(randn(T, n)), F(randn(T, n)), F(randn(T, n))
            y = pinv(x)
            Δy = copyto!(similar(y), randn(T, n))
            frule_test(pinv, (x, ẋ))
            y_fwd, ∂y_fwd = frule((Zero(),  ẋ), pinv, x)
            @test y_fwd isa typeof(y)
            @test ∂y_fwd isa typeof(y)
            rrule_test(pinv, Δy, (x, x̄))
            y_rev, back = rrule(pinv, x)
            @test y_rev isa typeof(y)
            @test back(Δy)[2] isa typeof(x)
        end
        @testset "Matrix{$T} with size ($m,$n)" for T in (Float64, ComplexF64),
            m in 1:3,
            n in 1:3

            X, Ẋ, X̄ = randn(T, m, n), randn(T, m, n), randn(T, m, n)
            ΔY = randn(T, size(pinv(X))...)
            frule_test(pinv, (X, Ẋ))
            rrule_test(pinv, ΔY, (X, X̄))
        end
    end
    @testset "$f" for f in (det, logdet)
        @testset "$f(::$T)" for T in (Float64, ComplexF64)
            b = (f === logdet && T <: Real) ? abs(randn(T)) : randn(T)
            test_scalar(f, b)
        end
        @testset "$f(::Matrix{$T})" for T in (Float64, ComplexF64)
            N = 3
            B = generate_well_conditioned_matrix(T, N)
            frule_test(f, (B, randn(T, N, N)))
            rrule_test(f, randn(T), (B, randn(T, N, N)))
        end
    end
    @testset "logabsdet(::Matrix{$T})" for T in (Float64, ComplexF64)
        N = 3
        B = randn(T, N, N)
        frule_test(logabsdet, (B, randn(T, N, N)))
        rrule_test(logabsdet, (randn(), randn(T)), (B, randn(T, N, N)))
        # test for opposite sign of determinant
        frule_test(logabsdet, (-B, randn(T, N, N)))
        rrule_test(logabsdet, (randn(), randn(T)), (-B, randn(T, N, N)))
    end
    @testset "tr" begin
        N = 4
        frule_test(tr, (randn(N, N), randn(N, N)))
        rrule_test(tr, randn(), (randn(N, N), randn(N, N)))
    end
    @testset "*" begin
        dims = [3,4,5]
        for n in dims, m in dims, p in dims
            n > 3 && n == m == p && continue  # don't need to test square case multiple times
            A = randn(m, n)
            B = randn(n, p)
            Ȳ = randn(m, p)
            rrule_test(*, Ȳ, (A, randn(m, n)), (B, randn(n, p)))
        end
    end
    @testset "$f" for f in [/, \]
        @testset "Matrix" begin
            for n in 3:5, m in 3:5
                A = randn(m, n)
                B = randn(m, n)
                Ȳ = randn(size(f(A, B)))
                rrule_test(f, Ȳ, (A, randn(m, n)), (B, randn(m, n)))
            end
        end
        @testset "Vector" begin
            x = randn(10)
            y = randn(10)
            ȳ = randn(size(f(x, y))...)
            rrule_test(f, ȳ, (x, randn(10)), (y, randn(10)))
        end
        if f == (/)
            @testset "$T on the RHS" for T in (Diagonal, UpperTriangular, LowerTriangular)
                RHS = T(randn(T == Diagonal ? 10 : (10, 10)))
                Y = randn(5, 10)
                Ȳ = randn(size(f(Y, RHS))...)
                rrule_test(f, Ȳ, (Y, randn(size(Y))), (RHS, randn(size(RHS))))
            end
        else
            @testset "$T on LHS" for T in (Diagonal, UpperTriangular, LowerTriangular)
                LHS = T(randn(T == Diagonal ? 10 : (10, 10)))
                y = randn(10)
                ȳ = randn(size(f(LHS, y))...)
                rrule_test(f, ȳ, (LHS, randn(size(LHS))), (y, randn(10)))
                Y = randn(10, 10)
                Ȳ = randn(10, 10)
                rrule_test(f, Ȳ, (LHS, randn(size(LHS))), (Y, randn(size(Y))))
            end
            @testset "Matrix $f Vector" begin
                X = randn(10, 4)
                y = randn(10)
                ȳ = randn(size(f(X, y))...)
                rrule_test(f, ȳ, (X, randn(size(X))), (y, randn(10)))
            end
            @testset "Vector $f Matrix" begin
                x = randn(10)
                Y = randn(10, 4)
                ȳ = randn(size(f(x, Y))...)
                rrule_test(f, ȳ, (x, randn(size(x))), (Y, randn(size(Y))))
            end
        end
    end
    @testset "norm" begin
        for dims in [(), (5,), (3, 2), (7, 3, 2)]
            A = randn(dims...)
            p = randn()
            ȳ = randn()
            rrule_test(norm, ȳ, (A, randn(dims...)), (p, randn()))
        end
    end
end


text:
struct ConstantLayer
    w
end

ConstantLayer(N::Integer) = ConstantLayer(zeros(N))

(m::ConstantLayer)(x) = m.w

Flux.@functor ConstantLayer

mutable struct A2TNetwork
  base::Chain
  attn::Chain
  solutions::Array{Any} # function takes in state and outputs vector of action values
  ignore_solution_computation::Bool
end

# By default, ignore the computation of the solutions
A2TNetwork(base, attn, solutions) = A2TNetwork(base, attn, solutions, true)

function (m::A2TNetwork)(input)
    b = m.base(input) #output is (Na, B)
    w = m.attn(input) #output is (Nt+1, B)
    Na, B, Nt  = size(b,1), size(b, 2), size(w,1) - 1

    if m.ignore_solution_computation
        qs = Array{Float32}(undef, Nt, Na, B)
        Zygote.ignore(() -> begin
            for i=1:B, s=1:Nt
                qs[s, :, i] .= m.solutions[s](input[:, i])
            end
        end)
    else
        qs = Flux.stack([m.solutions[s](input) for s in 1:Nt], 1)
    end
    sum(qs .* Flux.unsqueeze(w[1:Nt, :], 2), dims=1)[1,:,:] .+ w[Nt+1:Nt+1, :] .* b
end

Flux.@functor A2TNetwork

Flux.trainable(m::A2TNetwork) = (m.base, m.attn, (!m.ignore_solution_computation ? m.solutions : ())...)

function Base.iterate(m::A2TNetwork, i=1)
    layers = [m.base.layers..., m.attn.layers...]
    if !m.ignore_solution_computation
        for s in m.solutions
            for l in s
                push!(layers, l)
            end
        end
    end
    i > length(layers) && return nothing
    return (layers[i], i+1)
end

## A2T Network with state transform
mutable struct A2TSTNetwork
  base::Chain
  attn::Chain
  strans::Array{Chain}
  solutions::Array{Any} # function takes in state and outputs vector of action values
end

function (m::A2TSTNetwork)(input)
    b = m.base(input) #output is (Na, B)
    w = m.attn(input) #output is (Nt+1, B)
    Na, B, Nt  = size(b,1), size(b, 2), size(w,1) - 1
    qs = Zygote.Buffer(Array{Float32}(undef, Nt, Na, B))
    for s=1:Nt
        tinput = m.strans[s](input)
        for i=1:B
            qs[s, :, i] = m.solutions[s](tinput[:, i])
        end
    end
    sum(copy(qs) .* Flux.unsqueeze(w[1:Nt, :], 2), dims=1)[1,:,:] .+ w[Nt+1:Nt+1, :] .* b
end

Flux.@functor A2TSTNetwork

Flux.trainable(m::A2TSTNetwork) = (m.base, m.attn, m.strans...)

function Base.iterate(m::A2TSTNetwork, i=1)
    layers = [m.base.layers..., m.attn.layers...]
    for net in m.strans
        push!(layers, net.layers...)
    end
    i > length(layers) && return nothing
    return (layers[i], i+1)
end

## A2T network with fine tuning

mutable struct A2TFTNetwork
  base::Chain
  attn::Chain
  solutions::Array{Any} # function takes in state and outputs vector of action values
  finetune::Array{Chain} # fine tune network, one for each soluion
end

function (m::A2TFTNetwork)(input) # input is (Ni, B)
    Ni = size(input, 1)
    b = m.base(input) #output is (Na, B)
    w = m.attn(input) #output is (Nt+1, B)
    Na, B, Nt  = size(b,1), size(b, 2), size(w,1) - 1

    qs = Array{Float32}(undef, Nt, Na + Ni, B)
    Zygote.ignore(() -> begin
        for i=1:B, s=1:Nt
            qs[s, 1:Na, i] .= m.solutions[s](input[:, i])
            qs[s, Na+1:Na+Ni, i] .= input[:, i]
        end
    end)

    ft = Zygote.Buffer(Array{Float32}(undef, Nt, Na, B))
    for s=1:Nt
        ft[s, :, :] = m.finetune[s](qs[s, :, :])
    end

    sum(copy(ft) .* Flux.unsqueeze(w[1:Nt, :], 2), dims=1)[1,:,:] .+ w[Nt+1:Nt+1, :] .* b
end

Flux.@functor A2TFTNetwork

Flux.trainable(m::A2TFTNetwork) = (m.base, m.attn, m.finetune...)

function Base.iterate(m::A2TFTNetwork, i=1)
    layers = [m.base.layers..., m.attn.layers...]
    for net in m.finetune
        push!(layers, net.layers...)
    end
    i > length(layers) && return nothing
    return (layers[i], i+1)
end

## A Fine-tune network with choice of number of layers to freeze

mutable struct FTNetwork
  net::Chain
  train_layers::Array{Int}
end

(m::FTNetwork)(input) = m.net(input)
Flux.@functor FTNetwork

Flux.trainable(m::FTNetwork) = m.net.layers[m.train_layers]

function Base.iterate(m::FTNetwork, i=1)
    i > length(m.train_layers) && return nothing
    return (m.net.layers[m.train_layers[i]], i+1)
end


text:
cpu(x) = x
gpu(x) = x

cpu(x::CuArray) = adapt(Array, x)
gpu(x::Array) = CuArrays.cu(x)
cpu(x::StructArray) = replace_storage(Array, x)
gpu(x::StructArray) = replace_storage(x) do v
    typeof(v) <: StructArray ? gpu(v) :
    isbitstype(eltype(v)) ? gpu(v) : v
end

cpu(x::Neuron) = Neuron(cpu(x.synapses), x.soma)
gpu(x::Neuron) = Neuron(gpu(x.synapses), x.soma)

function cpu(x::STDP{T}) where T<:Real
    clastpre = cpu(x.lastpre)
    clastpost = cpu(x.lastpost)

    return STDP{T, typeof(clastpre)}(x.A₊, x.A₋, x.τ₊, x.τ₋, clastpre, clastpost)
end
function gpu(x::STDP{T}) where T<:Real
    glastpre = gpu(x.lastpre)
    glastpost = gpu(x.lastpost)

    return STDP{T, typeof(glastpre)}(x.A₊, x.A₋, x.τ₊, x.τ₋, glastpre, glastpost)
end

function cpu(x::Population)
    csynapses = cpu(x.synapses)
    csomas = cpu(x.somas)
    cweights = cpu(x.weights)
    clearner = cpu(x.learner)

    return Population(csomas, cweights, csynapses, clearner)
end

function gpu(x::Population)
    gsynapses = gpu(x.synapses)
    gsomas = gpu(x.somas)
    gweights = gpu(x.weights)
    glearner = gpu(x.learner)

    return Population(gsomas, gweights, gsynapses, glearner)
end

function cpu(x::InputPopulation)
    cinputs = cpu(x.inputs)

    return InputPopulation{typeof(cinputs)}(cinputs)
end

function gpu(x::InputPopulation)
    ginputs = gpu(x.inputs)

    return InputPopulation{typeof(ginputs)}(ginputs)
end

function cpu(x::Network)
    y = deepcopy(x)

    @inbounds for (k, v) in y.pops
        y.pops[k] = cpu(v)
    end

    @inbounds for (k, v) in y.connections
        y.pops[k] = NetworkEdge(cpu(v.weights), cpu(v.synapses), cpu(v.learner))
    end

    return y
end

function gpu(x::Network)
    y = deepcopy(x)

    @inbounds for (k, v) in y.pops
        y.pops[k] = gpu(v)
    end

    @inbounds for (k, v) in y.connections
        y.connections[k] = NetworkEdge(gpu(v.weights), gpu(v.synapses), gpu(v.learner))
    end

    return y
end


text:
"""
    conGrad(A,B,Xawal)
adalah fungsi yang digunakan untuk mencari solusi SPL `Ax=B` secara iteratif dengan
nilai tebakan awal `Xawal` menggunakan metode conjugate gradient.

# Example
```jl
julia> A = [2  3  3
        3  5  4
        3  4  6];

julia> B = [3,4,6];

julia> X,flag,err,M = conGrad(A,B,[0 0 0.9]);

julia> X
3-element Array{Float64,1}:
  2.380760976839332e-6
 -1.220333023839802e-6
  0.9999996142293243

julia> flag
0

julia> err
8.019009497939746e-8

julia> M
229×5 Array{Float64,2}:
   0.0  0.0          0.0          0.9       NaN
   1.0  0.0266764    0.0355685    0.953353    0.0694497
   2.0  0.0154769    0.000266424  0.975767    0.0432903
   3.0  0.0193622    0.00858077   0.976843    0.00924021
   ⋮
 225.0  5.10651e-6  -1.83459e-6   0.999999    1.32762e-6
 226.0  2.29335e-6  -1.29137e-6   1.0         3.08384e-6
 227.0  2.4522e-6   -1.23209e-6   1.0         1.73709e-7
 228.0  2.38076e-6  -1.22033e-6   1.0         8.01901e-8
```
return solusi `X` dengan `flag` bernilai 0 jika metode conjugate gradient
berhasil menemukan solusi dan gagal jika tidak 0.
Serta, matriks `M` yang berisi catatan proses tiap iterasi `[k, X', error]`
"""
function conGrad(A,B,Xawal)
    if ~(isposdef(A) && issymmetric(A))
        error("matriks A harus simetrik definit positif")
    end
    delta = 10^-7;
    maxi = 1000;
    flag = 1;
    X = Xawal[:];
    r = B-A*X
    d = r
    M = [0 X' NaN];
    for k = 1:maxi
        Xlama = X
        rlama = r
        dlama = d
        X = Xlama + ((rlama'rlama)/(dlama'A*dlama))*dlama
        r = rlama - ((rlama'rlama)/(dlama'A*dlama))*A*dlama
        d = r - ((r'r)/(rlama'rlama))*dlama
        err = norm(X-Xlama)
        M = [M; [k X' err] ];
        if err<delta
            flag = 0;
            break
        end
    end
    err = M[end,end]
    return X, flag, err, M
end


text: const DEFAULT_BRACKETING_nmax=100


text:
#!/usr/bin/env julia

using SimradEK60
using SimradEK60TestData
using EchogramPyPlot
using EchogramColorSchemes
using PyPlot
using LaTeXStrings

filename = EK60_SAMPLE
ps = pings(filename)
ps38 = [p for p in ps if p.frequency == 38000]
Sv38 = Sv(ps38)
r38 = R(ps38)

c = addwhite(EchogramColorSchemes.EK500).colors
echogram(Sv38, vmin=-95, vmax=-50, cmap=ColorMap("EK500", c, length(c)), range=maximum(r38))
cbar = PyPlot.colorbar()
cbar[:set_label](L" $S_v$ [dB re 1 $m^{-1}$] ")
xlabel("Pings")
ylabel("Range [m]")
savefig("example1.png")


text:
struct FeatureSuccessAndFailure
    feature::Feature
    n_success::UInt
    n_failure::UInt
end

mutable struct ResultAccumulator
    isaccumsuccess::Bool
    features::Vector{FeatureSuccessAndFailure}

    ResultAccumulator() = new(true, [])
end

function accumulateresult(acc::ResultAccumulator, result::FeatureResult)
    n_success::UInt = 0
    n_failure::UInt = 0
    for scenarioresult in result.scenarioresults
        arestepssuccessful = [issuccess(step) for step in scenarioresult.steps]
        isscenariosuccessful = all(arestepssuccessful)
        if isscenariosuccessful
            n_success += 1
        else
            n_failure += 1
        end

        acc.isaccumsuccess &= isscenariosuccessful
    end

    push!(acc.features, FeatureSuccessAndFailure(result.feature, n_success, n_failure))
end

issuccess(acc::ResultAccumulator) = acc.isaccumsuccess

function featureresults(accumulator::ResultAccumulator)
    accumulator.features
end


text:
using Documenter
using CausalForest

makedocs(
    sitename = "CausalForest.jl",
    format = Documenter.HTML(prettyurls = false),
    pages = [
        "Introduction" => "index.md"
    ]
)

deploydocs(
    repo = "github.com/BereniceAlexiaJocteur/CausalForest.jl.git",
    devbranch = "main"
)


text:
# DiagOp and FourierOp both <: AbstractLinearOp
# ======================================

# DiagOp has a Field as storage and represents
# the diagonal of a matrix  in the corresponding transformation basis
# determined by the Field type.
struct DiagOp{X<:Field} <: AbstractLinearOp
    f::X
end

# Interface methods for DiagOp
# ---------------------------------------

function Base.:inv(O::DiagOp{X}) where X<:Field
	DiagOp(X(fieldtransform(O.f), pinv.(fielddata(O.f))))
end

# This requires orthogonality of the transform
function LinearAlgebra.adjoint(O::DiagOp{X}) where X<:Field
	DiagOp(X(fieldtransform(O.f), conj.(fielddata(O.f))))
end

function Base.:*(O::DiagOp{X}, f::Y) where {F<:Transform, X<:Field{F}, Y<:Field{F}}
	Y(O.f * X(f))
end

function Base.:\(O::DiagOp{X}, f::Y) where {F<:Transform, X<:Field{F}, Y<:Field{F}}
	Y(O.f \ X(f))
end

_lmult(O::DiagOp{X}, f::Y) where {F<:Transform, X<:Field{F}, Y<:Field{F}} = O.f * X(f)

# getindex and basic operator functionality
# ------------------------------------------

LinearAlgebra.diag(O::DiagOp) = O.f

# indexing is propigated
Base.getindex(O::DiagOp, i) = getindex(O.f, i)

# Operations op(DiagOp), op(DiagOp, Number) and  op(Number, DiagOp)
# ------------------------------------------

# op(DiagOp, Number) and op(Number, DiagOp)
Base.:*(O::DiagOp, a::Number) = DiagOp(a * O.f)
Base.:*(a::Number, O::DiagOp) = DiagOp(a * O.f)

Base.:\(O::DiagOp, a::Number)  = DiagOp(O.f \ a)
Base.:\(a::Number, O::DiagOp)  = DiagOp(a \ O.f)

Base.:/(O::DiagOp, a::Number)  = DiagOp(O.f / a)
Base.:/(a::Number, O::DiagOp)  = DiagOp(a / O.f)

Base.:-(O::DiagOp)             = DiagOp(-O.f)

Base.:^(O::DiagOp, a::Number)  = DiagOp(O.f^a)
Base.:^(O::DiagOp, a::Integer) = DiagOp(O.f^a)

# op(DiagOp)
Base.:sqrt(O::DiagOp) = DiagOp(sqrt(O.f))

# Operations between DiagOp and UniformScaling
# ------------------------------------------

LinearAlgebra.:*(J::UniformScaling, O::DiagOp) = DiagOp(J.λ * O.f)
LinearAlgebra.:*(O::DiagOp, J::UniformScaling) = DiagOp(J.λ * O.f)
LinearAlgebra.:+(O::DiagOp, J::UniformScaling) = DiagOp(J.λ + O.f)
LinearAlgebra.:+(J::UniformScaling, O::DiagOp) = DiagOp(J.λ + O.f)
LinearAlgebra.:-(O::DiagOp, J::UniformScaling) = DiagOp(O.f - J.λ)
LinearAlgebra.:-(J::UniformScaling, O::DiagOp) = DiagOp(J.λ - O.f)

# Operations op(DiagOp, DiagOp)
# ------------------------------------------

# ops of the same type
Base.:+(O1::DiagOp{X}, O2::DiagOp{X}) where X<:Field = DiagOp(O1.f + O2.f)
Base.:-(O1::DiagOp{X}, O2::DiagOp{X}) where X<:Field = DiagOp(O1.f - O2.f)
Base.:*(O1::DiagOp{X}, O2::DiagOp{X}) where X<:Field = DiagOp(O1.f * O2.f)
Base.:\(O1::DiagOp{X}, O2::DiagOp{X}) where X<:Field = DiagOp(O1.f \ O2.f)
Base.:/(O1::DiagOp{X}, O2::DiagOp{X}) where X<:Field = DiagOp(O1.f / O2.f)


text:
@defcomp impactwaterresources begin
    regions = Index()

    watech = Variable(index=[time])
    watechrate = Parameter(default = 0.005)

    water = Variable(index=[time,regions])
    wrbm = Parameter(index=[regions])
    wrel = Parameter(default = 0.85)
    wrnl = Parameter(default = 1)
    wrpl = Parameter(default = 0.85)

    gdp90 = Parameter(index=[regions])
    income = Parameter(index=[time,regions])

    population = Parameter(index=[time,regions])
    pop90 = Parameter(index=[regions])

    temp = Parameter(index=[time,regions])

    function run_timestep(p, v, d, t)

        if gettime(t) > 2000
            v.watech[t] = (1.0 - p.watechrate)^(gettime(t) - 2000)
        else
            v.watech[t] = 1.0
        end

        for r in d.regions
            ypc = p.income[t, r] / p.population[t, r] * 1000.0
            ypc90 = p.gdp90[r] / p.pop90[r] * 1000.0

            water = p.wrbm[r] * p.gdp90[r] * v.watech[t] * (ypc / ypc90)^p.wrel * (p.population[t, r] / p.pop90[r])^p.wrpl * p.temp[t, r]^p.wrnl

            if water > 0.1 * p.income[t, r]
                v.water[t, r] = 0.1 * p.income[t, r]
            else
                v.water[t, r] = water
            end
        end
    end
end


text:
export NwtdirectionCG, NewtonCG
"""
Computes the Newton direction using the conjugate gradient iteravtive scheme
"""
function NwtdirectionCG(H, ∇f; verbose :: Bool = false)
    e = 1e-6
    n = length(∇f)
    τ = 0.5 # need parametrization
    cgtol = max(e, min(0.7, 0.01 * norm(∇f)^(1.0 + τ)))

    # cgTN doesn't work properly yet

    (d, cg_stats) = cg(H, -∇f,
                       atol=cgtol, rtol=0.0,
                       itmax=max(2 * n, 50),
                       verbose=verbose)

    return d
end

function NewtonCG(nlp          :: AbstractNLPModel,
                       nlp_stop     :: NLPStopping;
                       verbose      :: Bool=false,
                       verboseLS    :: Bool = false,
                       kwargs...)
    return  Newton(nlp,
                   nlp_stop;
                   verbose = verbose,
                   verboseLS = verboseLS,
                   Nwtdirection = NwtdirectionCG,
                   hessian_rep = hessian_operator,
                   kwargs...)
end


text:
function iminimal()
	p = Param()
	gammas = 1.0:0.1:3.0
	betas  = 0.5:0.05:1.0
	Rs  = 1.0:0.05:1.5
	sigmas = 0.0:0.05:1.0

	mp = @manipulate for γ in slider(gammas, label = "γ", value =p.gamma ),
						 β in slider(betas, label = "β", value =p.beta) ,
						 R in slider(Rs, label = "R", value =p.R) ,
						 σ in slider(sigmas, label = "σ", value =p.sigma)

		 p =Param(par = Dict(:beta => β, :sigma => σ,  :R => R))
		 minimal_EGM(p)

	end
	@layout! mp vbox(
		hbox(β, R, σ),
		observe(_))
	w = Window()
	body!(w,mp)
end

function iminimalb()
	p = Param()
	gammas = 1.0:0.1:3.0
	betas  = 0.5:0.05:1.0
	Rs  = 1.0:0.05:1.5
	sigmas = 0.0:0.05:1.0
	nus = 0.0:0.01:10
	bbars = 0.0:0.1:10

	mp = @manipulate for γ in slider(gammas, label = "γ", value =p.gamma ),
						 β in slider(betas, label = "β", value =p.beta) ,
						 R in slider(Rs, label = "R", value =p.R) ,
						 σ in slider(sigmas, label = "σ", value =p.sigma),
						 bbar in slider(bbars, label = "bbar", value =p.bbar),
						 nu in slider(nus, label = "ν", value =5.0)

		 p =Param(par = Dict(:beta => β, :sigma => σ,  :R => R , :bbar => bbar, :ν => nu))
		 minimal_EGM_bequest(p)

	end
	@layout! mp vbox(
		hbox(β, R, σ),
		hbox(bbar, nu),
		observe(_))
	w = Window()
	body!(w,mp)
end

# DCEGM.interact(DCEGM.rsgp) for sim
# DCEGM.interact(DCEGM.runfp) for fedors model
function ifedor()
	p = Param()
	gammas = 1.0:0.1:3.0
	betas  = 0.5:0.05:1.0
	Rs  = 1.0:0.05:1.5
	alphas = 0.0:0.05:0.5
	sigmas = 0.0:0.05:0.35
	lambdas = 0.0000002:0.05:1
	rhos = 0.1:0.05:1
	deltas = 0.0:0.05:0.5
	pensions = 0.0:0.1:1.0
	nus = 0.0:0.1:10
	bbars = 0.0:0.1:10

	mp = @manipulate for dosim = Dict("sim" => true, "sol" => false),
						 id = Dict("id=$id" => id for id in 1:2) ,
						 nsims = spinbox(label="nsims"; value=p.nsims) |> onchange,
						 nT = spinbox(label="nperiods"; value=p.nT) |> onchange,
						 # γ in slider(gammas, label = "γ", value =p.gamma ),
						 β in slider(betas, label = "β", value =p.beta) |> onchange ,
						 R in slider(Rs, label = "R", value =p.R) |> onchange,
						 α in slider(alphas, label = "α", value =p.alpha) |> onchange ,
						 σ in slider(sigmas, label = "σ", value =p.sigma) |> onchange ,
						 λ in slider(lambdas, label = "λ", value =p.lambda) |> onchange,
						 δ in slider(deltas, label= "δ",value = p.delta) |> onchange,
						 pens in slider(pensions, label= "pension",value = p.pension) |> onchange,
						 bbar in slider(bbars, label = "bbar", value =p.bbar) |> onchange,
						 nu in slider(nus, label = "ν", value =p.ν) |> onchange

		 m,p = runf(par = Dict(:nsims => nsims , :bbar => bbar, :ν => nu, :beta => β, :alpha => α, :sigma => σ, :lambda => λ, :R => R, :delta => δ, :pension => pens, :nT => nT))
		 if dosim
			 s = sim(m,p)
			 plot_s(s)
		 else
			 plot(m,p,id = id)
		 end
	end
	@layout! mp vbox(
		hbox(dosim, id, nsims),
		hbox(β, R, σ, nT),
		hbox(α, λ, nu),
		hbox(δ, pens, bbar),
		observe(_))
	w = Window()
	body!(w,mp)
end

function igmodel()
	p = Param()
	gammas = 1.0:0.1:3.0
	betas  = 0.5:0.05:1.0
	Rs  = 1.0:0.05:1.5
	alphas = 0.0:0.05:0.5
	sigmas = 0.0:0.05:0.35
	lambdas = 0.0000002:0.05:1
	rhos = 0.1:0.05:1
	deltas = 0.0:0.05:0.5
	pensions = 0.0:0.1:1.0
	nus = 0.0:0.1:10
	bbars = 0.0:0.1:10

	try
		mp = @manipulate for dosim = Dict("sim" => true, "sol" => false),
							 id = Dict("id=$id" => id for id in 1:2),
							 nsims = spinbox(label="nsims"; value=20) |> onchange,
							 # γ in slider(gammas, label = "γ", value =p.gamma ),

							 β in slider(betas, label = "β", value =p.beta) |> onchange ,
							 R in slider(Rs, label = "R", value =p.R)  |> onchange,
							 α in slider(alphas, label = "α", value =p.alpha) |> onchange ,
							 σ in slider(sigmas, label = "σ", value =p.sigma) |> onchange ,
							 λ in slider(lambdas, label = "λ", value =p.lambda) |> onchange,
							 ρ in slider(0:0.05:1, label = "ρ", value =p.ρ) |> onchange,
							 δ in slider(deltas, label= "δ",value = p.delta) |> onchange,
							 pens in slider(pensions, label= "pension",value = p.pension) |> onchange,
							 bbar in slider(bbars, label = "bbar", value =p.bbar) |> onchange,
							 nu in slider(nus, label = "ν", value =p.ν) |> onchange

			 m,p = rung(par = Dict(:nsims => nsims , :bbar => bbar, :ν => nu, :ρ => ρ, :beta => β, :alpha => α, :sigma => σ, :lambda => λ, :R => R, :delta => δ, :pension => pens))
			 if dosim
				 s = sim(m,p)
				 plot_s(s)
			 else
				 plot(m,p,id = id,iy = 1)
			 end
		end
		@layout! mp vbox(
			hbox(dosim, id, nsims),
			hbox(β, R, σ),
			hbox(α, λ, nu, ρ),
			hbox(δ, pens, bbar),
			observe(_))
		w = Window()
		body!(w,mp)
	catch e
		return e
	end

end

function ibkmodel(;it::Bool=false)
	p = Param(par = Dict(:nT => 20, :retage => 10) )
	gammas = 1.0:0.1:3.0
	betas  = 0.5:0.05:1.0
	Rs  = 1.0:0.05:1.5
	alphas = [0.0,0.1,0.35,0.5,1.0]
	alphaT = -1000:1:0.0
	lambdas = 0.2:0.1:1.0
	rhos = 0.1:0.05:1
	times = 1:p.nT

	if it
		mp = @manipulate for id = Dict("id=$id" => id for id in 1:2),
							 ti in slider(times, label = "period", value = 1),
			                 γ in slider(gammas, label = "γ", value =p.gamma ),
							 β in slider(betas, label = "β", value =p.beta) ,
							 a in slider(alphas, label = "α", value = p.alpha) ,
							 aT in slider(alphaT, label = "αT", value = p.alphaT) ,
							 alowT in slider(-5:0.5:0.0, label = "alow") ,
							 λ in slider(lambdas, label = "λ", value =p.lambda)#,
							 # ρ in slider(rhos, label = "ρ", value =p.ρ)

			m,p = runbk(par = Dict(:nT => 50,:a_low => -5.0,:a_lowT => alowT,:na =>501,:beta => β, :alphaT => aT, :alpha => a, :gamma => γ , :lambda => λ))
			plot(m,p,it = ti, id = id,ylims = (-5,13),size = (700,400))
		end
		w = Window()
		body!(w,mp)

	else
		mp = @manipulate for iy = OrderedDict("iy=$iy" => iy for iy in 1:p.ny) ,
			                 id = Dict("id=$id" => id for id in 1:2),
			                 γ in slider(gammas, label = "γ", value =p.gamma ),
							 β in slider(betas, label = "β", value =p.beta) ,
							 a in slider(alphas, label = "α", value = p.alpha) ,
							 aT in slider(alphaT, label = "αT", value = p.alphaT) ,
							 alow in slider(-5:0.5:0.0, label = "alow") ,
							 λ in slider(lambdas, label = "λ", value =p.lambda)#,
							 # ρ in slider(rhos, label = "ρ", value =p.ρ)

			m,p = runbk(par = Dict(:nT => 50,:a_low => -5.0,:a_lowT => alow,:na =>501,:beta => β, :alphaT => aT, :alpha => a, :gamma => γ , :lambda => λ))
			plot(m,p,iy = iy, id = id,ylims = (-5,13),size = (700,400))
		end
		w = Window()
		body!(w,mp)

	end

end

function ibksim(;pars = Dict(:nT => 50),it::Bool=false)
	p = Param(par = pars )
	gammas = 1.0:0.1:3.0
	betas  = 0.5:0.05:1.0
	Rs  = 1.0:0.05:1.5
	alphas = 0.0:0.05:1.0
	alphaT = -10:0.5:0.0
	lambdas = 0.02:0.1:1.0
	rhos = 0.1:0.05:1
	times = 1:p.nT

	if it
		mp = @manipulate for dosim = Dict("sim" => true, "sol" => false),
			id = Dict("id=$id" => id for id in 1:2),
			nsims = spinbox(label="nsims"; value=20) |> onchange,
			ti in slider(times, label = "period", value = 1) |> onchange,
			γ in slider(gammas, label = "γ", value =p.gamma ) |> onchange,
			β in slider(betas, label = "β", value =p.beta)  |> onchange,
			a in slider(alphas, label = "α", value = p.alpha)  |> onchange,
			aT in slider(alphaT, label = "αT", value = p.alphaT)  |> onchange,
			alowT in slider(-5:0.5:0.0, label = "alowT", value = p.a_lowT)  |> onchange,
			λ in slider(lambdas, label = "λ", value =p.lambda)  |> onchange
			# ρ in slider(rhos, label = "ρ", value =p.ρ)
			pp = Dict(:nsims => nsims,:nT => 70, :retage => 55, :pension => 0.5,:a_low => -5.0,:a_lowT => alowT:beta => β, :alphaT => aT, :alpha => a, :gamma => γ , :lambda => λ)
			m,p = runbk(par = pp)
			if dosim
				s = sim(m,p)
				plot_s(s)
			else
				plot(m,p,it = ti, id = id,ylims = (-5,13),size = (700,400))
			end

		end
		w = Window()
		body!(w,mp)

	else
		mp = @manipulate for dosim = Dict("sim" => true, "sol" => false),
			iy = OrderedDict("iy=$iy" => iy for iy in 1:p.ny) ,
			nsims = spinbox(label="nsims"; value=20) |> onchange,
			id = Dict("id=$id" => id for id in 1:2),
			γ in slider(gammas, label = "γ", value =p.gamma )|> onchange,
			β in slider(betas, label = "β", value =p.beta) |> onchange,
			a in slider(alphas, label = "α", value = p.alpha) |> onchange,
			aT in slider(alphaT, label = "αT", value = p.alphaT) |> onchange,
			alowT in slider(-5:0.5:0.0, label = "alowT", value = p.a_lowT) |> onchange,
			λ in slider(lambdas, label = "λ", value =p.lambda) |> onchange
			# ρ in slider(rhos, label = "ρ", value =p.ρ)

			pp = Dict(:nsims => nsims,:nT => 70, :retage => 55, :pension => 0.5, :a_low => -5.0,:a_lowT => alowT,:beta => β, :alphaT => aT, :alpha => a, :gamma => γ , :lambda => λ)
 			m,p = runbk(par = pp)
			if dosim
 				s = sim(m,p)
 				plot_s(s)
 			else
 				plot(m,p,iy = iy, id = id,ylims = (-5,13),size = (700,400))
 			end
		end
		w = Window()
		body!(w,mp)
	end
end


text:
module GeometricIntegratorsDiffEq

using Reexport
@reexport using DiffEqBase

using GeometricIntegrators

const warnkeywords =
    (:save_idxs, :d_discontinuities, :unstable_check, :save_everystep,
     :save_end, :initialize_save, :adaptive, :abstol, :reltol, :dtmax,
     :dtmin, :force_dtmin, :internalnorm, :gamma, :beta1, :beta2,
     :qmax, :qmin, :qsteady_min, :qsteady_max, :qoldinit, :failfactor,
     :maxiters, :isoutofdomain, :unstable_check,
     :calck, :progress, :tstops, :saveat, :dense)

function __init__()
    global warnlist = Set(warnkeywords)
end

include("algorithms.jl")
include("solve.jl")

export GeometricIntegratorAlgorithm, GIEuler, GIMidpoint, GIHeun, GIKutta,
       GIKutta, GIERK4, GIERK438, GIImplicitEuler, GIImplicitMidpoint,
       GIRadIIA2, GIRadIIA3, GISRK3, GIGLRK, GISymplecticEulerA, GISymplecticEulerB,
       GILobattoIIIAIIIB2, GILobattoIIIBIIIA2

end # module


text:
#==
All differentials need to define + and *.
That happens here.

We just use @eval to define all the combinations for AbstractDifferential
subtypes, as we know the full set that might be encountered.
Thus we can avoid any ambiguities.

Notice:
    The precidence goes: (:Wirtinger, :Zero, :DoesNotExist, :One, :AbstractThunk, :Any)
    Thus each of the @eval loops creating definitions of + and *
    defines the combination this type with all types of  lower precidence.
    This means each eval loops is 1 item smaller than the previous.
==#

function Base.:*(a::Wirtinger, b::Wirtinger)
    error("""
          Cannot multiply two Wirtinger objects; this error likely means a
          `WirtingerRule` was inappropriately defined somewhere. Multiplication
          of two Wirtinger objects is not defined because chain rule application
          often expands into a non-commutative operation in the Wirtinger
          calculus. To put it another way: simply given two Wirtinger objects
          and no other information, we can't know "locally" which components to
          conjugate in order to implement the chain rule. We could pick a
          convention; for example, we could define `a::Wirtinger * b::Wirtinger`
          such that we assume the chain rule application is of the form `f_a ∘ f_b`
          instead of `f_b ∘ f_a`. However, picking such a convention is likely to
          lead to silently incorrect derivatives due to commutativity assumptions
          in downstream generic code that deals with the reals. Thus, ChainRulesCore
          makes this operation an error instead.
          """)
end

function Base.:+(a::Wirtinger, b::Wirtinger)
    return Wirtinger(+(a.primal, b.primal), a.conjugate + b.conjugate)
end

for T in (:Zero, :DoesNotExist, :One, :AbstractThunk, :Any)
    @eval Base.:+(a::Wirtinger, b::$T) = a + Wirtinger(b, Zero())
    @eval Base.:+(a::$T, b::Wirtinger) = Wirtinger(a, Zero()) + b

    @eval Base.:*(a::Wirtinger, b::$T) = Wirtinger(a.primal * b, a.conjugate * b)
    @eval Base.:*(a::$T, b::Wirtinger) = Wirtinger(a * b.primal, a * b.conjugate)
end

Base.:+(::Zero, b::Zero) = Zero()
Base.:*(::Zero, ::Zero) = Zero()
for T in (:DoesNotExist, :One, :AbstractThunk, :Any)
    @eval Base.:+(::Zero, b::$T) = b
    @eval Base.:+(a::$T, ::Zero) = a

    @eval Base.:*(::Zero, ::$T) = Zero()
    @eval Base.:*(::$T, ::Zero) = Zero()
end

Base.:+(::DoesNotExist, ::DoesNotExist) = DoesNotExist()
Base.:*(::DoesNotExist, ::DoesNotExist) = DoesNotExist()
for T in (:One, :AbstractThunk, :Any)
    @eval Base.:+(::DoesNotExist, b::$T) = b
    @eval Base.:+(a::$T, ::DoesNotExist) = a

    @eval Base.:*(::DoesNotExist, ::$T) = DoesNotExist()
    @eval Base.:*(::$T, ::DoesNotExist) = DoesNotExist()
end

Base.:+(a::One, b::One) = extern(a) + extern(b)
Base.:*(::One, ::One) = One()
for T in (:AbstractThunk, :Any)
    @eval Base.:+(a::One, b::$T) = extern(a) + b
    @eval Base.:+(a::$T, b::One) = a + extern(b)

    @eval Base.:*(::One, b::$T) = b
    @eval Base.:*(a::$T, ::One) = a
end

Base.:+(a::AbstractThunk, b::AbstractThunk) = extern(a) + extern(b)
Base.:*(a::AbstractThunk, b::AbstractThunk) = extern(a) * extern(b)
for T in (:Any,)
    @eval Base.:+(a::AbstractThunk, b::$T) = extern(a) + b
    @eval Base.:+(a::$T, b::AbstractThunk) = a + extern(b)

    @eval Base.:*(a::AbstractThunk, b::$T) = extern(a) * b
    @eval Base.:*(a::$T, b::AbstractThunk) = a * extern(b)
end


text:
#=
Constructors for Str strings

Copyright 2017-2018 Gandalf Software, Inc., Scott P. Jones
Licensed under MIT License, see LICENSE.md
=#

function _encode_ascii_latin(str, len)
    buf, out = _allocate(UInt8, len)
    @inbounds for ch in str
        set_codeunit!(out, ch%UInt8)
        out += 1
    end
    buf
end

function _encode_ascii_latin(pnt::Ptr{UInt8}, len)
    buf, out = _allocate(UInt8, len)
    fin = out + len
    @inbounds while out < fin
        ch8 = get_codeunit(pnt)
        set_codeunit!(out,
                      (ch8 <= 0x7f ? ch8 : (((ch8 & 3) << 6) | (get_codeunit(pnt += 1) & 0x3f))))
        pnt += 1
        out += 1
    end
    buf
end

@inline function safe_copy(::Type{Vector{T}}, ::Type{C}, str) where {T<:CodeUnitTypes,C<:CSE}
    @preserve str buf begin
        len = ncodeunits(str)
        buf, out = _allocate(T, len)
        _memcpy(out, pointer(str), len)
        Str(C, buf)
    end
end

@inline safe_copy(::Type{<:Str}, ::Type{C}, str) where {C<:CSE} = Str(C, str.data)
@inline safe_copy(::Type{String}, ::Type{C}, str) where {C<:CSE} = Str(C, _data(str))

function _str(str::AbstractString)
    # handle zero length string quickly
    isempty(str) && return empty_ascii
    len, flags, num4byte, num3byte, num2byte, latin1byte = unsafe_check_string(str)
    if flags == 0
        Str(ASCIICSE, _encode_ascii_latin(str, len))
    elseif num4byte != 0
        Str(_UTF32CSE, _encode_utf32(str, len))
    elseif num3byte + num2byte != 0
        Str(_UCS2CSE, _encode_utf16(str, len))
    else
        Str(latin1byte == 0 ? ASCIICSE : _LatinCSE, _encode_ascii_latin(str, len))
    end
end

function _str(str::T) where {T<:Union{Vector{UInt8}, Str{<:Binary_CSEs}, String}}
    # handle zero length string quickly
    (siz = sizeof(str)) == 0 && return empty_ascii
    @preserve str begin
        pnt = pointer(str)
        len, flags, num4byte, num3byte, num2byte, latin1byte = fast_check_string(pnt, siz)
        if flags == 0
            buf, out = _allocate(UInt8, len)
            _memcpy(out, pnt, len)
            Str(ASCIICSE, buf)
        elseif num4byte != 0
            Str(_UTF32CSE, _encode_utf32(pnt, len))
        elseif num3byte + num2byte != 0
            Str(_UCS2CSE, _encode_utf16(pnt, len))
        else
            Str(latin1byte == 0 ? ASCIICSE : _LatinCSE, _encode_ascii_latin(pnt, len))
        end
    end
end

function _str_cpy(::Type{T}, str, len) where {T}
    @preserve str begin
        buf, pnt = _allocate(T, len)
        @inbounds for ch in str
            set_codeunit!(pnt, ch%T)
            pnt += sizeof(T)
        end
        buf
    end
end

convert(::Type{Str}, str::AbstractString) = _str(str)
convert(::Type{Str}, str::String)         = _str(str)
convert(::Type{Str}, str::Str)            = str

convert(::Type{<:Str{C}}, str::AbstractString) where {C} = convert(Str{C}, _str(str))
convert(::Type{<:Str{C}}, str::Str{C}) where {C} = str

convert(::Type{UniStr}, str::AbstractString) = _str(str)
convert(::Type{UniStr}, str::String)         = _str(str)
convert(::Type{UniStr}, str::Str{<:Union{ASCIICSE,SubSet_CSEs}}) = str

convert(::Type{<:Str{C}}, vec::AbstractVector{UInt8}) where {C<:Union{BinaryCSE,Text1CSE}} =
    Str(C, _str_cpy(UInt8, vec, length(vec)))
convert(::Type{<:Str{Text2CSE}}, vec::AbstractVector{UInt16}) =
    Str(Text2CSE, _str_cpy(UInt16, vec, length(vec)))
convert(::Type{<:Str{Text4CSE}}, vec::AbstractVector{UInt32}) =
    Str(Text4CSE, _str_cpy(UInt32, vec, length(vec)))

(::Type{Str})(str::AbstractString) = _str(str)
(::Type{Str})(str::String)         = _str(str)
(::Type{Str})(str::Str)            = str

(::Type{UniStr})(str::AbstractString) = _str(str)
(::Type{UniStr})(str::String)         = _str(str)
(::Type{UniStr})(str::Str{<:Union{ASCIICSE,SubSet_CSEs}}) = str

function convert(::Type{UniStr}, str::T) where {T<:Str}
    # handle zero length string quickly
    is_empty(str) && return empty_ascii
    len, flags = count_chars(T, pointer(str), ncodeunits(str))
    if flags == 0
        Str(ASCIICSE, codeunit(T) == UInt8 ? str.data : _str_cpy(UInt8, str, len))
    elseif (flags & ~UTF_LATIN1) == 0
        Str(_LatinCSE, codeunit(T) == UInt8 ? str.data : _str_cpy(UInt8, str, len))
    elseif (flags & UTF_UNICODE4) == 0
        Str(_UCS2CSE, codeunit(T) == UInt16 ? str.data : _str_cpy(UInt16, str, len))
    else
        Str(_UTF32CSE, codeunit(T) == UInt32 ? str.data : _str_cpy(UInt32, str, len))
    end
end

#convert(::Type{<:Str{C}}, str::String) where {C} = convert(C, _str(str))

convert(::Type{<:Str{Text1CSE}}, str::String) = Str(Text1CSE, str)
convert(::Type{<:Str{BinaryCSE}}, str::String) = Str(BinaryCSE, str)

convert(::Type{<:Str{LatinCSE}}, str::Str{_LatinCSE}) = Str(LatinCSE, str.data)
convert(::Type{<:Str{UCS2CSE}},  str::Str{_UCS2CSE})  = Str(UCS2CSE,  str.data)
convert(::Type{<:Str{UTF32CSE}}, str::Str{_UTF32CSE}) = Str(UTF32CSE, str.data)

convert(::Type{<:Str{LatinCSE}}, str::Str{ASCIICSE}) = Str(LatinCSE, str.data)

convert(::Type{String}, str::Str{<:Union{ASCIICSE,Text1CSE,BinaryCSE}}) = str.data

function convert(::Type{String}, vec::AbstractVector{<:Chr})
    out = get_iobuffer(sizeof(vec))
    @inbounds for ch in vec
        print(out, ch)
    end
    String(take!(out))
end

function convert(::Type{String}, vec::AbstractVector{<:Union{Text1Chr,ASCIIChr}})
    @preserve buf begin
        buf, pnt = _allocate(UInt8, length(vec))
        @inbounds for byt in vec
            set_codeunit!(pnt, byt)
            pnt += 1
        end
        buf
    end
end

function Str(vec::AbstractVector{T}) where {CS,BT,T<:Chr{CS,BT}}
    C = codepoint_cse(T)
    buf, pnt = _allocate(BT, length(vec))
    @inbounds for ch in vec
        set_codeunit!(pnt, ch%BT)
        pnt += sizeof(BT)
    end
    Str(C, buf)
end

"""Convert to a UniStr if valid Unicode, otherwise return a Text1Str"""
function unsafe_str(str::Union{Vector{UInt8}, T, SubString{T}};
                    accept_long_null  = false,
                    accept_surrogates = false,
                    accept_long_char  = false,
                    accept_invalids   = true
                    ) where {T <: Union{BinaryStr, Text1Str, String}}
    # handle zero length string quickly
    (siz = sizeof(str)) == 0 && return empty_ascii
    @preserve str begin
        pnt = pointer(str)
        len, flags, num4byte, num3byte, num2byte, latin1byte, invalids =
            unsafe_check_string(pnt, 1, siz;
                                accept_long_null  = accept_long_null,
                                accept_surrogates = accept_surrogates,
                                accept_long_char  = accept_long_char,
                                accept_invalids   = accept_invalids)
        if flags == 0
            # Don't allow this to be aliased to a mutable Vector{UInt8}
            safe_copy(T, ASCIICSE, str)
        elseif invalids != 0
            safe_copy(T, Text1CSE, str)
        elseif num4byte != 0
            Str(_UTF32CSE, _encode_utf32(pnt, len))
        elseif num2byte + num3byte != 0
            Str(_UCS2CSE, _encode_utf16(pnt, len))
        else
            Str(latin1byte == 0 ? ASCIICSE : _LatinCSE, _encode_ascii_latin(pointer(str), len))
        end
    end
end

"""Convert to a UniStr if valid Unicode, otherwise return a Text1Str/Text2Str/Text4Str"""
function unsafe_str(str::T;
                    accept_long_null  = false,
                    accept_surrogates = false,
                    accept_long_char  = false,
                    accept_invalids   = true
                    ) where {T<:Union{AbstractString,SubString{<:AbstractString},
                                      AbstractVector{<:Union{AbsChar,CodeUnitTypes}}}}
    # handle zero length string quickly
    is_empty(str) && return empty_ascii
    len, flags, num4byte, num3byte, num2byte, latin1byte, invalids =
        unsafe_check_string(str;
                            accept_long_null  = accept_long_null,
                            accept_surrogates = accept_surrogates,
                            accept_long_char  = accept_long_char,
                            accept_invalids   = accept_invalids)
    if flags == 0
        Str(ASCIICSE, _cvtsize(UInt8, str, len))
    elseif invalids != 0
        # Todo: Make sure this handles different sorts of SubStrings effectively
        siz = sizeof(eltype(T))
        C = siz == 4 ? Text4CSE : (siz == 2 ? Text2CSE : Text1CSE)
        S = codeunit(C)
        buf, pnt = _allocate(S, len)
        @inbounds for ch in str
            set_codeunit!(pnt, ch%S)
            pnt += siz
        end
        Str(C, buf)
    elseif num4byte != 0
        Str(_UTF32CSE, _encode_utf32(str, len))
    elseif num2byte + num3byte != 0
        Str(_UCS2CSE, _encode_utf16(str, len))
    else
        Str(latin1byte == 0 ? ASCIICSE : _LatinCSE, _encode_ascii_latin(str, len))
    end
end

# Fallback constructors for Str types, from any AbstractString
(::Type{T})(vec::S) where {T<:Str, S<:AbstractVector} = convert(T, vec)
(::Type{T})(str::S) where {T<:Str, S<:AbstractString} = convert(T, str)
(::Type{T})(str::S) where {T<:Str, S<:Str} = convert(T, str)
(::Type{T})(str::T) where {T<:Str} = str


text:
# https://adventofcode.com/2021/day/10
using AdventOfCode

input = readlines("2021/data/day_10.txt")
isopen(char::AbstractChar)=char in collect("([{<")
isclose(char::AbstractChar)=char in collect(")]}>")
corruptscore(char::AbstractChar)=Dict(collect(")]}>").=>(3,57,1197,25137))[char]
function readcorruptedline(line)
    memory=Char[]
    for po in 1:length(line)
        hp=line[po]
        if isclose(hp)
            if abs(hp-memory[end])>2 return(corruptscore(hp)) end
            pop!(memory)
        else
            push!(memory,hp)
        end
    end
    0
end
function part_1(input)
    sum(readcorruptedline, input)
end

function readincompleteline(line)
    memory=Char[]
    for po in 1:length(line)
        hp=line[po]
        if isclose(hp)
            if abs(hp-memory[end])>2 return(0) end
            pop!(memory)
        else
            push!(memory,hp)
        end
    end
    incompletescore(memory)
end
openscore(char::AbstractChar)=Dict(collect("([{<").=>(1:4))[char]
function incompletescore(openstr)
    if length(openstr)==0 return(0) end
    openscore(openstr[begin])+5incompletescore(openstr[begin+1:end])
end

function part_2(input)
    median(filter(>(0),readincompleteline.(input)))|>Int
end