text:
"""
    MRP{T} <: Rotation

Modified Rodrigues Parameter. Is a 3D parameterization of attitude, and is a sterographic
projection of the 4D unit sphere onto the plane tangent to the negative real pole. They
have a singularity at θ = ±360°.

# Constructors
MRP(x, y, z)
MRP(r::AbstractVector)
"""
struct MRP{T} <: Rotation{3,T}
    x::T
    y::T
    z::T
    @inline MRP{T}(x,y,z) where T = new{T}(x,y,z)
    @inline MRP{T}(p::MRP) where T = new{T}(p.x, p.y, p.z)
end

# ~~~~~~~~~~~~~~~ Constructors ~~~~~~~~~~~~~~~ #
MRP(x::X, y::Y, z::Z) where {X,Y,Z} = MRP{promote_type(X,Y,Z)}(x, y, z)
(::Type{M})(g::StaticVector) where M<:MRP= M(g[1], g[2], g[3])

# ~~~~~~~~~~~~~~~ Conversions ~~~~~~~~~~~~~~~ #
params(g::MRP) = SVector{3}(g.x, g.y, g.z)

# ~~~~~~~~~~~~~~~ Initializers ~~~~~~~~~~~~~~~ #
function Random.rand(rng::AbstractRNG, ::Random.SamplerType{RP}) where RP <: MRP
    RP(rand(rng, UnitQuaternion))
end
Base.one(::Type{RP}) where RP <: MRP = RP(0.0, 0.0, 0.0)

# ~~~~~~~~~~~~~~~~ Quaternion <=> MRP ~~~~~~~~~~~~~~~~~~ #
@inline function (::Type{Q})(p::MRP) where Q <: UnitQuaternion
    n2 = p.x * p.x + p.y * p.y + p.z * p.z
    M = 2/(1+n2)
    q = Q((1-n2)/(1+n2), M*p.x, M*p.y, M*p.z, false)
end

@inline function (::Type{RP})(q::UnitQuaternion) where RP<:MRP
    M = 1/(1+q.w)
    RP(q.x*M, q.y*M, q.z*M)
end

# ~~~~~~~~~~~~~~~ StaticArrays Interface ~~~~~~~~~~~~~~~ #
@inline (::Type{RP})(t::NTuple{9}) where RP<:MRP = convert(RP, UnitQuaternion(t))
@inline Base.getindex(p::MRP, i::Int) = convert(UnitQuaternion, p)[i]
@inline Base.Tuple(p::MRP) = Tuple(convert(UnitQuaternion, p))

# ~~~~~~~~~~~~~~~ Math Operations ~~~~~~~~~~~~~~~ #
Base.inv(p::MRP) = MRP(-p.x, -p.y, -p.z)

# ~~~~~~~~~~~~~~~ Composition ~~~~~~~~~~~~~~~ #
function (*)(p2::MRP, p1::MRP)
    p2, p1 = params(p2), params(p1)
    MRP(((1-p2'p2)*p1 + (1-p1'p1)*p2 - cross(2p1, p2) ) / (1+p1'p1*p2'p2 - 2p1'p2))
end

function (\)(p1::MRP, p2::MRP)
    q1,q2 = params(p1), params(p2)
    n1,n2 = q1'q1, q2'q2
    θ = 1/((1+n1)*(1+n2))
    s1,s2 = (1-n1), (1-n2)
    v1 = -2q1
    v2 =  2q2
    s = s1*s2 - v1'v2
    v = s1*v2 + s2*v1 + v1 × v2

    M = θ/(1+θ*s)
    return MRP(v*M)
end

function (/)(p1::MRP, p2::MRP)
    q1,q2 = params(p1), params(p2)
    n1,n2 = q1'q1, q2'q2
    θ = 1/((1+n1)*(1+n2))
    s1,s2 = (1-n1), (1-n2)
    v1 =  2q1
    v2 = -2q2
    s = s1*s2 - v1'v2
    v = s1*v2 + s2*v1 + v1 × v2

    M = θ/(1+θ*s)
    return MRP(v*M)
end

# ~~~~~~~~~~~~~~~ Rotation ~~~~~~~~~~~~~~~ #
@inline (*)(p::MRP, r::StaticVector) = UnitQuaternion(p)*r
@inline (\)(p::MRP, r::StaticVector) = UnitQuaternion(p)\r

# ~~~~~~~~~~~~~~~ Kinematics ~~~~~~~~~~~~~~~ #
function kinematics(p::MRP, ω)
    # From Fundamentals of Spacecraft Attitude Determination and Control, eq (3.24)
    #     here A = (I(3) + 2*(skew(p)^2 + skew(p))/(1+p'p)) * (1+p'p)
    #            = √R * (1 + p'p)
    #     where √R*√R = (√R')*(√R') = RotMatrix(p)
    p = params(p)
    A = SA[  #
        1 + p[1]^2 - p[2]^2 - p[3]^2  2(p[1]*p[2] - p[3])      2(p[1]*p[3] + p[2]);
        2(p[2]*p[1] + p[3])            1-p[1]^2+p[2]^2-p[3]^2   2(p[2]*p[3] - p[1]);
        2(p[3]*p[1] - p[2])            2(p[3]*p[2] + p[1])      1-p[1]^2-p[2]^2+p[3]^2]
    0.25*A*ω
end

# ~~~~~~~~~~~~~~~ Useful Jacobians ~~~~~~~~~~~~~~~ #
function ∇rotate(p::MRP, r)
    p = params(p)
    4( (1-p'p)*skew(r)*(4p*p'/(1+p'p) - I) - (4/(1+p'p)*skew(p) + I)*2*skew(p)*r*p'
      - 2*(skew(p)*skew(r) + skew(skew(p)*r)))/(1+p'p)^2
end

function ∇composition2(p2::MRP, p1::MRP)
    p2,p1 = params(p2), params(p1)
    n1 = p1'p1
    n2 = p2'p2
    D = 1 / (1+n1*n2 - 2p1'p2)
    d1 = (-2p1*p2' + (1-n1)*I - skew(2p1) ) * D
    d2 = -((1-n2)*p1 + (1-n1)*p2 - cross(2p1, p2) ) * D^2 *
        (2p2*n1 - 2p1)'
    return d1 + d2
end

function ∇composition1(p2::MRP, p1::MRP)
    p2,p1 = params(p2), params(p1)
    n1 = p1'p1
    n2 = p2'p2
    D = 1 / (1+n1*n2 - 2p1'p2)
    d1 = ((1-n2)*I + -2p2*p1' + 2skew(p2) ) * D
    d2 = -((1-n2)*p1 + (1-n1)*p2 - cross(2p1, p2) ) * D^2 *
        (2p1*n2 - 2p2)'
    d1 + d2
end

function ∇²composition1(p2::MRP, p1::MRP, b::AbstractVector)
    check_length(b, 3)
    p2,p1 = params(p2), params(p1)
    n1 = p1'p1
    n2 = p2'p2
    D = 1 / (1+n1*n2 - 2p1'p2)  # scalar
    dD = -D^2 * (n2*2p1 - 2p2)  # 3×1 (correct)
    A = -((1-n2)*p1 + (1-n1)*p2 - cross(2p1, p2) )  # 3×1
    dA = -I*(1-n2) + 2p2*p1' - 2skew(p2)  # 3×3 (correct)
    B = 2(p1*n2 -  p2)  # 3×1
    dB = 2n2*b  # 3×1

    # d1b = ((1-n2)*I + -2p1*p2' - 2skew(p2) ) * D * b
    dd1b = -2I*(p2'b)*D + ((1-n2)*I + -2p1*p2' - 2skew(p2) )*b*dD' # 3×3 (correct)

    dd2 = D^2 * B * b' * dA +
         D^2 * A'b*2*(n2*I) +
         B*A'b* 2D * dD' # 3×3 (correct)
    dd1b + dd2
end

function ∇differential(p::MRP)
    p = params(p)
    n = p'p
    # p = params(p)
    # (1-n)I + 2(skew(p) + p*p')
    # @SMatrix [n + 2p.x^2      2(p.x*p.y-p.z)  2(p.x*p.z+p.y);
    #           2(p.y*p.x+p.z)  n + 2p.y^2      2(p.y*p.z-p.x);
    #           2(p.z*p.x-p.y)  2(p.z*p.y+p.x)  n + 2p.z^2]
    #
    # p2 = params(p)
    # n2 = p2'p2
    return (1-n)*I + 2(skew(p) + p*p')
end

function ∇²differential(p2::MRP, b::AbstractVector)
    check_length(b, 3)
    p2 = params(p2)
    n2 = p2'p2
    A = -p2  # 3×1
    B = -2p2  # 3×1
    D = 1
    dD = 2p2

    dA = -I*(1-n2) - 2skew(p2)  # 3×3
    dB = 2n2*I  # 3×3

    d1 = (-2p2'b*I*D) - (dA'b * dD')
    d2 = dB * A'b * D^2 +
        (-(1-n2)*b + (skew(2p2))*b)*B' * D^2 +
        B*A'b * 2D * dD'
    d1 + d2'
end

"""
    ∇err(p1::MRP, p2::MRP)

Jacobian of `p1\\p2` wrt `p2`
"""
function ∇err(p1::MRP, p2::MRP)
    q1,q2 = params(p1), params(p2)
    n1,n2 = q1'q1,   q2'q2
    θ = 1/((1+n1)*(1+n2))
    s1,s2 = (1-n1), (1-n2)
    v1 = -2q1
    v2 =  2q2
    s = s1*s2 - v1'v2
    v = s1*v2 + s2*v1 + v1 × v2
    dsdp = -2s1*q2 - 2v1
    dvdp = 2s1*I + -2v1*q2' + 2skew(v1)
    dθdp = -θ^2*(1+n1)*2q2

    M = θ/(1+θ*s)
    dMdp = 1/(1+θ*s)*dθdp - θ/(1+θ*s)^2*(dθdp*s + θ*dsdp)
    return dvdp*M + v*dMdp'
end

"""
    ∇²err(p1::MRP, p2::MRP, b::StaticVector{3})

Jacobian of `(∂/∂p p1\\p2)'b` wrt `p2`
"""
function ∇²err(p1::MRP, p2::MRP, b::AbstractVector)
    check_length(b, 3)
    q1,q2 = params(p1), params(p2)
    n1,n2 = q1'q1, q2'q2
    θ = 1/((1+n1)*(1+n2))
    s1,s2 = (1-n1), (1-n2)
    v1 = -2q1
    v2 =  2q2
    s = s1*s2 - v1'v2
    v = s1*v2 + s2*v1 + v1 × v2

    dsdp = -2s1*q2 - 2v1  # 3×1
    dsdp2 = -2s1*I  # 3×3

    dvdp = 2s1*b + -2q2*v1'b - 2skew(v1)*b
    dvdp2 = -I*2v1'b

    dθdp = -θ^2*(1+n1)*2q2  # 3×1
    dθdp2 = -2θ*(1+n1)*2q2*dθdp' - θ^2*(1+n1)*2I # 3×3

    M = θ/(1+θ*s)  # scalar
    dMdp = 1/(1+θ*s)*dθdp - θ/(1+θ*s)^2*(dθdp*s + θ*dsdp) # 3×1
    dM2 = θ/(1+θ*s)^2  # scalar
    dM3 = dθdp*s + θ*dsdp  # 3×1
    dM2dp = dθdp'/(1+θ*s)^2 - 2θ/(1+θ*s)^3 * (dθdp*s + θ*dsdp)'
    dM3dp = dθdp2*s + dθdp*dsdp' + dsdp*dθdp' + θ*dsdp2

    dMdp2 = -1/(1+θ*s)^2*dθdp*(dθdp*s + dsdp*θ)' + 1/(1+θ*s)*dθdp2  # good
    dMdp2 -= dM3*dM2dp + dM2*dM3dp

    vb = s1*v2'b + s2*v1'b + b'skew(v1)*v2  # scalar
    vpdp = s1*2b' - 2q2' * (v1'b)  + b'skew(v1)*2 # good
    # vpdp = s1*2b' - 2q2 * (v1'b)

    d1 = M*dvdp2 + dvdp*dMdp'
    d2 = dMdp2*vb + dMdp*vpdp
    return d1 + d2
end


text:
"""
```
load_data(m::AbstractModel; try_disk::Bool = true, verbose::Symbol = :low)
```

Create a DataFrame with all data series for this model, fully transformed.

First, check the disk to see if a valid dataset is already stored in `inpath(m, \"data\")`. A
dataset is valid if every series in `m.observable_mappings` is present and the entire sample is
contained (from `date_presample_start` to `date_mainsample_end`. If no valid dataset is
already stored, the dataset will be recreated. This check can be eliminated by passing
`try_disk=false`.

If the dataset is to be recreated, in a preliminary stage,
intermediate data series as specified in `m.observable_mappings` are
loaded in levels using `load_data_levels`. See `?load_data_levels` for
more details.

Then, the series in levels are transformed as specified in `m.observable_mappings`. See
`?transform_data` for more details.

If `m.testing` is false, then the resulting DataFrame is saved to disk as `data_<yymmdd>.csv`.
The data are then returned to the caller.
"""
function load_data(m::AbstractModel; cond_type::Symbol = :none, try_disk::Bool = true, verbose::Symbol=:low)
    recreate_data = false

    # Check if already downloaded
    if try_disk && has_saved_data(m; cond_type=cond_type)
        filename = get_data_filename(m, cond_type)
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            print("Reading dataset $(filename) from disk...")
        end
        df = read_data(m; cond_type = cond_type)
        if isvalid_data(m, df; cond_type = cond_type)
            if VERBOSITY[verbose] >= VERBOSITY[:low]
                println("dataset from disk valid")
            end
        else
            if VERBOSITY[verbose] >= VERBOSITY[:low]
                println("dataset from disk not valid")
            end
            recreate_data = true
        end
    else
        recreate_data = true
    end

    # Download routines
    if recreate_data
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            println("Creating dataset...")
        end

        levels = load_data_levels(m; verbose=verbose)
        if cond_type in [:semi, :full]
            cond_levels = load_cond_data_levels(m; verbose=verbose)
            levels, cond_levels = reconcile_column_names(levels, cond_levels)
            levels = vcat(levels, cond_levels)
            na2nan!(levels)
        end
        df = transform_data(m, levels; cond_type=cond_type, verbose=verbose)

        # Ensure that only appropriate rows make it into the returned DataFrame.
        start_date = date_presample_start(m)
        end_date   = if cond_type in [:semi, :full]
            date_conditional_end(m)
        else
            date_mainsample_end(m)
        end
        df = df[start_date .<= df[:, :date] .<= end_date, :]

        if !m.testing
            save_data(m, df; cond_type=cond_type)
        end
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            println("dataset creation successful")
        end

        # NaN out conditional period variables not in `cond_semi_names(m)` or
        # `cond_full_names(m)` if necessary
        nan_cond_vars!(m, df; cond_type = cond_type)

        # check that dataset is valid
        isvalid_data(m, df)
    end

    return df
end

"""
```
load_data_levels(m::AbstractModel; verbose::Symbol=:low)
```

Load data in levels by appealing to the data sources specified for the model. Data from FRED
is loaded first, by default; then, merge other custom data sources.

Check on disk in `inpath(m, \"data\")` datasets, of the correct
vintage, corresponding to the ones required by the entries in
`m.observable_mappings`. Load the appropriate data series (specified
in `m.observable_mappings[key].input_series`) for each data source.

To accomodate growth rates and other similar transformations, more rows of data may be
downloaded than otherwise specified by the date model settings. (By the end of the process,
these rows will have been dropped.)

Data from FRED (i.e. the `:fred` data source) are treated separately. These are downloaded
using `load_fred_data`. See `?load_fred_data` for more details.

Data from non-FRED data sources are read from disk, verified, and merged.
"""
function load_data_levels(m::AbstractModel; verbose::Symbol=:low)
    # Start two quarters further back than `start_date` as we need these additional
    # quarters to compute differences.
    start_date = date_presample_start(m) - Dates.Month(6)
    end_date = date_mainsample_end(m)

    # Parse m.observable_mappings for data series
    data_series = parse_data_series(m)

    # Load FRED data
    df = load_fred_data(m; start_date=firstdayofquarter(start_date), end_date=end_date, verbose=verbose)

    # Set ois series to load
    if n_anticipated_shocks(m) > 0
        data_series[:OIS] = [Symbol("ant$i") for i in 1:n_anticipated_shocks(m)]
    end

    # For each additional source, search for the file with the proper name. Open
    # it, read it in, and merge it with fred_series

    vint = data_vintage(m)

    for source in keys(data_series)

        # Check that this source is actually used
        mnemonics = data_series[source]
        if isempty(mnemonics)
            warn("No series were specified from $(string(source))")
            continue
        end

        # Skip FRED sources, which have already been handled
        # Conditional data are handled in `load_cond_data_levels`
        if source in [:FRED, :conditional]
            continue
        end

        # Read and merge data from this source
        file = inpath(m, "raw", "$(lowercase(string(source)))_$vint.csv")

        if isfile(file)
            if VERBOSITY[verbose] >= VERBOSITY[:low]
                println("Reading data from $file...")
            end

            # Read in dataset and check that the file contains data for the proper dates
            addl_data = CSV.read(file)

            # Convert dates from strings to quarter-end dates for date arithmetic
            format_dates!(:date, addl_data)

            # Warn on sources with incomplete data; missing data will be replaced with NaN
            # during merge.
            if !in(lastdayofquarter(start_date), addl_data[:date]) ||
                !in(lastdayofquarter(end_date), addl_data[:date])

                warn("$file does not contain the entire date range specified; NaNs used.")
            end

            # Make sure each mnemonic that was specified is present
            for series in mnemonics
                if !in(series, names(addl_data))
                    error("$(string(series)) is missing from $file.")
                end
            end

            # Extract just the columns and rows of the dataset we want, and merge them with
            # data
            cols = [:date; mnemonics]
            rows = start_date .<= addl_data[:date] .<= end_date

            addl_data = addl_data[rows, cols]
            df = join(df, addl_data, on=:date, kind=:outer)
        else
            # If series not found, use all NaNs
            addl_data = DataFrame(fill(NaN, (size(df,1), length(mnemonics))))
            names!(addl_data, mnemonics)
            df = hcat(df, addl_data)
            warn("$file was not found; NaNs used")
        end
    end

    # turn NAs into NaNs
    na2nan!(df)

    sort!(df, :date)

    # print population level data to a file
    if !m.testing
        filename = inpath(m, "raw", "population_data_levels_$vint.csv")
        mnemonic = parse_population_mnemonic(m)[1]
        if !isnull(mnemonic)
            CSV.write(filename, df[:,[:date, get(mnemonic)]])
        end
    end

    return df
end

"""
```
load_cond_data_levels(m::AbstractModel; verbose::Symbol=:low)
```

Check on disk in `inpath(m, \"cond\")` for a conditional dataset (in levels) of the correct
vintage and load it.

The following series are also loaded from `inpath(m, \"raw\")` and either
appended or merged into the conditional data:

- The last period of (unconditional) data in levels
  (`data_levels_<yymmdd>.csv`), used to calculate growth rates
- The first period of forecasted population
  (`population_forecast_<yymmdd>.csv`), used for per-capita calculations
"""
function load_cond_data_levels(m::AbstractModel; verbose::Symbol=:low)

    # Prepare file name
    cond_vint = cond_vintage(m)
    cond_idno = lpad(cond_id(m), 2, 0) # print as 2 digits
    file = inpath(m, "cond", "cond_cdid=" * cond_idno * "_cdvt=" * cond_vint * ".csv")

    if isfile(file)
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            println("Reading conditional data from $file...")
        end

        # Read data
        cond_df = CSV.read(file)
        format_dates!(:date, cond_df)

        date_cond_end = cond_df[end, :date]

        # Use population forecast as population data
        population_forecast_file = inpath(m, "raw", "population_forecast_" * data_vintage(m) * ".csv")
        if isfile(population_forecast_file) && !isnull(get_setting(m, :population_mnemonic))
            pop_forecast = CSV.read(population_forecast_file)

            population_mnemonic = get(parse_population_mnemonic(m)[1])
            rename!(pop_forecast, :POPULATION =>  population_mnemonic)
            na2nan!(pop_forecast)
            format_dates!(:date, pop_forecast)

            cond_df = join(cond_df, pop_forecast, on=:date, kind=:left)

            # Turn NAs into NaNs
            na2nan!(cond_df)
            sort!(cond_df, :date)

            return cond_df
        else
            error("Population forecast data in $population_forecast_file not found, but required to load conditional data")
        end
    else
        # If series not found, throw an error
        error("Conditional data in $file not found")
    end
end

"""
```
save_data(m::AbstractModel, df::DataFrame; cond_type::Symbol = :none)
```

Save `df` to disk as CSV. File is located in `inpath(m, \"data\")`.
"""
function save_data(m::AbstractModel, df::DataFrame; cond_type::Symbol = :none)
    filename = get_data_filename(m, cond_type)
    CSV.write(filename, df)
end

"""
```
has_saved_data(m::AbstractModel; cond_type::Symbol = :none)
```

Determine if there is a saved dataset on disk for the required vintage and
conditional type.
"""
function has_saved_data(m::AbstractModel; cond_type::Symbol = :none)
    filename = get_data_filename(m, cond_type)
    return isfile(filename)
end

"""
```
read_data(m::AbstractModel; cond_type::Symbol = :none)
```

Read CSV from disk as DataFrame. File is located in `inpath(m, \"data\")`.
"""
function read_data(m::AbstractModel; cond_type::Symbol = :none)
    filename = get_data_filename(m, cond_type)
    df       = CSV.read(filename)

    # Convert date column from string to Date
    df[:date] = map(Date, df[:date])

    # NaN out conditional period variables not in `cond_semi_names(m)` or
    # `cond_full_names(m)` if necessary
    nan_cond_vars!(m, df; cond_type = cond_type)

    return df
end

"""
```
isvalid_data(m::AbstractModel, df::DataFrame; cond_type::Symbol = :none)
```

Return if dataset is valid for this model, ensuring that all observables are contained and
that all quarters between the beginning of the presample and the end of the mainsample are
contained. Also checks to make sure that expected interest rate data is available if `n_anticipated_shocks(m) > 0`.
"""
function isvalid_data(m::AbstractModel, df::DataFrame; cond_type::Symbol = :none)
    valid = true

    # Ensure that every series in m_series is present in df_series
    m_series = collect(keys(m.observable_mappings))
    df_series = names(df)
    coldiff = setdiff(m_series, df_series)
    valid = valid && isempty(coldiff)
    if !isempty(coldiff)
        println("Columns of 'df' do not match expected.")
        println(coldiff)
    end

    # Ensure the dates between date_presample_start and date_mainsample_end are contained.
    actual_dates = df[:date]

    start_date = date_presample_start(m)
    end_date   = if cond_type in [:semi, :full]
        date_conditional_end(m)
    else
        date_mainsample_end(m)
    end
    expected_dates = get_quarter_ends(start_date, end_date)
    datesdiff = setdiff(expected_dates, actual_dates)

    valid = valid && isempty(datesdiff)
    if !isempty(datesdiff)
        println("Dates of 'df' do not match expected.")
        println(datesdiff)
    end

    # Ensure that no series is all NaN
    for col in setdiff(names(df), [:date])
        if all(isnan.(df[col]))
            warn("df[$col] is all NaN.")
        end
    end

    return valid
end

"""
```
df_to_matrix(m, df; cond_type = :none, in_sample = true)
```

Return `df`, converted to matrix of floats, and discard date column. Also ensure
that rows are sorted by date and columns by `m.observables`, with the option to specify
whether or not the out of sample rows are discarded. The output of this
function is suitable for direct use in `estimate`, `posterior`, etc.

### Keyword Arguments:
- `include_presample::Bool`: indicates whether or not there are presample periods.
- `in_sample::Bool`: indicates whether or not to discard rows that are out of sample. Set this flag to false in
the case that you are calling filter_shocks! in the scenarios codebase.
"""
function df_to_matrix(m::AbstractModel, df::DataFrame; cond_type::Symbol = :none, include_presample::Bool = true, in_sample::Bool = true)
    # Sort rows by date
    df1 = sort(df, :date)

    if in_sample
        start_date = if include_presample
            date_presample_start(m)
        else
            date_mainsample_start(m)
        end
        end_date = if cond_type in [:semi, :full]
            date_conditional_end(m)
        else
            date_mainsample_end(m)
        end
        df1 = df1[start_date .<= df[:, :date] .<= end_date, :]
    end

    # Discard columns not used
    cols = collect(keys(m.observables))
    sort!(cols, by = x -> m.observables[x])
    df1 = df1[cols]

    y = convert(Matrix{Float64}, df1)'

    if any(isnan.(y)) & get_setting(m, :use_chand_recursion)
        error("Cannot use Chandrasekhar recursions with missing data.")
    end
    return y
end

"""
```
data_to_df(m, data, start_date)
```

Create a `DataFrame` out of the matrix `data`, including a `:date` column
beginning in `start_date`.  Variable names and indices are obtained from
`m.observables`.
"""
function data_to_df{T<:AbstractFloat}(m::AbstractModel, data::Matrix{T}, start_date::Date)
    # Check number of rows = number of observables
    nobs = n_observables(m)
    @assert size(data, 1) == nobs "Number of rows of data matrix ($(size(data, 1))) must equal number of observables ($nobs)"

    # Initialize DataFrame and add dates
    nperiods = size(data, 2)
    end_date = iterate_quarters(start_date, nperiods - 1)
    dates = quarter_range(start_date, end_date)
    df = DataFrame(date = dates)

    # Add observables
    for var in keys(m.observables)
        ind = m.observables[var]
        df[var] = vec(data[ind, :])
    end

    return df
end

"""
```
parse_data_series(m::AbstractModel)
```

Parse `m.observable_mappings` for the data sources and mnemonics to read in.

Returns a `Dict{Symbol, Vector{Symbol}}` mapping sources => mnemonics found in that data file.
"""
function parse_data_series(m::AbstractModel)

    data_series = Dict{Symbol, Vector{Symbol}}()

    # Parse vector of observable mappings into data_series dictionary
    for obs in values(m.observable_mappings)
        for series in obs.input_series
            mnemonic, source = map(Symbol, split(string(series), DSGE_DATASERIES_DELIM))

            if !in(source, keys(data_series))
                data_series[source] = Vector{Symbol}()
            end

            if !in(mnemonic, data_series[source])
                push!(data_series[source], mnemonic)
            end
        end
    end
    data_series
end

"""
```
read_population_data(m; verbose = :low)

read_population_data(filename; verbose = :low)
```

Read in population data stored in levels, either from
`inpath(m, \"raw\", \"population_data_levels_[vint].csv\"`) or `filename`.
"""
function read_population_data(m::AbstractModel; verbose::Symbol = :low)
    vint = data_vintage(m)
    filename = inpath(m, "raw", "population_data_levels_" * vint * ".csv")
    read_population_data(filename; verbose = verbose)
end

function read_population_data(filename::String; verbose::Symbol = :low)
    if VERBOSITY[verbose] >= VERBOSITY[:low]
        println("Reading population data from $filename...")
    end

    df = CSV.read(filename)

    na2nan!(df)
    format_dates!(:date, df)
    sort!(df, :date)

    return df
end

"""
```
read_population_forecast(m; verbose = :low)

read_population_forecast(filename, population_mnemonic, last_recorded_date; verbose = :low)
```

Read in population forecast in levels, either from
`inpath(m, \"raw\", \"population_forecast_[vint].csv\")` or `filename`.
If that file does not exist, return an empty `DataFrame`.

"""
function read_population_forecast(m::AbstractModel; verbose::Symbol = :low)
    population_forecast_file = inpath(m, "raw", "population_forecast_" * data_vintage(m) * ".csv")
    population_mnemonic = parse_population_mnemonic(m)[1]

    if isnull(population_mnemonic)
        error("No population mnemonic provided")
    else
        read_population_forecast(population_forecast_file, get(population_mnemonic); verbose = verbose)
    end
end

function read_population_forecast(filename::String, population_mnemonic::Symbol;
                                  verbose::Symbol = :low)
    if isfile(filename)
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            println("Loading population forecast from $filename...")
        end

        df = CSV.read(filename)

        rename!(df, :POPULATION => population_mnemonic)
        na2nan!(df)
        format_dates!(:date, df)
        sort!(df, :date)

        return df[:, [:date, population_mnemonic]]
    else
        if VERBOSITY[verbose] >= VERBOSITY[:low]
            warn("No population forecast found")
        end

        return DataFrame()
    end
end

ext: jl


text:
using HDF5
const PRINT_MEMORY = `ps -p $(getpid()) -o rss=`
const DATA = zeros(1000)
macro memtest(ex)
    @info :Memory
    quote
        for i in 1:100
            for _ in 1:100
                $ex
            end
            # HDF5.h5_garbage_collect()
            GC.gc(); print(rpad(i, 8)); run(PRINT_MEMORY)
        end
    end
end
@memtest h5open("/tmp/memtest.h5", "w") do file
    dset = d_create(file, "A", datatype(DATA), dataspace(DATA), chunk=(100,))
    dset[:] = DATA[:]
end
@memtest h5open("/tmp/memtest.h5", "w") do file
    file["A", chunk=(100,)] = DATA[:]
end
@memtest h5open("/tmp/memtest.h5", "r") do file
    file["A","dxpl_mpio", 0]
end


text:
using DataFrames
# Dataset should be there as mentioned in the chapter
DfTRoadSafety_Accidents_2014 = readtable("DfTRoadSafety_Accidents_2014.csv")
DfTRoadSafety_Vehicles_2014 = readtable("DfTRoadSafety_Vehicles_2014.csv")
left_DfTRoadSafety_2014 = join(DfTRoadSafety_Accidents_2014, DfTRoadSafety_Vehicles_2014, on = :_Accident_Index, kind = :left)
by(DfTRoadSafety_Accidents_2014, :Location_Northing_OSGR, size)

max_stars_repo_path: Chapter 2/6.jl

max_stars_repo_name: PacktPublishing/Julia-for-data-science

max_issues_repo_path: Chapter 2/6.jl

max_issues_repo_name: PacktPublishing/Julia-for-data-science

max_forks_repo_path: Chapter 2/6.jl

max_forks_repo_name: PacktPublishing/Julia-for-data-science


text:
module GradDescent

using LinearAlgebra
export
    Optimizer,
    VanillaGradDescent,
    Inversedecay,
    Momentum,
    Adagrad,
    Adadelta,
    RMSprop,
    Adam,
    Adamax,
    Nadam,
    update,
    t

include("AbstractOptimizer.jl")
include("VanillaGradDescent.jl")
include("InverseDecayOptimizer.jl")
include("MomentumOptimizer.jl")
include("AdaGradOptimizer.jl")
include("AdaDeltaOptimizer.jl")
include("RMSpropOptimizer.jl")
include("AdamOptimizer.jl")
include("AdamaxOptimizer.jl")
include("NadamOptimizer.jl")

end # module


text:
using Genie
using HTTP
using OpenTelemetryInstrumentationGenie

route("/hello") do
    "world"
end

up(7777)

sleep(1)

HTTP.get("http://localhost:7777/hello")
HTTP.get("http://localhost:7777/hello")
HTTP.get("http://localhost:7777/hello")

down()


text:
include("init.jl")

@testset "Combination Federate" begin

    broker = createBroker();

    cfed = h.helicsCreateCombinationFederateFromConfig(joinpath(@__DIR__, "combinationfederate.json"))

    @test h.helicsFederateIsValid(cfed)

    @test h.helicsFederateGetEndpointCount(cfed) == 6
    @test h.helicsFederateGetFilterCount(cfed) == 6
    @test h.helicsFederateGetInputCount(cfed) == 7

    ept = h.helicsFederateGetEndpointByIndex(cfed, 0);

    @test h.helicsEndpointGetName(ept) == "EV_Controller/EV6"

    filt = h.helicsFederateGetFilterByIndex(cfed, 3)

    @test h.helicsFilterGetName(filt) == "EV_Controller/filterEV3"

    f = h.helicsFederateGetFilter(cfed, "EV_Controller/filterEV3")
    @test h.helicsFilterGetName(f) == "EV_Controller/filterEV3"

    ipt = h.helicsFederateGetInputByIndex(cfed, 4)
    @test h.helicsInputGetExtractionUnits(ipt) == ""
    @test h.helicsSubscriptionGetKey(ipt) == "IEEE_123_feeder_0/charge_EV3"

    h.helicsEndpointClearMessages(ept)

    h.helicsFederateDestroy(cfed)
    h.helicsBrokerDestroy(broker)

end


text:
"Kamada Kawai Layout"
function kamada_kawai(G::PowerModelsGraph, dist::Union{Nothing,Matrix{Float64}}=nothing, pos::Union{Nothing,Matrix{Float64}}=nothing, weight="weight", scale=1, center=nothing, dim=2)
    graph = G.graph # convert to undirected graph
    nNodes = LightGraphs.nv(graph)
    if nNodes == 0
        return
    end

    if dist===nothing
        dist=Dict()
        for i in 1:nNodes
            dist[i]=LightGraphs.dijkstra_shortest_paths(graph, i).dists
        end
    end
    dist_mtx = 1e6 * ones(nNodes, nNodes) #reform dist into matrix, probably easier??
    for nr in LightGraphs.vertices(graph)
        rdist = dist[nr]
        for nc in LightGraphs.vertices(graph)
            dist_mtx[nr,nc] = rdist[nc]
        end
    end
    if pos===nothing
        if dim >= 3
            pos= 2 .* rand(Float64,dim, nNodes) .- 1 #??? make a matrix n x dim
        elseif dim == 2
            a = LightGraphs.adjacency_matrix(graph)
            pos = convert(Array,RecursiveArrayTools.VectorOfArray(NetworkLayout.Circular.layout(a)))
        else
            pos = [pt for pt in range(0, 1, length=nNodes)]
        end
    end

    pos = _kamada_kawai_solve(dist_mtx, pos, dim, nNodes)
    positions = [[pos[1,i],pos[2,i]] for i in 1:size(pos,2)]
    return positions
end

"""
Anneal node locations based on the Kamada-Kawai cost-function,
using the supplied matrix of preferred inter-node distances,
and starting locations.
"""
function _kamada_kawai_solve(dist_mtx::Array{Float64,2}, pos_arr::Array{Float64,2}, dim::Int, nNodes::Int)::Matrix{Float64}
    pos_vec = pos_arr[:]
    meanwt = 1e-3

    opt = NLopt.Opt(:LD_LBFGS, length(pos_vec))
    opt.xtol_rel = 1e-4
    opt.min_objective = (pos_vec, grad)->_kamada_kawai_costfn(pos_vec,grad,1.0./(dist_mtx + LinearAlgebra.I(nNodes) * 1e-3), meanwt, dim, nNodes)
    (minf,minx,ret) = NLopt.optimize(opt, pos_vec)

    numevals = opt.numevals # the number of function evaluations
    #println("got $minf at $minx after $numevals iterations (returned $ret)")
    return reshape(minx, dim, nNodes)
end

"Cost-function and gradient for Kamada-Kawai layout algorithm"
function _kamada_kawai_costfn(pos_vec::Vector{Float64}, grad::Vector{Float64}, invdist::Array{Float64,2}, meanweight::Float64, dim::Int, nNodes::Int)
    pos_arr = reshape(pos_vec,dim,nNodes)
    delta = zeros(Float64,dim,nNodes,nNodes)
    for i in 1:nNodes
        delta[:,:,i] = pos_arr .- pos_arr[:,i]
    end

    nodesep = zeros(Float64, dim, nNodes)
    nodesep =  reshape(sqrt.(sum(x -> x^2, delta; dims=1)),nNodes,nNodes)
    direction = OMEinsum.ein"ijk,jk -> ijk"(delta, 1.0./(nodesep+LinearAlgebra.I(nNodes)*1e-3))

    offset = nodesep .* invdist .- 1.0
    for i in 1:nNodes
        offset[i,i] = 0.0
    end

    cost = 0.5 * sum(offset.^2)
    gradient = OMEinsum.ein"jk,jk,ijk->ij"(invdist, offset, direction) - OMEinsum.ein"jk,jk,ijk->ik"(invdist, offset, direction)

    # # Additional parabolic term to encourage mean position to be near origin:
    sumpos = sum(pos_arr, dims=2)
    cost += 0.5 .* meanweight .* sum(sumpos.^2)
    origin_penalty = meanweight*sumpos
    for i in 1:nNodes
        gradient[:,i] += origin_penalty
    end
    if length(grad) > 0
        grad[:] = gradient[:]
    end

    return cost
end

function layout_graph_kamada_kawai!(G) #return type must be dictionary
    return positions = kamada_kawai(G)
end


text:
using UKrig
using Test

#@testset "UKrig.jl" begin
    # Write your own tests here.
    @test 1==1


text:
# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Script to compile data for some descriptive statistics
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

using DataFrames, CSV, XLSX, LinearAlgebra, Statistics, StatsBase, StatsPlots, Dates, Pipe

# other scripts
dir_home = "x:/VIVES/1-Personal/Florian/git/2022_01_comext/src/"
include(dir_home * "functions.jl")

# location of data input/output (io)
dir_io = "C:/Users/u0148308/data/comext/"
dir_dropbox = "C:/Users/u0148308/Dropbox/BREXIT/"

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# VLAIO
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

df_VLAIO = CSV.read(dir_dropbox * "rawdata/" * "df_VLAIO" * ".csv", DataFrame)
transform!(df_VLAIO, ["PRODUCT_NC", "PRODUCT_NC_LAB1", "PRODUCT_NC_LAB2", "DECLARANT_ISO", "TRADE_TYPE", "PARTNER_ISO", "FLOW"] .=> ByRow(string), renamecols=false)
transform!(df_VLAIO, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> convert(Union{Missing, Float64}, x)), renamecols=false)
df_VLAIO.PRODUCT_NC = lpad.(string.(df_VLAIO.PRODUCT_NC), 8, '0') # needs to be done again

# remove missing, only for VALUES
#   - drop missing [VALUE, QUANTITY]: lose 1-16027346/21667553 ~26% observations
#   - drop missing VALUE: lose 1-21666543/21667553 ~ 0.1% observations
subset!(df_VLAIO, :VALUE_IN_EUROS .=> ByRow(x -> !ismissing(x)))

# not needed in case of VALUES!!
# MAD adjustment
#   - cutoff = 3: lose further 1-15177297/21666543 ~30% observations (~30% of total)
#   - cutoff = 4: lose further 1-15812254/21666543 ~27% observations (~27% of total)
# outlier_cutoff = 4
# gdf = groupby(df_VLAIO, "PRODUCT_NC")
# df_VLAIO = transform(gdf, :VALUE_IN_EUROS => MAD_method => :MAD)
# subset!(df_VLAIO, :MAD => ByRow(x -> x < outlier_cutoff))
# df_VLAIO = df_VLAIO[:, Not(:MAD)]

# add EU/WORLD again
EU27 = ["Roemenië", "Griekenland", "Oostenrijk", "Polen", "Duitsland", "Spanje", "Hongarije", "Slovakije", "Italië", "Nederland",
       "Frankrijk", "Letland", "Kroatië", "Cyprus", "Malta", "Litouwen", "Slovenië", "Estland", "Portugal", "Finland", "Tsjechië",
       "Luxemburg", "Zweden", "Denemarken", "Bulgarije", "Ierland", "België"]

# slightly modified functions
function append_WORLD(df::DataFrame)

    # sum over PARTNER
    cols_grouping = ["DECLARANT_ISO", "PRODUCT_NC", "PRODUCT_NC_LAB1", "PRODUCT_NC_LAB2", "TRADE_TYPE", "FLOW", "PERIOD"]
    gdf = groupby(df, cols_grouping)
    df_WORLD = combine(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> (x -> sum(skipmissing(x))), renamecols=false)

    df_WORLD.TRADE_TYPE .= "total"
    df_WORLD.PARTNER_ISO .= "WORLD"

    df = vcat(df, df_WORLD)

    return df
end

function append_EU(df::DataFrame, EU::Vector{String})

    # subset and sum over EU ctrys
    df_EU = subset(df, :PARTNER_ISO => ByRow(x -> x in EU))
    cols_grouping = ["DECLARANT_ISO", "PRODUCT_NC", "PRODUCT_NC_LAB1", "PRODUCT_NC_LAB2", "TRADE_TYPE", "FLOW", "PERIOD"]
    gdf = groupby(df_EU, cols_grouping)
    df_EU = combine(gdf, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> (x -> sum(skipmissing(x))), renamecols=false)

    df_EU.TRADE_TYPE .= "intra"
    df_EU.PARTNER_ISO .= "EU"

    df = vcat(df, df_EU)

    return df
end

df_VLAIO = append_WORLD(df_VLAIO)
df_VLAIO = append_EU(df_VLAIO, EU27)

# double check if no zeros/missing introduced by WORLD/EU
transform!(df_VLAIO, [:VALUE_IN_EUROS, :QUANTITY_IN_KG] .=> ByRow(x -> ifelse(iszero(x) | ismissing(x), missing, x)), renamecols=false)
subset!(df_VLAIO, :VALUE_IN_EUROS => ByRow(x -> !ismissing(x))) # lose 0% observations

cols_subset = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PERIOD", "PRODUCT_NC", "VALUE_IN_EUROS", "QUANTITY_IN_KG"]
df_VLAIO = df_VLAIO[:, cols_subset]

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Figure 5
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# prepare data
df = copy(df_VLAIO)
partners = ["Verenigd Koninkrijk", "Duitsland", "Nederland", "Frankrijk", "Italië"]
subset!(df, :DECLARANT_ISO => ByRow(x -> x == "Vlaanderen"), :PARTNER_ISO => ByRow(x -> x in partners))

cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PERIOD"]
gdf = groupby(df, cols_grouping)
df = combine(gdf, nrow => :COUNT)

transform!(df, :PERIOD => ByRow(x -> Date(string(x), DateFormat("yyyymm"))) => :DATE)
sort!(df)

# HP filter
λ = 20
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW"]
gdf = groupby(df, cols_grouping)
df = transform(gdf, :COUNT => (x -> HP(x, λ)) => :COUNT_HP)

for flow in ["imports", "exports"]

    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :COUNT,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="number of products", title="Flemish "* flow*": number of products")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/VLAIO/summary_stats/" * "fig5_" * flow * "_product_count" * ".png") # export image dropbox

    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :COUNT_HP,
    group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="number of products", title="Flemish "* flow*": number of products (HP, λ=$λ")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/VLAIO/summary_stats/" * "fig5_" * flow * "_product_count_HP" * ".png") # export image dropbox

end

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# comext
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# prepare data
df = copy(df_comext)
partners = ["GB", "DE", "NL", "FR", "IT"]
subset!(df, :DECLARANT_ISO => ByRow(x -> x == "BE"), :PARTNER_ISO => ByRow(x -> x in partners))
subset!(df, :PRODUCT_NC => ByRow(x -> x != "TOTAL"))

cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW", "PERIOD"]
gdf = groupby(df, cols_grouping)
df = combine(gdf, nrow => :COUNT)

transform!(df, :PERIOD => ByRow(x -> Date(string(x), DateFormat("yyyymm"))) => :DATE)
sort!(df)

# HP filter
λ = 20
cols_grouping = ["DECLARANT_ISO", "PARTNER_ISO", "FLOW"]
gdf = groupby(df, cols_grouping)
df = transform(gdf, :COUNT => (x -> HP(x, λ)) => :COUNT_HP)

for flow in ["imports", "exports"]

    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :COUNT,
            group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="number of products", title="Belgian "* flow*": number of products")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/VLAIO/summary_stats/" * "fig5_BE_" * flow * "_product_count" * ".png") # export image dropbox

    p = @df subset(df, :FLOW => ByRow(x -> x == flow)) plot(:DATE, :COUNT_HP,
    group=:PARTNER_ISO, lw=2, legend=:bottomleft, ylabel="number of products", title="Belgian "* flow*": number of products (HP, λ=$λ")
    vline!([Date(2016,6,23)], label="vote", color=:black, lw=1, ls=:solid) # refer
    vline!([Date(2020,01,31)], label="exit", color=:black, lw=1, ls=:dash) # exit
    vline!([Date(2020,12,31)], label="trans end", color=:black, lw=1.5, ls=:dot) # trans end
    savefig(p, dir_dropbox * "results/images/VLAIO/summary_stats/" * "fig5_BE_" * flow * "_product_count_HP" * ".png") # export image dropbox

end


text:
## Julia 1.1.0

# Initialize
using Plots
plotlyjs() # Backend PlotlyJS

left = -4 * pi
steps = 100

x = left

anim = @animate for i = 1:steps
    x = range(left, stop=i * abs(left) * 2 / (steps), length=10*i)
    sinx = sin.(x)
    plot(x, sinx)
end

gif(anim, "./gif/sin.gif", fps = 30)

max_stars_repo_name: fengyfei/sapphire

max_issues_repo_name: fengyfei/sapphire

max_forks_repo_name: fengyfei/sapphire


text:
using ProblemParser
using Utils

using Bijections
using Chain
using FunctionalCollections
using Graphs

const file = find_input(@__FILE__)
const input = parse(Lines(Split('-', Map(Symbol))), slurp(file))
const node2num = @chain input begin
    Iterators.flatten
    unique
    enumerate
    map(reverse, _)
    Dict
    Bijection
end

function make_graph(input)
    g = Graph(length(node2num))
    for (a,b) in input
        add_edge!(g, node2num[a], node2num[b])
    end
    g
end

smallcave(name::Symbol) = @chain name String _[1] islowercase

smallcave(num::Int) = node2num(num) |> smallcave

part1(g::Graph) = part1(g, node2num[:start])

part1(input) = @chain input make_graph part1

function part1(g::Graph, cur::Integer, visited=PersistentSet())
    current::Symbol = node2num(cur)
    current == :end && return 1
    withme = conj(visited, cur)
    paths = 0
    for n in neighbors(g, cur)
        (n in visited && (!smallcave(current) || smallcave(n))) && continue
        paths += part1(g, n, withme)
    end
    paths
end

assertequal(part1(input), 3779)

part2(g::Graph) = part2(g, node2num[:start], 1)

part2(input) = @chain input make_graph part2

function part2(g::Graph, cur::Integer, extra, visited=PersistentSet())
    current::Symbol = node2num(cur)
    current == :end && return 1
    withme = conj(visited, cur)
    paths = 0
    for n in neighbors(g, cur)
        if n ∉ visited || smallcave(current) && !smallcave(n)
            paths += part2(g, n, extra, withme)
        elseif extra > 0 && node2num(n) != :start && smallcave(n)
            paths += part2(g, n, extra - 1, withme)
        end
    end
    paths
end

assertequal(part2(input), 96988)

max_stars_repo_name: CmdQ/AoC

max_issues_repo_name: CmdQ/AoC

max_forks_repo_name: CmdQ/AoC


text:
export TransientCmeAlgorithm

abstract type TransientCmeAlgorithm end

include("sparse/fspsolve.jl")


text:
using Resonance

kneadread = Resonance.load_knead()

DataFrames.transform!(kneadread,
    AsTable(Cols(r"raw")) => ByRow(sum) => :raw,
    AsTable(Cols(r"trimmed")) => ByRow(sum) => :trimmed,
    AsTable(Cols(r"decontaminated")) => ByRow(sum) => :decontaminated,
    AsTable(Cols(r"final")) => ByRow(sum) => :final
)

DataFrames.transform!(kneadread,
    AsTable([:raw, :trimmed]) => ByRow(x-> ismissing(x[2]) ? 0 : (x[1] - x[2]) / x[1] * 100) => :perc_trimmed,
    AsTable([:raw, :decontaminated]) => ByRow(x-> ismissing(x[2]) ? 0 : (x[1] - x[2]) / x[1] * 100) => :perc_decontaminated
)

kneadlong = stack(select(kneadread, [:Sample, :raw, :trimmed, :final]), [:raw, :trimmed, :final])

##

using CairoMakie
using AlgebraOfGraphics

##

fig = Figure(; resolution=(1100, 600))
lims = extrema(skipmissing(kneadread.raw))

ax1 = Axis(fig[1,1]; xlabel="Raw reads", ylabel="Final reads", yscale=log10, xscale=log10)
ax2 = Axis(fig[1,3]; xlabel="Raw reads", ylabel="Final reads", yscale=log10, xscale=log10)

sc1 = scatter!(ax1, kneadread.raw, kneadread.final; color = kneadread.perc_trimmed)
lines!(ax1, Point2f.(zip(lims, lims)); color = :gray)

sc2 = scatter!(ax2, kneadread.raw, kneadread.final; color = kneadread.perc_decontaminated, colormap=:plasma)
lines!(ax2, Point2f.(zip(lims, lims)); color = :gray)

Colorbar(fig[1,2], sc1; label = "Percent trimmed")
Colorbar(fig[1,4], sc2; label = "Percent human")

save("figures/kneaddata_counts.png", fig)
fig

#-

max_stars_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_issues_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_forks_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis


text:
using FlippingModel
using Test

@testset "FlippingModel.jl" begin
    # Write your own tests here.
end


text:
# conversion methods for DataFrames

function rcopy(::Type{T}, s::Ptr{VecSxp};
               normalizenames::Bool=true,
               sanitize::Union{Bool,Nothing}=nothing) where T<:AbstractDataFrame
    if sanitize !== nothing
        Base.depwarn("The `sanitize` keyword argument is deprecated. Use `normalizenames` instead.", :rcopy)
        normalizenames = sanitize
    end
    isFrame(s) || error("s is not an R data frame")
    vnames = rcopy(Array{Symbol},getnames(s))
    if normalizenames
        vnames = [Symbol(replace(string(v), '.' => '_')) for v in vnames]
    end
    DataFrame([vec(rcopy(AbstractArray, c)) for c in s], vnames)
end

## DataFrame to sexp conversion.
function sexp(::Type{RClass{:list}}, d::AbstractDataFrame)
    nr,nc = size(d)
    nv = names(d)
    rd = protect(allocArray(VecSxp, nc))
    try
        for i in 1:nc
            rd[i] = sexp(d[!, nv[i]])
        end
        setattrib!(rd,Const.NamesSymbol, [string(n) for n in nv])
        setattrib!(rd,Const.ClassSymbol, "data.frame")
        setattrib!(rd,Const.RowNamesSymbol, 1:nr)
    finally
        unprotect(1)
    end
    rd
end


text:
# This file is available under the terms of the MIT License

@doc """
Functions for evaluating the objective and gradient.
These are mostly wrappers of the tools in DMDUtil.jl.
"""

###########################################################
# DMD objective function
# #   make sure you already update all alpha, B related vars
# function objective(ar, vars, params, svars;
#     updateP=true, updateB=true, updateR=true)

#     copy!(vars.ar, ar);
#     # update phi
#     updateP && (update_P!(vars, params);
#         update_PQR!(vars, params, svars));
#     # update B
#     updateB && update_B!(vars, params, svars);
#     # update Residual
#     updateR && update_R!(vars, params);

#     # calculate objective value
#     return 0.5*sum(abs2, vars.R);
# end

function BFunc(params)
    T = eltype(params.X);
    # update residual
    BLAS.gemm!('N', 'N', T(1.0), params.P, params.B, T(0.0), params.R);
    params.R .-= params.X
    #
    return params.lossFunc(params.R)
end

function bFunc(params, id)
    T = eltype(params.X);
    # update residual
    BLAS.gemv!('N', T(1.0), params.P, params.b[id], T(0.0), params.r[id]);
    params.r[id] .-= params.x[id]
    #
    return params.lossFunc(params.r[id])
end

function bGrad(gbr, params, id)
    T = eltype(params.X);
    Tr = typeof(real(params.X[1]));
    # update residual
    BLAS.gemv!('N', T(1.0), params.P, params.b[id], T(0.0), params.r[id]);
    params.r[id] .-= params.x[id];
    #
    params.lossGrad(params.r[id]);
    #
    gb = vr2vc(gbr);
    BLAS.gemv!('T', T(1.0), params.P, params.r[id], T(0.0), gb);
    gbr[1:2:end] .*= Tr(2.0);
    gbr[2:2:end] .*= Tr(-2.0);
end

function projbl2(params, id)
    copyto!(params.b[id],params.P\params.x[id])
end

function projb(params, id)
    # define the function and gradient interface for optim
    function f(br)
        copyto!(params.br[id], br);
        return bFunc(params, id)
    end
    #
    function g!(gbr, br)
        copyto!(params.br[id], br);
        bGrad(gbr, params, id);
    end

    res = optimize(f, g!, params.br[id], BFGS(),
                   params.inner_opts);

    copyto!(params.br[id], res.minimizer);

end

function projB(params)
    for id = 1:params.n
        projb(params, id);
    end
end

function abFunc(params, id)
    # update P
    update_P!(params);
    # partially minimize over b
    projb(params, id);
    #
    return bFunc(params, id)
end

function aBFunc(params)
    # update P
    update_P!(params);
    # partially minimize over B
    projB(params);
    #
    return BFunc(params)
end

function abGrad(gar, params, id)
    # update P
    update_P!(params);
    #
    T = eltype(params.X);
    Tr = typeof(real(params.X[1]));
    # partially minimize over b
    projb(params, id);
    # update residual
    BLAS.gemv!('N', T(1.0), params.P, params.b[id], T(0.0), params.r[id]);
    params.r[id] .-= params.x[id];
    #
    params.lossGrad(params.r[id]);
    #
    ga = vr2vc(gar);
    #
    params.r[id] .*= t;
    BLAS.gemv!('T', T(1.0), params.P, params.r[id], T(0.0), ga);
    ga .*= params.b[id];
    #
    gar[1:2:end] .*= Tr(2.0);
    gar[2:2:end] .*= Tr(-2.0);
end

function aBGrad(gar, params)
    # update P
    update_P!(params);
    #
    T = eltype(params.X);
    Tr = typeof(real(params.X[1]));
    # partially minimize over B
    projB(params);
    # update residual
    BLAS.gemm!('N', 'N', T(1.0), params.P, params.B, T(0.0), params.R);
    params.R .-= params.X
    #
    params.lossGrad(params.R);
    #
    ga = vr2vc(gar);
    #
    params.R .*= t;
    BLAS.gemm!('T', 'N', T(1.0), params.P, params.R, T(0.0), params.tM);
    params.tM .*= params.B;
    sum!(ga, params.tM);
    #
    gar[1:2:end] .*= Tr(2.0);
    gar[2:2:end] .*= Tr(-2.0);
end

###########################################################
# gradient w.r.t. alpha
# function gradient!(ar, gar, vars, params, svars;
#     updateP=true, updateB=true, updateR=true)

#     copy!(vars.ar, ar);
#     # update phi
#     updateP && (update_P!(vars, params);
#         update_PQR!(vars, params, svars));
#     # update B
#     updateB && update_B!(vars, params, svars);
#     # update Residual
#     updateR && update_R!(vars, params);
#     # calculate gradient
#     grad_ar!(gar, vars, params, svars);
# end

# function gradient!(ar, gar, vars, params, svars, id;
#     updateP=true, updateB=true, updateR=true)

#     copy!(vars.ar, ar);
#     # update phi
#     updateP && (update_P!(vars, params);
#         update_PQR!(vars, params, svars));
#     # update B
#     updateB && update_b!(vars, params, svars, id);
#     # update Residual
#     updateR && update_r!(vars, params, id);
#     # calculate gradient
#     grad_ar!(gar, vars, params, svars, id);
# end


text:
using Test, LinearAlgebra

using XPORTA: write_ieq, IEQ

using BellScenario

@testset "N-(N-1)-N scenario" begin

@time for N in 3:7
    d = N-1

    filepath = "./data/quick_adjacency_decomposition/"
    filename = "$N-$d-$N"

    println("$filename scenario")

    PM = LocalSignaling(N,N,d)
    BG_seed = BellGame(Matrix{Int64}(I,(N,N)), d)

    vertices = LocalPolytope.vertices(PM)
    canonical_facets = LocalPolytope.adjacency_decomposition(vertices, BG_seed, PM, max_vertices = 100)
    games = collect(keys(canonical_facets))

    pretty_print_txt(games, filepath*filename)
    write_ieq(filepath*"ieq/"*filename, convert(IEQ, games))
end

end


text:
# debug functionality

"""
    enable_timings()

Enable the recording of debug timings.
"""
enable_timings() = (TimerOutputs.enable_debug_timings(CUDA); return)
disable_timings() = (TimerOutputs.disable_debug_timings(CUDA); return)


text:
using CairoMakie, Random
Random.seed!(14)
n = 25
m = 50
xPos = LinRange(-π/3,2π,n)
startPos = sin.(xPos)/3
widths = 0.1 .+ 0.2rand(n)
colors = rand(n)
cmap = to_colormap(:starrynight)[5:end]

with_theme(theme_dark()) do
    fig = Figure(resolution = (600,400))
    ax = Axis(fig[1,1])
    for i in 0:3.5:m
        indxs = shuffle(1:n)
        barplot!(xPos[indxs], rand(1:3,n) .+ i; fillto = startPos .+ i,
        color = colors, colormap = cmap, width = widths, strokewidth = 0.5,
        strokecolor = colors[indxs])
    end
    hidedecorations!(ax)
    save("barplot.png", fig, px_per_unit = 2)
    display(fig)
end


text:
using Concorde
using Statistics
using Match
using PyCall

const MAX_TIME_LIMIT = Inf
const MAX_DRONE_RANGE = Inf

mutable struct TSPDroneResult
    total_cost::Float64
    truck_route::Vector{Int}
    drone_route::Vector{Int}
    Ct::Matrix{Float64}
    Cd::Matrix{Float64}
    flying_range::Float64
end

include("tspd_utils.jl")
include("tsp_ep_all.jl")
include("DPS.jl")
include("PythonRL/wrapper.jl")
include("summary.jl")

function local_search_functions(method::String)
    local_search_methods = Function[two_point_move, one_point_move, two_opt_move]
    local_search_methods = @match method begin
        "TSP-ep"        => Function[]
        "TSP-ep-1p"     => Function[one_point_move]
        "TSP-ep-2p"     => Function[two_point_move]
        "TSP-ep-2opt"   => Function[two_opt_move]
        "TSP-ep-all"    => Function[two_point_move, one_point_move, two_opt_move]
    end
    return local_search_methods
end

function solve_tspd(
    x::Vector{Float64},
    y::Vector{Float64},
    truck_cost_factor::Float64,
    drone_cost_factor::Float64;
    n_groups::Int = 1,
    method::String = "TSP-ep-all",
    flying_range::Float64 = MAX_DRONE_RANGE,
    time_limit::Float64 = MAX_TIME_LIMIT
)
    local_search_methods = local_search_functions(method)
    Ct, Cd = cost_matrices_with_dummy(x, y, truck_cost_factor, drone_cost_factor)

    total_tspd_len, truck_route, drone_route = divide_partition_search(Ct, Cd; local_search_methods=local_search_methods, n_groups=n_groups, flying_range=flying_range, time_limit=time_limit)

    result = TSPDroneResult(total_tspd_len, truck_route, drone_route, Ct, Cd, flying_range)
    return result
end

function solve_tspd(
    truck_cost_mtx::Matrix{Float64},
    drone_cost_mtx::Matrix{Float64};
    n_groups::Int = 1,
    method::String = "TSP-ep-all",
    flying_range::Float64 = MAX_DRONE_RANGE,
    time_limit::Float64 = MAX_TIME_LIMIT
)
    Ct, Cd = cost_matrices_with_dummy(truck_cost_mtx, drone_cost_mtx)
    local_search_methods = local_search_functions(method)

    total_tspd_len, truck_route, drone_route = divide_partition_search(Ct, Cd; local_search_methods=local_search_methods, n_groups=n_groups, flying_range=flying_range, time_limit=time_limit)

    result = TSPDroneResult(total_tspd_len, truck_route, drone_route, Ct, Cd, flying_range)
    return result
end


text:
using Test
using LinearAlgebra
include("run_calculation.jl")

@testset "Basic LDA basis construction" begin
    state      = load_state("mg.json")
    parameters = Dict(state[:parameters]...,
                      :xc   => "LDA",
                      :kpts => [1, 2, 3],
                      :ecut => 300)
    basis = get_dftk_basis(parameters, state[:extra])

    lattice = basis.model.lattice
    @test lattice[:, 1] ≈ austrip(1u"Å") * [3.21, 0.0, 0.0]
    @test lattice[:, 2] ≈ austrip(1u"Å") * [-1.605, 2.7799415461480477, 0.0]
    @test lattice[:, 3] ≈ austrip(1u"Å") * [0.0, 0.0, 5.21304]

    atoms = basis.model.atoms
    @test length(atoms) == 1
    @test all(at isa ElementPsp for (at, positions) in atoms)
    @test atoms[1][1].symbol == :Mg
    @test atoms[1][2][1] ≈ [0, 0, 0]
    @test atoms[1][2][2] ≈ [1/3, 2/3, 1/2] atol=1e-5
    @test atoms[1][1].psp.identifier == "hgh/lda/mg-q2.hgh"
    @test atoms[1][1].psp.Zion == 2

    @test basis.model.temperature == 0.0
    @test basis.model.smearing isa DFTK.Smearing.None
    @test basis.model.spin_polarization == :none

    xcterm = only(tt for tt in basis.model.term_types if tt isa Xc)
    @test only(xcterm.functionals) == :lda_xc_teter93

    @test length(basis.kpoints) == 4
    @test basis.Ecut ≈ 300 * austrip(1u"eV") atol=1e-8
end

@testset "PBE semicore basis construction" begin
    state      = load_state("mg.json")
    parameters = Dict(state[:parameters]...,
                      :xc       => "PBE",
                      :kpts     => 3.333,
                      :smearing => ("Gaussian", 10),
                      :pps      => "hgh.k")  # semicore psps
    basis = get_dftk_basis(parameters, state[:extra])

    atoms = basis.model.atoms
    @test length(atoms) == 1
    @test length(atoms) == 1
    @test all(at isa ElementPsp for (at, positions) in atoms)
    @test atoms[1][1].psp.identifier == "hgh/pbe/mg-q10.hgh"
    @test atoms[1][1].psp.Zion == 10

    @test basis.model.temperature ≈ 10 * austrip(1u"eV") atol=1e-8
    @test basis.model.smearing isa DFTK.Smearing.Gaussian
    @test basis.model.spin_polarization == :none

    xcterm = only(tt for tt in basis.model.term_types if tt isa Xc)
    @test length(xcterm.functionals) == 2
    @test xcterm.functionals[1] == :gga_x_pbe
    @test xcterm.functionals[2] == :gga_c_pbe

    kgrid = DFTK.kgrid_size_from_minimal_spacing(basis.model.lattice,
                                                 1 / 3.333austrip(1u"Å"))
    kcoords, ksymops = bzmesh_ir_wedge(kgrid, basis.model.symmetries)
    @test kcoords ≈ [kpt.coordinate for kpt in basis.kpoints]
end

@testset "Custom functional construction" begin
    kpts = [(0, 0, 0), (0.5, 0.7, 0.3)]
    state      = load_state("mg.json")
    parameters = Dict(state[:parameters]...,
                      :xc       => "PBE",
                      :kpts     => kpts,
                      :smearing => ("Fermi-Dirac", 5))
    basis = get_dftk_basis(parameters, state[:extra])

    @test basis.model.temperature ≈ 5 * austrip(1u"eV") atol=1e-8
    @test basis.model.smearing isa DFTK.Smearing.FermiDirac
    @test basis.model.spin_polarization == :none

    xcterm = only(tt for tt in basis.model.term_types if tt isa Xc)
    @test length(xcterm.functionals) == 2
    @test xcterm.functionals[1] == :gga_x_pbe
    @test xcterm.functionals[2] == :gga_c_pbe

    @test length(basis.kpoints) == 2
    @test Float64.(basis.kpoints[1].coordinate) ≈ [kpts[1]...]
    @test Float64.(basis.kpoints[2].coordinate) ≈ [kpts[2]...]
end

@testset "Test kpoint options" begin
    kpointoptions = [
        (ase=2.0,                length=18, kpt1=[0.0, 0.0, 0.0], kpt2=[ 1/5, 0.0, 0.0]),
        (ase=[2, 3, 4],          length=12, kpt1=[0.0, 0.0, 0.0], kpt2=[-0.5, 0.0, 0.0]),
        (ase=[2, 3, 4, "gamma"], length=6,  kpt1=[1/4, 0.0, 1/8], kpt2=[ 1/4, 1/3, 1/8]),
        (ase=[(0, 0.2, 0), (0.5, 0.2, 0.3)], length=2,
         kpt1=[0, 0.2, 0], kpt2=[0.5, 0.2, 0.3]),
    ]

    for params in kpointoptions
        state      = load_state("mg.json")
        parameters = Dict(state[:parameters]..., :kpts => params.ase)
        basis = get_dftk_basis(parameters, state[:extra])
        @test length(basis.kpoints) == params.length
        @test Float64.(basis.kpoints[1].coordinate) ≈ params.kpt1
        @test Float64.(basis.kpoints[2].coordinate) ≈ params.kpt2
    end
end

@testset "Test smearing options" begin
    smearingoptions = [
        (ase=("Fermi-Dirac", 10), temperature=10 * austrip(1u"eV"),
         smearing=DFTK.Smearing.FermiDirac),
        (ase=("Gaussian", 5), temperature=5 * austrip(1u"eV"),
         smearing=DFTK.Smearing.Gaussian),
        (ase=("Methfessel-Paxton", 5, 0), temperature=5 * austrip(1u"eV"),
         smearing=DFTK.Smearing.Gaussian),
        (ase=("Methfessel-Paxton", 5, 1), temperature=5 * austrip(1u"eV"),
         smearing=DFTK.Smearing.MethfesselPaxton1),
        (ase=("Methfessel-Paxton", 5, 2), temperature=5 * austrip(1u"eV"),
         smearing=DFTK.Smearing.MethfesselPaxton2),
    ]

    for params in smearingoptions
        state      = load_state("mg.json")
        parameters = Dict(state[:parameters]..., :smearing => params.ase)
        model = get_dftk_model(parameters, state[:extra])
        @test model.smearing isa params.smearing
        @test model.temperature ≈ params.temperature atol=1e-8
    end
end

@testset "Test mixing options" begin
    mixingoptions = [
        (ase="SimpleMixing()",               mixing=SimpleMixing, α=0.8),
        (ase="SimpleMixing(α=0.4)",          mixing=SimpleMixing, α=0.4),
        (ase="KerkerMixing(α=0.4, kTF=0.7)", mixing=KerkerMixing, α=0.4),
        (ase="KerkerMixing()",               mixing=KerkerMixing, α=0.8),
        (ase="DielectricMixing()",           mixing=DielectricMixing, α=0.8),
        (ase="DielectricMixing(α=0.2, kTF=0.7, εr=8.0)",
                                             mixing=DielectricMixing, α=0.2),
        # HybridMixing gets translated to DFTK.χ0Mixing internally
        (ase="HybridMixing()",               mixing=χ0Mixing, α=0.8),
        (ase="HybridMixing(α=0.2, kTF=0.7, εr=8.0)",
                                             mixing=χ0Mixing, α=0.2),
    ]

    for params in mixingoptions
        state      = load_state("mg.json")
        parameters = Dict(state[:parameters]..., :mixing => params.ase)
        mixing = get_dftk_mixing(parameters, state[:extra])
        @test (mixing isa params.mixing)
        @test mixing.α == params.α
    end
end

# TODO Missing checks for mandatory ASE parameters:
#      See https://wiki.fysik.dtu.dk/ase/development/calculators.html
#      and https://wiki.fysik.dtu.dk/ase/development/proposals/calculators.html
#      - charge

@testset "Silicon calculation" begin
    label = "silicon_dftk"
    ENERGY_PBE = -213.12688268374683  # eV
    FORCES_PBE = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]

    state = load_state("si.json")
    state[:parameters] = Dict(state[:parameters]...,
                               :xc => "PBE",
                               :kpts => (3, 3, 3),
                               :ecut => 190,
                               :scftol => 1e-4,
                               :nbands => 12)
    save_state(label * ".json", state)

    rm(state[:scfres], force=true)
    state = run_calculation(["energies", "forces"], label * ".json")
    @test state[:results][:energy] ≈ ENERGY_PBE atol=1e-4 rtol=1e-4
    @test state[:results][:forces] ≈ FORCES_PBE atol=1e-2

    orig_ene    = state[:results][:energy]
    orig_forces = state[:results][:forces]
    orig_atoms  = state[:atoms]

    save_state(label * ".json", state)
    state = load_state(label * ".json")

    # Check we got the old parameters and results back:
    @test state[:parameters][:xc]   == "PBE"
    @test state[:parameters][:kpts] == [3, 3, 3]
    @test state[:parameters][:ecut] == 190
    @test state[:results][:energy] ≈ orig_ene
    @test state[:results][:forces] ≈ orig_forces
    @test state[:atoms] == orig_atoms
end

max_stars_repo_name: mfherbst/ase-dftk

max_issues_repo_name: mfherbst/ase-dftk

max_forks_repo_name: mfherbst/asedftk


text:
mutable struct VMCgutzwiller
    model :: UnitCellQModel

    wavefunction :: GutzwillerSlater

    total_steps :: Int

    n_proposed :: Int
    n_accepted :: Int

    data :: Vector{Float64}
end

function report(sim::VMCgutzwiller)
    println("Finished with ", sim.total_steps)
    println("Number of proposed moves = ", sim.n_proposed)
    println("Number of accepted moves = ", sim.n_accepted)
    n_sites = numofsites(sim.model)
    Dict(zip([(i,j) for i=1:n_sites for j=i+1:n_sites],
             sim.data./sim.n_accepted))
end

function runVMC(model::UnitCellQModel,
                total_steps::Int)
    n_sites = numofsites(model)

    hmat = generatebdg(model)
    symmetrize = eltype(hmat) <: Complex ? Hermitian : Symmetric
    fact = eigen(symmetrize(hmat))

    n_occupied = div(n_sites, 2)
    println(fact.values[n_occupied-2:n_occupied+2])
    e_fermi = fact.values[n_occupied]
    n=1
    while fact.values[n_occupied+n] - e_fermi < 1.e-10
        n=n+1
    end
    states = fact.vectors[:, 1:n_occupied]

    wavefunction = random_gutzwiller_half(states)

    sim = VMCgutzwiller(model, wavefunction,
                        total_steps, 0, 0,
                        zeros(Float64, binomial(n_sites, 2)))

    while sim.n_accepted < sim.total_steps
        sim.n_proposed +=1
        if propose_step!(sim.wavefunction)
            sim.n_accepted += 1

            corrdata = measure(sim.wavefunction, :ZZ)
            sim.data += corrdata

            # display progress every 1000 accepted steps
            if sim.n_accepted % 1000 == 0
                check_and_update_gutzwiller!(sim.wavefunction)
                print("*")
            end
        end
    end
    println()
    report(sim)
end

function measure(wf::GutzwillerSlater, m::Symbol)
    n_sites = size(wf.states, 1)
    n_operators = binomial(n_sites, 2)
    correlations = zeros(Float64, n_operators)
    index = 1
    for i=1:n_sites, j=i+1:n_sites
        ni = wf.configuration[i] - 1/2
        nj = wf.configuration[j] - 1/2
        correlations[index] = ni * nj
        index += 1
    end
    correlations
end

max_stars_repo_name: amiragha/mps

max_issues_repo_name: amiragha/mps

max_forks_repo_name: amiragha/mps


text:
# matrix functions of dense matrices
# https://en.wikipedia.org/wiki/Matrix_function

# NOTE: for a matrix function f, the pushforward and pullback can be computed using the
# Fréchet derivative and its adjoint, respectively.
# https://en.wikipedia.org/wiki/Fréchet_derivative

# The pushforwards and pullbacks are related by matrix adjoints. If the pushforward of f(A)
# at A is (f_*)_A(ΔA), then the pullback at A is (f^*)_A(ΔY) = ((f_*)_A(ΔY'))'.
# If f has a power series representation with real coefficients, then this simplifies to
# (f^*)_Y(ΔY) = (f_*)_{A'}(ΔY)
# So we reuse the code from the pushforward to implement the pullback.

#####
##### interface function definitions
#####

"""
    _matfun(f, A) -> (Y, intermediates)

Compute the matrix function `Y=f(A)` for matrix `A`.
The function returns a tuple containing the result and a tuple of intermediates to be
reused by [`_matfun_frechet`](@ref) to compute the Fréchet derivative.
"""
_matfun

"""
    _matfun!(f, A) -> (Y, intermediates)

Similar to [`_matfun`](@ref), but where `A` may be overwritten.
"""
_matfun!

"""
    _matfun_frechet(f, E, A, Y, intermediates)

Compute the Fréchet derivative of the matrix function ``Y = f(A)`` at ``A`` in the direction
of ``E``, where `intermediates` is the second argument returned by [`_matfun`](@ref).

The Fréchet derivative is the unique linear map ``L_f \\colon E → L_f(A, E)``, such that
```math
L_f(A, E) = f(A + E) - f(A) + o(\\lVert E \\rVert).
```

[^Higham08]:
    > Higham, Nicholas J. Chapter 3: Conditioning. Functions of Matrices. 2008, 55-70.
    > doi: 10.1137/1.9780898717778.ch3
"""
_matfun_frechet

"""
    _matfun_frechet!(f, E, A, Y, intermediates)

Similar to [`_matfun_frechet`](@ref), but where `E` may be overwritten.
"""
_matfun_frechet!

"""
    _matfun_frechet_adjoint(f, E, A, Y, intermediates)

Compute the adjoint of the Fréchet derivative of the matrix function ``Y = f(A)`` at ``A``
in the direction of ``E``, where `intermediates` is the second argument returned by
[`_matfun`](@ref).

Given the Fréchet ``L_f(A, E)`` computed by [`_matfun_frechet`](@ref), then its adjoint
``L_f^⋆(A, E)`` is defined by the identity
```math
\\langle B, L_f(A, C) \\rangle = \\langle L_f^⋆(A, B), C \\rangle.
```
This identity is satisfied by ``L_f^⋆(A, E) = L_f(A, E')'``.

[^Higham08]:
    > Higham, Nicholas J. Chapter 3: Conditioning. Functions of Matrices. 2008, 55-70.
    > doi: 10.1137/1.9780898717778.ch3
"""
function _matfun_frechet_adjoint(f, E, A, Y, intermediates)
    E′ = E'
    # avoid passing an Adjoint to _matfun_frechet in case it can't handle it
    E′ = E′ isa Adjoint ? copy(E′) : E′
    LE = adjoint(_matfun_frechet(f, E′, A, Y, intermediates))
    # avoid returning an Adjoint
    return LE isa Adjoint ? copy(LE) : LE
end

"""
    _matfun_frechet_adjoint!(f, E, A, Y, intermediates)

Similar to [`_matfun_frechet_adjoint`](@ref), but where `E` may be overwritten.
"""
function _matfun_frechet_adjoint!(f, E, A, Y, intermediates)
    E′ = E'
    # avoid passing an Adjoint to _matfun_frechet in case it can't handle it
    E′ = E′ isa Adjoint ? copy(E′) : E′
    LE = adjoint(_matfun_frechet!(f, E′, A, Y, intermediates))
    # avoid returning an Adjoint
    return LE isa Adjoint ? copy(LE) : LE
end

#####
##### `exp`/`exp!`
#####

function frule((_, ΔA), ::typeof(LinearAlgebra.exp!), A::StridedMatrix{<:BlasFloat})
    if ishermitian(A)
        hermA = Hermitian(A)
        hermX, intermediates = _matfun(exp, hermA)
        ∂hermX = _matfun_frechet(exp, ΔA, hermA, hermX, intermediates)
        X = Matrix(hermX)
        ∂X = Matrix(∂hermX)
    else
        X, intermediates = _matfun!(exp, A)
        ∂X = _matfun_frechet!(exp, ΔA, A, X, intermediates)
    end
    return X, ∂X
end

function rrule(::typeof(exp), A0::StridedMatrix{<:BlasFloat})
    # TODO: try to make this more type-stable
    if ishermitian(A0)
        # call _matfun instead of the rrule to avoid hermitrizing ∂A in the pullback
        hermA = Hermitian(A0)
        hermX, hermX_intermediates = _matfun(exp, hermA)
        function exp_pullback_hermitian(ΔX)
            ∂hermA = _matfun_frechet_adjoint(exp, ΔX, hermA, hermX, hermX_intermediates)
            return NO_FIELDS, Matrix(∂hermA)
        end
        return Matrix(hermX), exp_pullback_hermitian
    else
        A = copy(A0)
        X, intermediates = _matfun!(exp, A)
        function exp_pullback(ΔX)
            # Ensures ∂X is mutable. The outer `adjoint` is unwrapped without copy by
            # the default _matfun_frechet_adjoint!
            ∂X = ChainRulesCore.is_inplaceable_destination(ΔX) ? ΔX : convert(Matrix, ΔX')'
            ∂A = _matfun_frechet_adjoint!(exp, ∂X, A, X, intermediates)
            return NO_FIELDS, ∂A
        end
        return X, exp_pullback
    end
end

## Destructive matrix exponential using algorithm from Higham, 2008,
## "Functions of Matrices: Theory and Computation", SIAM
## Adapted from LinearAlgebra.exp! with return of intermediates
## https://github.com/JuliaLang/julia/blob/f613b55/stdlib/LinearAlgebra/src/dense.jl#L583-L666
function _matfun!(::typeof(exp), A::StridedMatrix{T}) where {T<:BlasFloat}
    n = LinearAlgebra.checksquare(A)
    ilo, ihi, scale = LAPACK.gebal!('B', A)  # modifies A
    nA = opnorm(A, 1)
    Inn = Matrix{T}(I, n, n)
    ## For sufficiently small nA, use lower order Padé-Approximations
    if (nA <= 2.1)
        if nA > 0.95
            C = T[
                17643225600.0,
                8821612800.0,
                2075673600.0,
                302702400.0,
                30270240.0,
                2162160.0,
                110880.0,
                3960.0,
                90.0,
                1.0,
            ]
        elseif nA > 0.25
            C = T[17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0, 1.0]
        elseif nA > 0.015
            C = T[30240.0, 15120.0, 3360.0, 420.0, 30.0, 1.0]
        else
            C = T[120.0, 60.0, 12.0, 1.0]
        end
        si = 0
    else
        C = T[
            64764752532480000.0,
            32382376266240000.0,
            7771770303897600.0,
            1187353796428800.0,
            129060195264000.0,
            10559470521600.0,
            670442572800.0,
            33522128640.0,
            1323241920.0,
            40840800.0,
            960960.0,
            16380.0,
            182.0,
            1.0,
        ]
        s = log2(nA / 5.4)  # power of 2 later reversed by squaring
        si = ceil(Int, s)
    end

    if si > 0
        A ./= convert(T, 2^si)
    end

    A2 = A * A
    P = copy(Inn)
    W = C[2] * P
    V = C[1] * P
    Apows = typeof(P)[]
    for k in 1:(div(size(C, 1), 2) - 1)
        k2 = 2 * k
        P *= A2
        push!(Apows, P)
        W += C[k2 + 2] * P
        V += C[k2 + 1] * P
    end
    U = A * W
    X = V + U
    F = lu!(V - U)  # NOTE: use lu! instead of LAPACK.gesv! so we can reuse factorization
    ldiv!(F, X)
    Xpows = typeof(X)[X]
    if si > 0  # squaring to reverse dividing by power of 2
        for t in 1:si
            X *= X
            push!(Xpows, X)
        end
    end

    _unbalance!(X, ilo, ihi, scale, n)
    return X, (ilo, ihi, scale, C, si, Apows, W, F, Xpows)
end

# Application of the chain rule to exp!, also Algorithm 6.4 from
# Al-Mohy, Awad H. and Higham, Nicholas J. (2009).
# Computing the Fréchet Derivative of the Matrix Exponential, with an application to
# Condition Number Estimation", SIAM. 30 (4). pp. 1639-1657.
# http://eprints.maths.manchester.ac.uk/id/eprint/1218
function _matfun_frechet!(
    ::typeof(exp), ΔA, A::StridedMatrix{T}, X, (ilo, ihi, scale, C, si, Apows, W, F, Xpows)
) where {T<:BlasFloat}
    n = LinearAlgebra.checksquare(A)
    _balance!(ΔA, ilo, ihi, scale, n)

    if si > 0
        ΔA ./= convert(T, 2^si)
    end

    ∂A2 = mul!(A * ΔA, ΔA, A, true, true)
    A2 = first(Apows)
    # we will repeatedly overwrite ∂temp and ∂P below
    ∂temp = Matrix{eltype(∂A2)}(undef, n, n)
    ∂P = copy(∂A2)
    ∂W = C[4] * ∂P
    ∂V = C[3] * ∂P
    for k in 2:(length(Apows) - 1)
        k2 = 2 * k
        P = Apows[k - 1]
        ∂P, ∂temp = mul!(mul!(∂temp, ∂P, A2), P, ∂A2, true, true), ∂P
        axpy!(C[k2 + 2], ∂P, ∂W)
        axpy!(C[k2 + 1], ∂P, ∂V)
    end
    ∂U, ∂temp = mul!(mul!(∂temp, A, ∂W), ΔA, W, true, true), ∂W
    ∂temp .= ∂U .- ∂V
    ∂X = add!!(∂U, ∂V)
    mul!(∂X, ∂temp, first(Xpows), true, true)
    ldiv!(F, ∂X)

    if si > 0
        for t in 1:(length(Xpows) - 1)
            X = Xpows[t]
            ∂X, ∂temp = mul!(mul!(∂temp, X, ∂X), ∂X, X, true, true), ∂X
        end
    end

    _unbalance!(∂X, ilo, ihi, scale, n)
    return ∂X
end

#####
##### utils
#####

# Given (ilo, ihi, iscale) returned by LAPACK.gebal!('B', A), apply same balancing to X
function _balance!(X, ilo, ihi, scale, n)
    n = size(X, 1)
    if ihi < n
        for j in (ihi + 1):n
            LinearAlgebra.rcswap!(j, Int(scale[j]), X)
        end
    end
    if ilo > 1
        for j in (ilo - 1):-1:1
            LinearAlgebra.rcswap!(j, Int(scale[j]), X)
        end
    end

    for j in ilo:ihi
        scj = scale[j]
        for i in 1:n
            X[j, i] /= scj
        end
        for i in 1:n
            X[i, j] *= scj
        end
    end
    return X
end

# Reverse of _balance!
function _unbalance!(X, ilo, ihi, scale, n)
    for j in ilo:ihi
        scj = scale[j]
        for i in 1:n
            X[j, i] *= scj
        end
        for i in 1:n
            X[i, j] /= scj
        end
    end

    if ilo > 1
        for j in (ilo - 1):-1:1
            LinearAlgebra.rcswap!(j, Int(scale[j]), X)
        end
    end
    if ihi < n
        for j in (ihi + 1):n
            LinearAlgebra.rcswap!(j, Int(scale[j]), X)
        end
    end
    return X
end


text:
import FactCheck2TestSets: convert_line

@testset "convert_line" begin
    @test convert_line("asdf") == "asdf"

    s = """context("foo") do """
    @test convert_line(s) == "@testset \"foo\" begin"

    s = """facts("GeometryTypes") do"""
    @test convert_line(s) == "@testset \"GeometryTypes\" begin"

    s = """@fact_throws ndims(HyperCube)"""
    @test convert_line(s) == "@test_throws Exception ndims(HyperCube)"

    s = """@fact min_euclidean(c1,c2) --> roughly(5)"""
    @test convert_line(s) == "@test min_euclidean(c1,c2) ≈ 5"

    s = " @fact Inf --> not(isfinite)"
    @test convert_line(s) == " @test !(isfinite(Inf))"

    s = "  @fact make_checked_nan(S) --> true"
    @test convert_line(s) == "  @test make_checked_nan(S)"

    s = "@fact_throws ErrorException colorim(rand(UInt8, 3, 5, 3))"
    @test convert_line(s) == "@test_throws ErrorException colorim(rand(UInt8,3,5,3))"

    s = "@fact Images.ssd(imgs, (A.-mnA)/(mxA-mnA)) --> less_than(eps(UFixed16))"
    @test convert_line(s) == "@test Images.ssd(imgs,(A .- mnA) / (mxA - mnA)) < eps(UFixed16)"

    s = "for id in ids @fact corners[id] --> true end"
    @test convert_line(s) == "for id = ids\n    @test corners[id]\nend"

    s = "FactCheck.exitstatus()"
    @test convert_line(s) == ""

    s = """using FactCheck"""
    @test convert_line(s) == "if VERSION >= v\"0.5.0-dev+7720\"\n    using Base.Test\nelse\n    using BaseTestNext\n    const Test = BaseTestNext\nend"

    # s = "@fact abs(ovr[1, 2] - RGB{Float32}(a[1, 2], b[1, 2], a[1, 2])) --> roughly(0, atol=1e-5)"
    s = """@fact dimindex(imgds, "z") --> 0"""
    @test convert_line(s) == "@test dimindex(imgds,\"z\") == 0"

    s = """@fact update(l, 1, 1)   --> (leaf) -> hasindex(leaf[1], 1)"""
    @test convert_line(s) == "@test ((leaf->begin \n            hasindex(leaf[1],1)\n        end))(update(l,1,1))"

    s = "@fact PersistentVector([1])   --> not(isempty)"
    @test convert_line(s) == "@test !(isempty(PersistentVector([1])))"

    s = """    @fact_throws FileIO.WriterError save(Stream(format"BROKEN", STDOUT))"""
    @test convert_line(s) == "    @test_throws FileIO.WriterError save(Stream(format\"BROKEN\",STDOUT))"

end


text:
# Location in the input dictionary where the architecture specification
# can be found.
const _options_path_ = KeyChain(("mapper_options",))

# Don't rank input or output nodes.
isnonranking(t) = isinput(t) || isoutput(t)

struct PMConstructor{T <: Union{String,Dict}} <: MapConstructor
    file    ::T
    options ::NamedTuple

    #--inner constructor
    function PMConstructor(file::T, options::NamedTuple = NamedTuple()) where T
        # Iterate through each opion in "kwargs" - ensure it is in the list of
        # options provided by "_defult_options_"
        default_options = _get_default_options()
        for key in keys(options)
            if !haskey(default_options, key)
                error("Unrecognized Mapper Override option $(key).")
            end
        end
        return new{T}(file, options)
    end
end

"""
    Project Manager File Constructor

This controls dispatch to the constructor for the taskgraph and options parser.

The constructor will attach the following pieces of metadata to the constructed
taskgraph.

# Nodes

## Direct metadata - retrieved from Project Manager.
* `"type"` - The task type of this task given by the Project Manager.
    Preserves captalization and punctuation.

    Expected types: `String`

## Computed metadata - used for internal computation.
* `"mapper_type"` - Mapper specific type for this task. Generally, there
    will be a 1-to-1 mapping between Mapper types and Project Manager types,
    except that Mapper types will all lower case.

    Exceptions: Memories will be converted to either "memory_1port" or
    "memory_2port" depending on the number of neighbors they have. This ensures
    that memory tasks needing 2 ports will not be mapped to a single ported
    memory on Asap 4.

    Expected types: `String`

# Edges

## Direct metadata - retrieved from Project Manager
* `"source_task"` - The name of the source task for this edge.
    Expected types: `String`

* `"source_index"` - The original index in the Project Manager's datastructure
    where this edge originated. Used to disambiguate between multiple links
    and allows the Project Manager to reconstruct routings.

    Expected types: `String, Int`

* `"dest_task"` - The name of the destination task. Expected types: `String`.

* `"dest_index"` - Same as `source_index` except for destination.

* `"measurements_dict"` - Dictionary containing measurements. This is optional
    and will be marked as `missing` if not present.

* `"pm_class"` - The class of this link given by the Project Manager.

## Computed metadata
* `"mapper_class"` - The mapper equivalend of the given `pm_class`. Right now,
    it is just the lowercase version of `pm_class`.

* `"cost"` - Mapper assigned cost to the link. Currently, only assigns a higher
    weight to memory links.

* `"preserve_dest"` - Preserve the destination index. This keeps the mapper from
    swapping out destination fifos.
"""
function _get_default_options()
    return Dict(
        # General Options
        # ---------------
        :verbosity => "info",

        # Architecture Generation
        # -----------------------

        # Default to nothing. If this is not set in the input file, a run
        # time error will be generated somewhere.
        :architecture => nothing,

        # Default the number of links to 2 to match Asap3/4 architectures.
        :num_links => 2,

        # Mapping Options
        # ---------------
        # Number of times to retry place and route if congested.
        :num_retries => 3,

        # Set to true to use packet links during placement.
        :use_packet_links => false,

        # Set to "true" to assign weights to inter-task communication links
        # according to the normalized number of writes on that link.
        :use_profiled_links => false,

        # The rule set to use
        :ruleset => nothing,

        # Load Existing Maps
        # ------------------

        # Path (string) to an existing map. Must be compatible with the rest of
        # the options provided or bad things might happen.
        :existing_map => nothing,
      )
end

function Base.parse(c::PMConstructor{String})
    json_dict = open(c.file, "r") do f
        JSON.parse(f)
    end
    # Parse options from the project manager.
    final_options = parse_options(c.options, json_dict["mapper_options"])
    json_dict["mapper_options"] = final_options
    @debug """
        Final Options Dict:
        $final_options
    """
    return json_dict
end

function Base.parse(c::PMConstructor{<:Dict})
    dict = deepcopy(c.file)
    final_options = parse_options(c.options, dict["mapper_options"])
    dict["mapper_options"] = final_options
    return dict
end

"""
    parse_options(internal::NamedTuple, external::Dict)

Parse the options passed internally and externally. Prune all external options
that are not in the `default_options` dict and return a final options
dictionary with the following option precedences from highest to lowest:

`internal`, `external`, `default`.
"""
function parse_options(internal::NamedTuple, external::Dict = Dict{Symbol,Any}())
    # Convert the keys of 'external' to symbols for uniformity.
    external_sym = Dict(Symbol(k) => v for (k,v) in external)

    default_options = _get_default_options()

    # Generally, we don't know what could come in the externally parsed options
    # dictionary. Filter out unrecognized symbols to make sure that we only
    # have options that are implemented in AsapMapper.
    #
    # Delete all other keys.
    keys_to_delete = Symbol[]

    for key in keys(external_sym)
        if !haskey(default_options,key)
            @warn "Unrecognized Mapper option: $key."
            push!(keys_to_delete, key)
        end
    end
    for key in keys_to_delete
        delete!(external_sym, key)
    end

    # Merge all results together. Use the precedence in the `merge` operation to
    # get this correct.
    final_dict = Dict{Symbol,Any}(
        merge(default_options, external_sym, Dict(pairs(internal)))
    )

    # Do any global actions with side-effects here.
    parse_verbosity(final_dict[:verbosity])

    return final_dict
end

function parse_verbosity(verbosity)
    # Just redirect to "set_logging"
    if verbosity in ("debug", "info", "warn", "error")
        set_logging(Symbol(verbosity))
    else
        @warn "Unknown Verbosity option: $verbosity"
        set_logging(:info)
    end
end

function getrule(options)
    ruleset = get(options, :ruleset, nothing)
    if ruleset == nothing
        return KC()
    elseif isa(ruleset, AbstractKC)
        return ruleset
    else
        error("Unrecognized ruleset: $(ruleset)")
    end
end

################################################################################
# build_map - Top level function for creating Map objects from the Project
#   Manager.
################################################################################
function build_map(c::PMConstructor)
    # Parse the input json file
    json_dict = parse(c)
    a = build_architecture(c, json_dict)
    t = build_taskgraph(c, json_dict)

    # build the map and attach the options dictionary to it.
    options = json_dict[_options_path_]

    rule = getrule(options)
    m = Map(rule,a,t)
    m.options = options

    # Load an existing map if provided with one.
    if typeof(options[:existing_map]) <: String
        Mapper2.MapperCore.load(m, options[:existing_map])
    end

    return m
end

# Placement for Asap2
function asap_pnr(m::Map{D,Asap2}; kwargs...) where {D}
    for _ in 1:m.options[:num_retries]
        # Build the aux storage for this type.
        aux = build_aux(m)
        place!(m, aux = aux; kwargs...)
        route!(m)
        check_routing(m; quiet = true) && break
    end
    return m
end

function asap_pnr(m::Map{A,D}) where {A,D}
    for _ in 1:m.options[:num_retries]
        place!(m)
        route!(m)
        check_routing(m; quiet = true) && break
    end

    return m
end


text:
# calibration readouts
# --------------------
bucketsize = 1.0 # calibration bucket size in liters
solution = 1.0 # calibration solution concentration (g/l)
# calis: total calibration ml solution, 1st column, vs sensor readout (μS/cm), 2nd column
calis = Dict( 309 => (# [0  1                    # 23.06.20 (only a test before going to the field)
                      # 1  2
                      # 6  10
                      # 16 25],
                      [ 0  1                    # 05.08.20
                        1 2
                        6 8
                        16 22
                        36 47
                        76 90
                        116 131
                        166 191
                        191 229],
                      [ 0 1                     # 10.08.20
                        1 2
                        6 8
                        16 27
                        36 48
                        76 95
                        116 144],
                      [ 0 1                     # 13.08.20
                        1 2
                        6 8
                        16 26
                        36 53
                        76 98
                        116 143
                        156 186
                        256 345],  # disturbs linear fit quite a lot
                      [ 0  1                    # 20.08.20
                        2  3
                        12 19
                        42 57
                        102 121
                        162 182
                        222 251
                        262 270]),
              265 => (# [ 0 1                     # 23.06.20  (only a test before going to the field)
                      #   1 3
                      #   6 11
                      #   16 26],
                      [ 0  1                    # 05.08.20
                        1 2
                        6 8
                        16 23
                        36 47
                        76 91
                        116 133
                        166 193
                        191 235],
                      [ 0 1                     # 10.08.20
                        1 2
                        6 8
                        16 27
                        36 54
                        76 100
                        116 147],
                      [ 0 1                     # 13.08.20
                        1 2
                        6 8
                        16 26
                        36 54
                        76 99
                        116 140
                        156 188
                        256 352],  # disturbs linear fit quite a lot
                      [ 0  1                    # 20.08.20
                        2  3
                        12 19
                        42 57
                        102 122
                        162 183
                        222 252
                        262 267]),
              145 => ([0  1                     # 23.06.20  (only a test before going to the field, but not enough data points for a proper calibration anyway)
                        1  2
                        6  10
                        16 25],
                       [0  0                     # 20.08.20
                        2  3
                        12 18
                        42 57
                        102 121
                        162 182
                        222 251])
            )

# fit a linear function through calibration readouts
ctd309["cond2conc"] = Rhone2020.fit_calibration(bucketsize, solution, e_cond, calis[309][1], calis[309][2], calis[309][3], calis[309][4])
ctd265["cond2conc"] = Rhone2020.fit_calibration(bucketsize, solution, e_cond, calis[265][1], calis[265][2], calis[265][3], calis[265][4])
ctd145["cond2conc"] = Rhone2020.fit_calibration(bucketsize, solution, e_cond, calis[145][1], calis[145][2])

# temperature calibration (see calibration_plots.md)
# only possible for lab calibrations of CTD-309 and CTD-265; take mean value between Christophe and Annegret
T309_shift = -0.035
T265_shift = -0.335
T145_shift = 0.033

for date in keys(indices)
  ctd309[date][:temp_PT] = ctd309[date][:temp_PT] .- T309_shift
  ctd265[date][:temp_PT] = ctd265[date][:temp_PT] .- T265_shift
  if haskey(ctd145, date)
    ctd145[date][:temp_PT] = ctd145[date][:temp_PT] .- T145_shift
  end
end


text:
const!(builder, x::Union{AbstractArray,XScalar,Bool}) = xlaclient.ops.Constant(builder, x)
const!(builder, xs::Tuple) = build!(builder, XTuple(), const!.((builder,), xs)...)

struct Lambda
  vars::Vector{Any}
  func::IR
end

for op in :[Neg, Sign, Floor, Ceil, Round, Exp, Log, Expm1, Log1p, Tanh,
            Sin, Cos, Lgamma, Digamma, Erf, Erfc, ErfInv, Sqrt, Rsqrt, Not].args
  @eval begin
    struct $op end
    build!(builder, ::$op, x) = xlaclient.ops.$op(x)
  end
end

for op in :[Atan2, Pow, And, Or, Xor, Add, Sub, Mul, SafeMul, Div, Rem, Dot,
            Max, Min, ShiftLeft, ShiftRightArithmetic, ShiftRightLogical,
            Gt, Ge, Lt, Le].args
  @eval begin
    struct $op end
    build!(builder, ::$op, x, y) = xlaclient.ops.$op(x, y)
  end
end

struct Rev
  dims
end

build!(builder, op::Rev, x) =
  xlaclient.ops.Rev(x, op.dims .- 1)

struct Conv
  strides
  padding
  lhs_dilation
  rhs_dilation
end

function build!(builder, op::Conv, x, w)
  dims = xlaclient.make_convolution_dimension_numbers(("WHCN", "WHIO", "WHCN"), 2)
  xlaclient.ops.ConvGeneralDilated(x, w, op.strides, op.padding, op.lhs_dilation, op.rhs_dilation, dims)
end

struct DynamicSlice
  size
end

function build!(builder, slice::DynamicSlice, x, start...)
  one = const!(builder, 1)
  start = map(i -> xlaclient.ops.Sub(i, one), start)
  xlaclient.ops.DynamicSlice(x, start, slice.size)
end

struct BroadcastInDim
  size
  map
end

build!(builder, bc::BroadcastInDim, x) = xlaclient.ops.BroadcastInDim(x, bc.size, bc.map.-1)

alldims(builder, xs) = [0:ndims(shapeof(builder, xs))-1;]

struct Map end

function build!(builder, ::Map, f, args...)
  settypes!(builder, f, args..., with = eltype)
  xlaclient.ops.Map(builder, args, build(f), alldims(builder, args[1]))
end

struct Reduce
  dims
end

function build!(builder, op::Reduce, f, xs, init)
  settypes!(builder, f, xs, xs, with = eltype)
  dims = op.dims == (:) ? alldims(builder, xs) : sort([op.dims...].-1)
  xlaclient.ops.Reduce(builder, [xs], [init], build(f), dims)
end

struct ReduceWindow
  window
  stride
  base_dilation
  window_dilation
  padding
end

ReduceWindow(window) = ReduceWindow(window, window, one.(window), one.(window), map(_ -> (0, 0), window))

function build!(builder, op::ReduceWindow, f, xs, init)
  xlaclient.ops.ReduceWindowWithGeneralPadding(
    xs, init, build(f), op.window, op.stride,
    op.base_dilation, op.window_dilation, op.padding)
end

struct SelectAndScatter
  window
  stride
  padding
end

SelectAndScatter(window) = SelectAndScatter(window, window, map(_ -> (0, 0), window))

function build!(builder, op::SelectAndScatter, select, scatter, xs, source, init)
  xlaclient.ops.SelectAndScatterWithGeneralPadding(
    xs, build(select), op.window, op.stride, op.padding, source, init, build(scatter))
end

struct XTuple end

build!(builder, ::XTuple, xs...) = xlaclient.ops.Tuple(builder, xs)

struct GetTupleElement
  idx::Int
end

build!(builder, op::GetTupleElement, x) = xlaclient.ops.GetTupleElement(x, op.idx)

struct ConvertElementType
  to::Type{<:XScalar}
end

build!(builder, op::ConvertElementType, x) =
  xlaclient.ops.ConvertElementType(x, primitivetype(op.to))

struct Reshape
  dims
  size
end

build!(builder, op::Reshape, x) =
  xlaclient.ops.Reshape(x, op.dims.-1, op.size)

struct Conditional end

function build!(builder, ::Conditional, pred,
                true_operand, true_computation,
                false_operand, false_computation)
  settypes!(builder, true_computation, true_operand)
  settypes!(builder, false_computation, false_operand)
  xlaclient.ops.Conditional(pred, true_operand, build(true_computation),
                            false_operand, build(false_computation))
end

struct While end

function build!(builder, ::While, condition, body, init)
  settypes!(builder, condition, init)
  settypes!(builder, body, init)
  xlaclient.ops.While(build(condition), build(body), init)
end


text:
# random ordering
@mesh(HEXMESH, n)
random_elements([n,n,n]);
random_nodes();

solve([q1,q2,q3,q4,q5,q6]);

regtime = Base.Libc.time();
for iter=1:times
    solve([q1,q2,q3,q4,q5,q6]);
end
regtime = Base.Libc.time() - regtime;
regtime /= times;
timings[3] = regtime;
println("Finished random");


text:
"""
    initialise(chr; kwargs...)

Initialise drawing and return a `Dict` containing properties.
"""
function initialise(chr; drawingsize = "A0landscape",
        outfile = "genomap.svg", nbreaks = 30, textangle = -pi/6,
        numberstyle = 'k', defaultcolour = "gray75", arrowwidth = 12,
        genetextsize = 15, colourfunction = x -> x.feature, colourmap = Dict(),
        features = ["CDS", "rRNA", "tRNA"], legend = :categorical,
        legend_high = "1.0", legend_low = "0.0",
        rightlimit = 0.95, leftlimit = 0.01, annotate = true,
        extrafunction = (p) -> nothing,
        genetextfunction = x -> get(x, :gene, ""))
    (xmax, ymax) = drawingdimensions(drawingsize)
    drawing = Drawing(xmax, ymax, outfile)
    background("white")
    p = Dict()
    p[:genetextsize] = genetextsize
    p[:chromosome] = chr
    p[:drawing] = drawing
    p[:outfile] = outfile
    p[:interval] = Int(round(length(chr.sequence) / nbreaks, RoundUp))
    p[:textangle] = textangle
    p[:annotate] = annotate
    breaks = 1:p[:interval]:length(chr.sequence)
    left = Point[]
    right = Point[]
    p[:rightlimit] = rightlimit
    p[:leftlimit] = leftlimit
    @inbounds for i in 1:length(breaks)
        push!(left,  Point(p[:leftlimit] * xmax, i * ymax / (length(breaks) + 1)))
        push!(right, Point(p[:rightlimit] * xmax, i * ymax / (length(breaks) + 1)))
    end
    lastpoint = between(left[end], right[end], (length(chr.sequence) - breaks[end]) / (p[:interval] - 1))
    p[:breaks] = breaks
    p[:left] = left
    p[:right] = right
    p[:lastpoint] = lastpoint
    p[:numberstyle] = numberstyle
    p[:offset_y] = ymax / 150
    p[:defaultcolour] = defaultcolour
    p[:arrowwidth] = arrowwidth
    p[:xmax] = xmax
    p[:ymax] = ymax
    p[:legend] = legend
    p[:legend_high] = legend_high
    p[:legend_low] = legend_low
    p[:genetextoffset] = -p[:arrowwidth] - ymax / 20nbreaks
    p[:features] = features
    p[:colourfunction] = colourfunction
    if !isempty(colourmap)
        p[:colourmap] = colourmap
    else
        p[:colourmap] = generatecolours(p)
    end
    ### genetextfunction can be either a function that, given a Gene, returns
    # the text to plot above the gene, or a Symbol showing which property of the
    # gene to plot
    if genetextfunction isa Symbol
        p[:genetextfunction] = g -> getproperty(g, genetextfunction, "")
    else
        p[:genetextfunction] = genetextfunction
    end
    p[:extrafunction] = extrafunction
    return p
end

function generatecolours(p)
    chr = p[:chromosome]
    features = p[:features]
    plottedgenes = @genes(chr, :feature in features)
    categories = unique(vcat(p[:colourfunction].(plottedgenes)...))
    colours = Dict()
    cmap = Colors.distinguishable_colors(length(categories))
    for i in eachindex(categories)
        colours[categories[i]] = cmap[i]
    end
    return colours
end