text:
"""$(README)

$(EXPORTS)
"""
module ActinRingsMC

export SystemParams, SimulationParams, System, Lattice
export calc_lf, calc_Lf, calc_max_lattice_height, calc_min_lattice_height, calc_radius
export generate_starting_config, update_occupancies!, run!, run_us!

using DelimitedFiles
using DocStringExtensions
using JSON
using Random

const kb = 1.380649e-23

"""
System parameters.

$(TYPEDFIELDS)
"""
struct SystemParams
    """Dissociation constant for single crosslinker binding"""
    ks::Float64
    """Dissociation constant for double crosslinker binding"""
    kd::Float64
    """Temperature (K)"""
    T::Float64
    """Lattice spacing (m)"""
    delta::Float64
    """Crosslinker concentration (M)"""
    Xc::Float64
    """Bending rigidity (N m^2)"""
    EI::Float64
    """Filament length (m)"""
    Lf::Float64
    """Binding sites per filament"""
    lf::Int
    """Total number of filaments"""
    Nfil::Int
    """Number of scaffold filaments"""
    Nsca::Int
end

"""
Simulation parameters.

$(TYPEDFIELDS)
"""
struct SimulationParams
    """Number of iterations for US"""
    iters::Int
    """MC steps for a run or single iteration"""
    steps::Int
    """Maximum change in bias energy (kbT)"""
    max_bias_diff::Float64
    """Number of steps between writing to output files"""
    write_interval::Int
    """Frequency of radius moves (the remainder are filament translation moves)"""
    radius_move_freq::Float64
    """Output filebase (include filepath)"""
    filebase::String
    """Use analytical model to generation starting biases"""
    analytical_biases::Bool
    """Read starting biases from file"""
    read_biases::Bool
    """File to read biases from"""
    biases_filename::String
    """Iteration to to start on"""
    restart_iter::Int
    """Width of bins for lattice height bias (set to 1 for no binning)"""
    binwidth::Int
end

mutable struct Filament
    lf::Int
    using_current::Bool
    coors::Array{Int,2}
    current_coors::Array{Int,2}
    trial_coors::Array{Int,2}
    index::Int
end

mutable struct System
    parms::SystemParams
    filaments::Vector{Filament}
    radius::Float64
    energy::Float64
end

"""
$(TYPEDSIGNATURES)

System data, including parameters and positions of filaments.

The given radius should be that corresponding to the lattice height used to create the
lattice and the starting configuration of the filaments.
"""
function System(parms::SystemParams, filaments::Vector{Filament}, radius::Float64)
    return System(parms, filaments, radius, 0)
end

mutable struct Lattice
    occupancy::Dict{Vector{Int},Tuple{Int,Int}}
    using_current::Bool
    current_occupancy::Dict{Vector{Int},Tuple{Int,Int}}
    trial_occupancy::Dict{Vector{Int},Tuple{Int,Int}}
    height::Int
    current_height::Int
    trial_height::Int
    max_height::Int
    min_height::Int
end

"""
$(TYPEDSIGNATURES)

2D lattice with periodic conditions on y.

The max and min heights should be calculated based on the number of scaffold filaments and
the length of the filaments.
"""
function Lattice(height::Int, max_height::Int, min_height::Int)
    occupancy::Dict{Vector{Int},Tuple{Int,Int}} = Dict()
    return Lattice(
        occupancy,
        true,
        occupancy,
        copy(occupancy),
        height,
        height,
        height,
        max_height,
        min_height
)
end

mutable struct Biases
    barriers::Vector{Int}
    counts::Vector{Int}
    freqs::Vector{Float64}
    probs::Vector{Float64}
    enes::Vector{Float64}
    numbins::Int
    binwidth::Int
end

function create_bin_barriers(lattice::Lattice, numbins::Int)
    binsize = div((lattice.max_height - lattice.min_height + 1), numbins)
    barriers::Vector{Int} = []
    barrier = lattice.min_height
    for _ in 1:(numbins - 1)
        barrier += binsize
        push!(barriers, barrier)
    end

    return barriers
end

function Biases(lattice::Lattice, binwidth::Int)
    numstates = lattice.max_height - lattice.min_height + 1
    numbins = div(numstates, binwidth)
    barriers = create_bin_barriers(lattice, numbins)

    return Biases(
        barriers,
        zeros(numbins),
        zeros(numbins),
        zeros(numbins),
        zeros(numbins),
        numbins,
        binwidth
   )
end

function read_biases!(filename::String, itr::Int, biases::Biases)
    biases.enes = readdlm(filename, skipstart=1)[itr, :]
end

function get_bin(barriers::Vector{Int}, height::Int)
    for (bin, barrier) in enumerate(barriers)
        if height < barrier
            return bin
        end
    end

    return length(barriers) + 1
end

"""
$(TYPEDSIGNATURES)

Calculate the number of sites per filament.
"""
calc_lf(Lf::Float64, delta::Float64)::Int = div(Lf, delta)

"""
$(TYPEDSIGNATURES)

Calculate the length of a filament.
"""
calc_Lf(lf::Int, delta::Float64) = delta * lf

"""
$(TYPEDSIGNATURES)

Calculate the maximum allowable height of a lattice given Nsca.
"""
calc_max_lattice_height(Nsca::Int, lf::Int) = Nsca * lf - Nsca - 1

"""
$(TYPEDSIGNATURES)

Calculate the minimum allowable height of a lattice given Nsca.
"""
calc_min_lattice_height(Nsca::Int, lf::Int) = div(Nsca, 2) * lf - 1

"""
$(TYPEDSIGNATURES)

Calculate the radius of a ring for a given lattice height.
"""
calc_radius(delta::Float64, height::Int) = delta * (height + 1) / (2pi)

"""Update the radius/height variables in the system and lattice."""
function update_radius!(system::System, lattice::Lattice, height::Int)
    lattice.height = height
    system.radius = calc_radius(system.parms.delta, height)

    return nothing
end

"""Use the current coordinates and occupancies."""
function use_current_coors!(system::System, lattice::Lattice)
    lattice.using_current = true
    lattice.occupancy = lattice.current_occupancy
    update_radius!(system, lattice, lattice.current_height)
    for filament in system.filaments
        filament.using_current = true
        filament.coors = filament.current_coors
    end

    return nothing
end

"""Use the trial coordinates and occupancies."""
function use_trial_coors!(system::System, lattice::Lattice)
    lattice.using_current = false
    lattice.occupancy = lattice.trial_occupancy
    update_radius!(system, lattice, lattice.trial_height)
    for filament in system.filaments
        filament.using_current = false
        filament.coors = filament.trial_coors
    end

    return nothing
end

"""Replace the current coordinates/occupancies with the trial ones."""
function accept_trial!(filament::Filament, lattice::Lattice)
    filament.current_coors = copy(filament.trial_coors)
    lattice.current_occupancy = copy(lattice.trial_occupancy)

    return nothing
end

function accept_trial!(system::System, lattice::Lattice)
    for filament in system.filaments
        filament.current_coors = copy(filament.trial_coors)
    end

    lattice.current_occupancy = copy(lattice.trial_occupancy)
    lattice.current_height = lattice.trial_height

    return nothing
end

"""Replace the trial coordinates/occupancies with the current ones."""
function accept_current!(filament::Filament, lattice::Lattice)
    filament.trial_coors = copy(filament.current_coors)
    lattice.trial_occupancy = copy(lattice.current_occupancy)

    return nothing
end

function accept_current!(system::System, lattice::Lattice)
    for filament in system.filaments
        filament.trial_coors = copy(filament.current_coors)
    end

    lattice.trial_occupancy = copy(lattice.current_occupancy)
    lattice.trial_height = lattice.current_height

    return nothing
end

"""Apply periodic boundary conditions to position."""
function wrap_pos!(lattice::Lattice, pos::Vector{Int})
    if pos[2] > lattice.height
        pos[2] = -(lattice.height + 1) + pos[2]
    elseif pos[2] < 0
        pos[2] = (lattice.height + 1) + pos[2]
    end

    return nothing
end

"""Distance between two points on the lattice along y given direction."""
function periodic_distance(lattice::Lattice, y1::Int, y2::Int, dir::Int)
    d = 0
    if y2 > y1
        if dir == 1
            d = y2 - y1
        else
            d = -(lattice.height - y2 + y1 + 1)
        end
    elseif y2 < y1
        if dir == 1
            d = lattice.height - y1 + y2 + 1
        else
            d = y2 - y1
        end
    end

    return d
end

function system_in_boundaries(system::System, lattice::Lattice)
    for filament in system.filaments
        if any(filament.coors[2, :] .< 0) || any(filament.coors[2, :] .> lattice.height)
            return false
        end
    end

    return true
end

"""
$(TYPEDSIGNATURES)

Generate a starting configuration with uniform overlaps.

The number of scaffold filament Nsca and the number of lattices lf must be even.
"""
function generate_starting_config(
    lattice::Lattice,
    Nfil::Int,
    Nsca::Int,
    lf::Int,
    overlap::Int
)
    if Nsca % 2 != 0
        throw(DomainError("Can currently only handle even Nsca"))
    end
    if lf % 2 != 0
        throw(DomainError("Uniform overlaps only possible with even lf"))
    end

    filaments::Vector{Filament} = []
    pos = [0, 0]
    Ni = 1
    while Ni <= Nfil
        for _ in 1:div(Nsca, 2)
            coors = zeros(2, lf)
            for j in 1:lf
                coors[:, j] = pos
                pos[2] += 1
                wrap_pos!(lattice, pos)
            end
            filament = Filament(lf, true, coors, coors, copy(coors), Ni)
            push!(filaments, filament)
            pos[2] += lf - 2 * overlap
            wrap_pos!(lattice, pos)
            Ni += 1
            if Ni > Nfil
                break
            end
        end

        pos[1] += 1
        pos[1] % 2 == 0 ? pos[2] = 0 : pos[2] = lf - overlap
        wrap_pos!(lattice, pos)
    end

    return filaments
end

"""
$(TYPEDSIGNATURES)

Clear occupancies and fully update.
"""
function update_occupancies!(filaments::Vector{Filament}, lattice::Lattice)
    lattice.current_occupancy = Dict()
    lattice.trial_occupancy = Dict()
    for filament in filaments
        for i in 1:filament.lf
            pos = filament.coors[:, i]
            lattice.current_occupancy[copy(pos)] = (filament.index, i)
            lattice.trial_occupancy[copy(pos)] = (filament.index, i)
        end
    end

    lattice.occupancy = lattice.current_occupancy

    return nothing
end

"""Calculate the bending energy (J) for a single filament."""
function filament_bending_energy(system::System)
    return system.parms.EI * system.parms.Lf / (2 * system.radius^2)
end

function filament_bending_energy(system::System, radius::Float64)
    return system.parms.EI * system.parms.Lf / (2 * radius^2)
end

"""Calculate the bending energy (J) for the whole system."""
function total_bending_energy(system::System)
    ene = 0
    for _ in 1:system.parms.Nfil
        ene += filament_bending_energy(system)
    end

    return ene
end

"""Calculate the overlap energy (J)."""
function overlap_energy(system::System, L::Float64)
    factor_t = L * kb * system.parms.T / system.parms.delta
    ks = system.parms.ks
    kd = system.parms.kd
    Xc = system.parms.Xc
    log_t = log(1 + ks^2 * Xc / (kd * (ks + Xc)^2))

    return -factor_t * log_t
end

function overlap_energy(system::System, lattice::Lattice, filament::Filament)
    l = 0
    for i in 1:system.parms.lf
        pos = filament.coors[:, i]
        for dx in [-1, 1]
            adj_pos = [pos[1] + dx, pos[2]]
            if adj_pos in keys(lattice.occupancy)
                l += 1
            end
        end
    end
    L = system.parms.delta * l

    return overlap_energy(system, L)
end

function overlap_energy(system::System, lattice::Lattice)
    ene = 0
    for filament in system.filaments
        ene += overlap_energy(system, lattice, filament) / 2
    end

    return ene
end

"""Calculate the bias energy (J)."""
function bias_energy(lattice::Lattice, biases::Biases)
    return biases.enes[get_bin(biases.barriers, lattice.height)]
end

"""Calculate total energy (J)."""
function total_energy(system::System, lattice::Lattice)
    ene = 0
    for filament in system.filaments
        ene += overlap_energy(system, lattice, filament) / 2
        ene += filament_bending_energy(system)
    end

    return ene
end

function total_energy(system::System, lattice::Lattice, biases::Biases)
    ene = total_energy(system, lattice)
    ene += bias_energy(lattice, biases)

    return ene
end

"""Calculate total energy difference (J)."""
function energy_diff(system::System, lattice::Lattice, filament::Filament)
    using_current = lattice.using_current
    if !using_current
        use_current_coors!(system, lattice)
    end

    cur_overlap_ene = overlap_energy(system, lattice, filament)
    cur_bending_ene = filament_bending_energy(system)

    use_trial_coors!(system, lattice)
    trial_overlap_ene = overlap_energy(system, lattice, filament)
    trial_bending_ene = filament_bending_energy(system)

    if using_current
        use_current_coors!(system, lattice)
    end

    return trial_overlap_ene + trial_bending_ene - cur_overlap_ene - cur_bending_ene
end

function energy_diff(system::System, lattice::Lattice, biases::Biases)
    using_current = lattice.using_current
    if !using_current
        use_current_coors!(system, lattice)
    end

    current_ene = total_energy(system, lattice, biases)

    use_trial_coors!(system, lattice)
    trial_ene = total_energy(system, lattice, biases)

    if using_current
        use_current_coors!(system, lattice)
    end

    return trial_ene - current_ene
end

"""Check if system connected and ring is unbroken with correct Nsca."""
function ring_and_system_connected(system::System, lattice::Lattice, check_consistency=false)
    ring_connected = []
    for filament in system.filaments
        push!(ring_connected, ring_and_system_connected(system, lattice, filament))
        if !check_consistency
            break
        end
    end
    if all(ring_connected)
        return true
    elseif all(.!ring_connected)
        return false
    else
        throw(ErrorException("Inconsistent connectivity check"))
    end
end

function ring_and_system_connected(system::System, lattice::Lattice, filament::Filament)
    # if system.parms.Nfil == system.parms.Nsca == 2
    #    first_pos = filament.coors[:, 1]
    #    last_pos = filament.coors[:, end]
    #    if (([first_pos[1] - 1, first_pos[2]] in keys(lattice.occupancy) ||
    #        [first_pos[1] + 1, first_pos[2]] in keys(lattice.occupancy)) &&
    #        ([last_pos[1] - 1, last_pos[2]] in keys(lattice.occupancy) ||
    #        [last_pos[1] + 1, last_pos[2]] in keys(lattice.occupancy)))
    #        return true
    #    else
    #        return false
    #    end
    # end
    searched_filaments::Set{Int} = Set()
    pos = filament.coors[:, 1]
    site_i = 1
    path::Vector{Vector{Int}} = [[], [], []]
    path_length = 0
    ring_contig = false
    connected_filaments::Set{Int} = Set(filament.index)
    Nsca::Int = system.parms.Nfil
    while site_i <= filament.lf
        for dx in [-1, 1]
            adj_pos = [pos[1] + dx, pos[2]]
            if adj_pos in keys(lattice.occupancy)
                adj_filament_i, adj_site_i = lattice.occupancy[adj_pos]
                if adj_filament_i in searched_filaments
                    continue
                end
                union!(searched_filaments, adj_filament_i)
                union!(connected_filaments, adj_filament_i)
                adj_filament = system.filaments[adj_filament_i]
                push!(path[1], filament.index)
                push!(path[2], site_i)
                push!(path[3], path_length)
                ring_contig, connected_filaments, Nsca = search_filament_for_path(
                    system,
                    lattice,
                    adj_filament,
                    adj_site_i,
                    path,
                    ring_contig,
                    connected_filaments,
                    Nsca
                )
                if (ring_contig &&
                    length(connected_filaments) == system.parms.Nfil &&
                    Nsca == system.parms.Nsca)

                    return true
                end
                pop!(path[1])
                pop!(path[2])
                pop!(path[3])
            end
        end
        site_i += 1
        pos += [0, 1]
        wrap_pos!(lattice, pos)
        path_length += 1
    end

    return false
end

function path_completed(
    lattice::Lattice,
    path_length::Int,
    entry_site::Int,
    exit_site::Int,
    )
    remainder_dist = exit_site - entry_site
    # println((entry_site, exit_site, remainder_dist, path_length, abs(path_length + remainder_dist), lattice.height + 1))

    return abs(path_length + remainder_dist) == lattice.height + 1
end

function search_filament_for_path(
    system::System,
    lattice::Lattice,
    filament::Filament,
    site_i::Int,
    path::Vector{Vector{Int}},
    ring_contig::Bool,
    connected_filaments::Set{Int},
    Nsca::Int
)
    searched_filaments::Set{Int} = Set()
    initial_pos = filament.coors[:, site_i]
    path_length = path[3][end]
    initial_path_length = path_length
    pos = copy(initial_pos)
    dir = -1
    initial_site_i = site_i
    # println(filament.index)
    # println(path)
    while true

        # Check if end of filament reached
        if site_i < 1 || site_i > filament.lf
            if dir == -1
                dir = 1
                path_length = initial_path_length + dir
                site_i = initial_site_i + dir
                pos = initial_pos + [0, dir]
                wrap_pos!(lattice, pos)
                continue
            else
                break
            end
        end

        for dx in [-1, 1]
            adj_pos = [pos[1] + dx, pos[2]]
            if adj_pos in keys(lattice.occupancy)
                # println(adj_pos)
                adj_filament_i, adj_site_i = lattice.occupancy[adj_pos]
                adj_filament = system.filaments[adj_filament_i]

                # This will check the neighbouring filaments at every site
                # This seems inefficient, but then I can't skip after one check if Nsca = 2
                if adj_filament_i in path[1]
                    path_i = indexin(adj_filament_i, path[1])[1]
                    sub_path_length = path_length - path[3][path_i]
                    exit_site = path[2][path_i]
                    if path_completed(lattice, sub_path_length, adj_site_i, exit_site)
                        ring_contig = true
                        Nsca = min(Nsca, length(path[1]) - path_i + 2)
                    end
                end

                if (ring_contig &&
                    length(connected_filaments) == system.parms.Nfil &&
                    Nsca == system.parms.Nsca)

                    return (ring_contig, connected_filaments, Nsca)
                end

                if (adj_filament_i in searched_filaments ||
                    adj_filament_i in path[1])

                    continue
                end

                union!(connected_filaments, adj_filament_i)
                union!(searched_filaments, adj_filament_i)
                push!(path[1], filament.index)
                push!(path[2], site_i)
                push!(path[3], path_length)
                ring_contig, connected_filaments, Nsca = search_filament_for_path(
                    system,
                    lattice,
                    adj_filament,
                    adj_site_i,
                    path,
                    ring_contig,
                    connected_filaments,
                    Nsca
                )
                if (ring_contig &&
                    length(connected_filaments) == system.parms.Nfil &&
                    Nsca == system.parms.Nsca)

                    return (ring_contig, connected_filaments, Nsca)
                end
                pop!(path[1])
                pop!(path[2])
                pop!(path[3])
                # println(filament.index)
                # println(path)
            end
        end

        site_i += dir
        pos += [0, dir]
        wrap_pos!(lattice, pos)
        path_length += dir
    end

    return (ring_contig, connected_filaments, Nsca)
end

"""Check filaments are contiguous, i.e. on adjacent lattice sites."""
function filaments_contiguous(system::System, lattice::Lattice)
    for filament in system.filaments
        prev_pos = filament.coors[:, 1]
        for i in 2:filament.lf
            pos = filament.coors[:, i]
            exp_pos = [prev_pos[1], prev_pos[2] + 1]
            wrap_pos!(lattice, exp_pos)
            if exp_pos != pos
                return false
            end
            prev_pos = pos
        end
    end

    return true
end

"""Test acceptance with Metropolis criterion."""
function accept_move(system::System, delta_energy::Float64, mult::Float64=1.)
    p_accept = min(1, mult * exp(-delta_energy / kb / system.parms.T))
    accept = false
    if p_accept == 1 || p_accept > rand(Float64)
        accept = true;
    end

    return accept;
end

function recenter!(system::System, lattice::Lattice)
    origin = system.filaments[1].coors[2, 1]
    if origin == 0
        return nothing
    end

    for filament in system.filaments
        for i in 1:filament.lf
            pos = filament.coors[:, i]
            delete!(lattice.occupancy, pos)
            pos[2] -= origin
            wrap_pos!(lattice, pos)
            lattice.occupancy[copy(pos)] = (filament.index, i)
            filament.coors[:, i] = pos
        end
    end

    return nothing
end

"""
Translate a filament with given vector.

Returns false if steric clash occurs.
"""
function translate_filament!(filament::Filament, lattice::Lattice, move_vector::Vector{Int})
    for i in 1:filament.lf
        pos = filament.coors[:, i]
        delete!(lattice.occupancy, pos)
    end

    for i in 1:filament.lf
        pos = filament.coors[:, i]
        pos += move_vector
        wrap_pos!(lattice, pos)
        if pos in keys(lattice.occupancy)
            return false
        end
        filament.coors[:, i] = pos
        lattice.occupancy[copy(pos)] = (filament.index, i)
    end

    return true
end

"""Attempt to translate a filament up or down by one site."""
function attempt_translation_move!(system::System, lattice::Lattice, ::Biases)
    use_trial_coors!(system, lattice)

    # Do not move the first filament as it is the reference
    filament = rand(system.filaments[2:end])
    move_vector = [0, rand([-1, 1])]
    if !translate_filament!(filament, lattice, move_vector)
        accept_current!(system, lattice)
        use_current_coors!(system, lattice)

        return false
    end

    if !ring_and_system_connected(system, lattice)
        accept_current!(filament, lattice)
        use_current_coors!(system, lattice)

        return false
    end

    delta_energy = energy_diff(system, lattice, filament)
    accept = accept_move(system, delta_energy)
    if accept
        accept_trial!(filament, lattice)
    else
        accept_current!(filament, lattice)
    end

    use_current_coors!(system, lattice)

    return accept
end

"""Find filament sites that are below upper periodic boundary."""
function find_split_points(system::System, lattice::Lattice)
    split_points::Vector{Int} = []
    for filament in system.filaments
        push!(split_points, 0)
        for i in 1:filament.lf
            y = filament.coors[2, i]
            if y == lattice.height
                if i != filament.lf
                    split_points[end] = i
                    break
                end
            end
        end
    end

    return split_points
end

"""
Translate segments of filaments below split point only.

Returns false if steric clash occurs.
"""
function translate_filaments_with_split_points!(
    system::System,
    lattice::Lattice,
    split_points::Vector{Int},
    dir::Int
)
    for (filament, split_point) in zip(system.filaments, split_points)
        for i in 1:split_point
            pos = filament.coors[:, i]
            delete!(lattice.occupancy, pos)
        end
    end

    for (filament, split_point) in zip(system.filaments, split_points)
        for i in 1:split_point
            filament.coors[2, i] += dir
            pos = filament.coors[:, i]
            if !(pos in keys(lattice.occupancy))
                lattice.occupancy[pos] = (filament.index, i)
            else
                return false
            end
        end
    end

    return true
end

"""Attempt to increase or decrease radius by one lattice site."""
function attempt_radius_move!(system::System, lattice::Lattice, biases::Biases)
    dir = rand([-1, 1])
    if (dir == 1 && lattice.height == lattice.max_height ||
        dir == -1 && lattice.height == lattice.min_height)
        accept_current!(system, lattice)
        use_current_coors!(system, lattice)

        return false
    end

    use_trial_coors!(system, lattice)
    split_points = find_split_points(system, lattice)
    if !translate_filaments_with_split_points!(system, lattice, split_points, dir)
        accept_current!(system, lattice)
        use_current_coors!(system, lattice)

        return false
    end

    lattice.trial_height += dir
    update_radius!(system, lattice, lattice.trial_height)
    # if lattice.trial_height > lattice.max_height
    #    println(dir)
    #    println(lattice.height)
    #    println(lattice.current_height)
    #    println(lattice.trial_height)
    #    throw(DomainError)
    # end
    if !filaments_contiguous(system, lattice)
        accept_current!(system, lattice)
        use_current_coors!(system, lattice)

        return false
    end

    if dir == 1
        if !ring_and_system_connected(system, lattice)
            accept_current!(system, lattice)
            use_current_coors!(system, lattice)

            return false
        end
    end

    delta_energy = energy_diff(system, lattice, biases)
    # mult = (lattice.trial_height + 1) / (lattice.current_height + 1)
    # accept = accept_move(system, delta_energy, mult)
    accept = accept_move(system, delta_energy)
    if accept
        accept_trial!(system, lattice)
    else
        update_radius!(system, lattice, lattice.height - dir)
        accept_current!(system, lattice)
    end

    use_current_coors!(system, lattice)

    return accept
end

"""Return an attempt move function according to move type weights."""
function select_move(simparms::SimulationParams)
    if simparms.radius_move_freq > rand(Float64)
        return attempt_radius_move!
    else
        return attempt_translation_move!
    end
end

"""Open and write header for order parameters file."""
function prepare_ops_file(filename::String)
    file = open(filename, "w")
    header = "step energy height radius"
    println(file, header)

    return file
end

"""Open and write header for umbrella sampling output file."""
function prepare_us_file(filename::String, lattice::Lattice)
    file = open(filename, "w")
    header = ""
    for h in lattice.min_height:lattice.max_height
        header *= "$h "
    end

    println(file, header)

    return file
end

"""Open and write header for VTF output file."""
function prepare_vtf_file(filename::String, system::System)
    file = open(filename, "w")
    atom_i = 0
    for filament in system.filaments
        end_atom_i = atom_i + filament.lf - 1
        println(file, "a $atom_i:$end_atom_i c $(filament.index) r 2.5")
        atom_i += filament.lf
    end

    println(file, "")

    return file
end

"""Write configuration in VTF format to file."""
function write_vtf(system::System, file::IOStream)
    println(file, "t")
    for filament in system.filaments
        for i in 1:filament.lf
            # Widen the aspect ratio for easier viewing
            x = filament.coors[1, i] * 10
            y = filament.coors[2, i]
            println(file, "$x $y 0")
        end
    end

    println(file, "")

    return nothing
end

"""Write current order parameters to file."""
function write_ops(system::System, lattice::Lattice, step::Int, file::IOStream)
    println(file, "$step $(system.energy) $(lattice.height) $(system.radius)")

    return nothing
end

"""Write current umbrella sampling output data to file."""
function write_us_data(data::Vector, file::IOStream)
    for d in data
        print(file, "$d ")
    end

    println(file, "")

    return nothing
end

"""Write simulation parameters to file."""
function write_params(system::System, simparms::SimulationParams, file::IOStream)
    parms = Dict(
        "lf" => system.parms.lf,
        "T" => system.parms.T,
        "kd" => system.parms.kd,
        "ks" => system.parms.ks,
        "EI" => system.parms.EI,
        "Lf" => system.parms.Lf,
        "Xc" => system.parms.Xc,
        "Nfil" => system.parms.Nfil,
        "Nsca" => system.parms.Nsca,
        "delta" => system.parms.delta,
        "steps" => simparms.steps,
        "write_interval" => simparms.write_interval,
        "filebase" => simparms.filebase,
        "max_bias_diff" => simparms.max_bias_diff,
        "radius_move_freq" => simparms.radius_move_freq,
        "iters" => simparms.iters,
        "analytical_biases" => simparms.analytical_biases,
        "binwidth" => simparms.binwidth
    )
    println(file, json(parms))

    return nothing
end

"""Add one to count of visits of current state."""
function update_counts(biases::Biases, lattice::Lattice)
    bin = get_bin(biases.barriers, lattice.height)
    biases.counts[bin] += 1

    return nothing
end

"""
$(TYPEDSIGNATURES)

Run an MC simulation.
"""
function run!(
    system::System,
    lattice::Lattice,
    simparms::SimulationParams,
    biases::Biases,
    ops_file::IOStream,
    vtf_file::IOStream
)

    # Overly simple way to record move acceptance frequencies
    attempts = [0, 0]
    accepts = [0, 0]
    recenter!(system, lattice)
    for step in 1:simparms.steps
        attempt_move! = select_move(simparms)
        accepted = attempt_move!(system, lattice, biases)
        attempt_move! === attempt_translation_move! ? attempts[1] += 1 : attempts[2] += 1
        if accepted
            attempt_move! === attempt_translation_move! ? accepts[1] += 1 : accepts[2] += 1
        end
        update_counts(biases, lattice)
        if step % simparms.write_interval == 0
            ring_and_system_connected(system, lattice, true)
            println("Step: $step")
            system.energy = total_energy(system, lattice)
            write_ops(system, lattice, step, ops_file)
            write_vtf(system, vtf_file)
        end
    end

    println()
    println("Move info:")
    println()
    println("Filament translation")
    println("Attempts: $(attempts[1])")
    println("Accepts: $(accepts[1])")
    println("Ratio: $(accepts[1] / attempts[1])")
    println()
    println("Radius move")
    println("Attempts: $(attempts[2])")
    println("Accepts: $(accepts[2])")
    println("Ratio: $(accepts[2] / attempts[2])")

    return nothing
end

"""
$(TYPEDSIGNATURES)

Run an MC simulation.
"""
function run!(system::System, lattice::Lattice, simparms::SimulationParams)
    ops_file = prepare_ops_file("$(simparms.filebase).ops")
    vtf_file = prepare_vtf_file("$(simparms.filebase).vtf", system)
    biases = Biases(lattice, simparms.binwidth)
    run!(system, lattice, simparms, biases, ops_file, vtf_file)
    close(ops_file)
    close(vtf_file)
    open("$(simparms.filebase).parms", "w") do file
        write_params(system, simparms, file)
    end

    return nothing
end

"""
Calculate new biases from counts.

Updates the passed arrays in the process.
"""
function update_biases!(biases::Biases, T::Float64, max_bias_diff::Float64)
    reduced_enes = biases.enes ./ (kb * T)
    steps = sum(biases.counts)
    norm = sum(biases.counts .* exp.(reduced_enes))
    max_bias_diff *= kb * T
    for i in 1:length(biases.counts)
        if biases.counts[i] == 0
            biases.freqs[i] = 0
            biases.probs[i] = 0
            bias_diff = -max_bias_diff
        else
            biases.freqs[i] = biases.counts[i] / steps
            biases.probs[i] = biases.counts[i] * exp(reduced_enes[i]) / norm
            bias_diff = kb * T * log(biases.probs[i]) - biases.enes[i]
            if bias_diff > max_bias_diff
                bias_diff = max_bias_diff
            elseif bias_diff < -max_bias_diff
                bias_diff = -max_bias_diff
            end
        end

        biases.enes[i] += bias_diff
        biases.counts[i] = 0
    end

    return nothing
end

"""
Calculate biases from analytical model.

Interpolate between the bin endpoints.
"""
function analytical_biases(system::System, lattice::Lattice, biases::Biases)
    radius_max = calc_radius(system.parms.delta, lattice.max_height)
    for bin in 1:biases.numbins
        if bin == biases.numbins
            lower_barrier = biases.barriers[bin - 1]
            upper_barrier = lower_barrier + biases.binwidth
        else
            upper_barrier = biases.barriers[bin]
            lower_barrier = upper_barrier - biases.binwidth
        end

        lower_upper_biases = []
        for h in [lower_barrier, upper_barrier]
            radius = calc_radius(system.parms.delta, h)
            L = 2 * pi * (radius_max - radius) / system.parms.Nsca
            overlaps = system.parms.Nsca + 2 * (system.parms.Nfil - system.parms.Nsca)
            sliding_energy = overlaps * overlap_energy(system, L)
            bending_energy = system.parms.Nfil * filament_bending_energy(system, radius)
            push!(lower_upper_biases, -(sliding_energy + bending_energy))
        end
        biases.enes[bin] = sum(lower_upper_biases) / 2
    end

    return nothing
end

"""
$(TYPEDSIGNATURES)

Run an umbrella sampling MC simulation.
"""
function run_us!(system::System, lattice::Lattice, simparms::SimulationParams)
    open("$(simparms.filebase).parms", "w") do file
        write_params(system, simparms, file)
    end

    biases = Biases(lattice, simparms.binwidth)

    start_iter = 1
    end_iter = simparms.iters
    if simparms.read_biases
        read_biases!(simparms.biases_filename, simparms.restart_iter, biases)
        start_iter = simparms.restart_iter + 1
        end_iter += simparms.restart_iter
    elseif simparms.analytical_biases
        analytical_biases(system, lattice, biases)
    end
    counts_file = prepare_us_file("$(simparms.filebase).counts", lattice)
    freqs_file = prepare_us_file("$(simparms.filebase).freqs", lattice)
    biases_file = prepare_us_file("$(simparms.filebase).biases", lattice)
    for i in start_iter:end_iter
        println()
        println("Iter: $i")
        iter_filebase = "$(simparms.filebase)_iter-$i"
        ops_file = prepare_ops_file("$iter_filebase.ops")
        vtf_file = prepare_vtf_file("$iter_filebase.vtf", system)
        run!(system, lattice, simparms, biases, ops_file, vtf_file)
        write_us_data(biases.counts, counts_file)
        update_biases!(biases, system.parms.T, simparms.max_bias_diff)
        write_us_data(biases.freqs, freqs_file)
        write_us_data(biases.enes, biases_file)
        close(ops_file)
        close(vtf_file)
    end

    close(counts_file)
    close(freqs_file)
    close(biases_file)

    println()

    return nothing
end

end

max_stars_repo_name: cumberworth/ActinRingsMC

max_issues_repo_name: cumberworth/ActinRingsMC

max_forks_repo_name: cumberworth/ActinRingsMC


text:
# run: ~/.julia/bin/mpiexecjl -n 4 julia --project scripts/heat_1D_mpi.jl
using Plots, Printf, MAT
import MPI

do_save = true

# MPI functions
@views function update_halo(A, neighbors_x, comm)
    if neighbors_x[1] >= 0 # MPI_PROC_NULL?
        sendbuf = A[2]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[1], 0, comm)
        MPI.Recv!(recvbuf, neighbors_x[1], 1, comm)
        A[1] = recvbuf[1]
    end
    if neighbors_x[2] >= 0 # MPI_PROC_NULL?
        sendbuf = A[end-1]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[2], 1, comm)
        MPI.Recv!(recvbuf, neighbors_x[2], 0, comm)
        A[end] = recvbuf[1]
    end
    return
end

@views function heat_1D_mpi()
    # MPI
    MPI.Init()
    dims        = [0]
    comm        = MPI.COMM_WORLD
    nprocs      = MPI.Comm_size(comm)
    MPI.Dims_create!(nprocs, dims)
    comm_cart   = MPI.Cart_create(comm, dims, [0], 1)
    me          = MPI.Comm_rank(comm_cart)
    coords      = MPI.Cart_coords(comm_cart)
    neighbors_x = MPI.Cart_shift(comm_cart, 0, 1)
    if (me==0) println("nprocs=$(nprocs), dims[1]=$(dims[1])") end
    # Physics
    lx   = 10.0
    λ    = 1.0
    ρCp  = 1.0
    nt   = 100
    # Numerics
    nx   = 32   # local
    nx_g = dims[1]*(nx-2) + 2 # global
    # Derived numerics
    dx   = lx/nx_g            # global
    dt   = dx^2/ρCp/λ/2.1
    # Array allocation
    qx   = zeros(nx-1)
    # Initial condition
    x0   = coords[1]*(nx-2)*dx
    xc   = [x0 + ix*dx - dx/2 - 0.5*lx  for ix=1:nx]
    T    = exp.(.-xc.^2)
    # Time loop
    for it = 1:nt
        if (it==11) global t0 = Base.time() end
        qx         .= .-λ.*diff(T)./dx
        T[2:end-1] .= T[2:end-1] .- dt./ρCp.*diff(qx)./dx
        update_halo(T, neighbors_x, comm_cart)
     end
    time_s = (Base.time()-t0)
    if (me==0) @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", time_s, round((2/1e9*nx*sizeof(lx))/(time_s/(nt-10)), sigdigits=2)) end
    if do_save file = matopen("$(@__DIR__)/T_$(me).mat", "w"); write(file, "T", Array(T)); close(file) end
    MPI.Finalize()
    return
end

heat_1D_mpi()

ext: jl

max_stars_repo_name: luraess/geo-hpc-course

max_issues_repo_name: luraess/geo-hpc-course

max_forks_repo_name: luraess/geo-hpc-course


text:
VERSION >= v"0.4.0-dev+6521" && __precompile__()

module OnlineMoments

import StatsBase

#-----------------------------------------------------------------------------

immutable Mean
  N::Float64
  M::Float64
end

Base.zero(::Type{Mean}) = Mean(0.0, 0.0)

Base.mean(a::Mean) = a.M

StatsBase.nobs(a::Mean) = convert(Int64, a.N)

function Base.(:+)(a::Mean, x::Real)
  N = a.N + 1.0
  Δ = x - a.M
  δ = Δ / N
  M = a.M + δ
  Mean(N, M)
end

function Base.(:+)(a::Mean, b::Mean)
  N = a.N + b.N
  Δ = b.M - a.M
  δ = Δ / N
  M = a.M + b.N*δ
  Mean(N, M)
end

#-----------------------------------------------------------------------------

immutable Variance
  N::Float64
  M::Float64
  V::Float64
end

Base.zero(::Type{Variance}) = Variance(0.0, 0.0, 0.0)

Base.mean(a::Variance) = a.M

Base.var(a::Variance) = a.V / (a.N - 1.0)

Base.std(a::Variance) = sqrt(var(a))

StatsBase.sem(a::Variance) = sqrt(a.V / ((a.N - 1.0) * a.N))

StatsBase.nobs(a::Variance) = convert(Int64, a.N)

function Base.(:+)(a::Variance, x::Real)
  N = a.N + 1.0
  Δ = x - a.M
  δ = Δ / N
  ξ = Δ * a.N
  M = a.M + δ
  V = a.V + δ*ξ
  Variance(N, M, V)
end

function Base.(:+)(a::Variance, b::Variance)
  N = a.N + b.N
  Δ = b.M - a.M
  δ = Δ / N
  ξ = Δ * (a.N*b.N)
  M = a.M + b.N*δ
  V = a.V + b.V + δ*ξ
  Variance(N, M, V)
end

#-----------------------------------------------------------------------------

immutable Moments
  N::Float64
  M::Float64
  V::Float64
  S::Float64
  K::Float64
end

Base.zero(::Type{Moments}) = Moments(0.0, 0.0, 0.0, 0.0, 0.0)

Base.mean(a::Moments) = a.M

Base.var(a::Moments) = a.V / (a.N - 1.0)

Base.std(a::Moments) = sqrt(var(a))

StatsBase.sem(a::Moments) = sqrt(a.V / ((a.N - 1.0) * a.N))

StatsBase.skewness(a::Moments) = sqrt(a.N) * a.S / a.V^1.5

StatsBase.kurtosis(a::Moments) = a.N * a.K / (a.V * a.V) - 3.0

StatsBase.nobs(a::Moments) = convert(Int64, a.N)

function Base.(:+)(a::Moments, x::Real)
  N = a.N + 1.0
  Δ = x - a.M
  δ = Δ / N
  δ² = δ * δ
  δ³ = δ² * δ
  ξ = Δ * a.N
  M = a.M + δ
  V = a.V + δ*ξ
  S = a.S + δ²*ξ*(a.N - 1.0) - 3.0*δ*a.V
  K = a.K + δ³*ξ*(a.N*a.N - a.N + 1.0) + 6.0*δ²*a.V - 4.0*δ*a.S
  Moments(N, M, V, S, K)
end

function Base.(:+)(a::Moments, b::Moments)
  N = a.N + b.N
  Δ = b.M - a.M
  δ = Δ / N
  δ² = δ * δ
  δ³ = δ² * δ
  ξ = Δ * (a.N*b.N)
  M = a.M + b.N*δ
  V = a.V + b.V + δ*ξ
  S = a.S + b.S + δ²*ξ*(a.N - b.N) + 3.0*δ*(a.N*b.V - b.N*a.V)
  K = a.K + b.K + δ³*ξ*(a.N*a.N - a.N*b.N + b.N*b.N) + 6.0*δ²*(a.N*a.N*b.V + b.N*b.N*a.V) + 4.0*δ*(a.N*b.S - b.N*a.S)
  Moments(N, M, V, S, K)
end

#-----------------------------------------------------------------------------

end


text:
function create_proj_dirs(Proj::Module)
    _check_is_proj(Proj)
    for dirfun in (
            datdir, plotsdir, scriptsdir,
            papersdir, procdir, rawdir, cachedir, devdir
        )
        dir = dirfun(Proj)
        !isdir(dir) && mkpath(dir)
    end
end

# ---------------------------------------------------------------------
macro create_proj_dirs()
    quote $(create_proj_dirs)(@__MODULE__) end
end

# TODO: Transform to your layout
# function print_proj_layout()
#     """
#     │projectdir          <- Project's main folder. It is initialized as a Git
#     │                       repository with a reasonable .gitignore file.
#     │
#     ├── _research        <- WIP scripts, code, notes, comments,
#     │   |                   to-dos and anything in an alpha state.
#     │   └── tmp          <- Temporary data folder.
#     │
#     ├── data             <- **Immutable and add-only!**
#     │   ├── sims         <- Data resulting directly from simulations.
#     │   ├── exp_pro      <- Data from processing experiments.
#     │   └── exp_raw      <- Raw experimental data.
#     │
#     ├── plots            <- Self-explanatory.
#     ├── notebooks        <- Jupyter, Weave or any other mixed media notebooks.
#     │
#     ├── papers           <- Scientific papers resulting from the project.
#     │
#     ├── scripts          <- Various scripts, e.g. simulations, plotting, analysis,
#     │   │                   The scripts use the `src` folder for their base code.
#     │   └── intro.jl     <- Simple file that uses DrWatson and uses its greeting.
#     │
#     ├── src              <- Source code for use in this project. Contains functions,
#     │                       structures and modules that are used throughout
#     │                       the project and in multiple scripts.
#     │
#     ├── README.md        <- Optional top-level README for anyone using this project.
#     ├── .gitignore       <- by default ignores _research, data, plots, videos,
#     │                       notebooks and latex-compilation related files.
#     │
#     ├── Manifest.toml    <- Contains full list of exact package versions used currently.
#     └── Project.toml     <- Main project file, allows activation and installation.
#                             Includes DrWatson by default.
#     """ |> println
# end


text:
import Dates
using AutoHashEquals

const WEEKDAYS = Dict(map(reverse,enumerate(map(Symbol ∘ uppercase, Dates.ENGLISH.days_of_week_abbr))))
const NTH_PERIODS = ["1st", "2nd", "3rd", "4th", "5th"]
const NTH_LAST_PERIODS = ["Last", "2nd Last", "3rd Last", "4th Last", "5th Last"]
const PERIODS = merge(Dict(map(reverse,enumerate(NTH_PERIODS))), Dict(map(reverse,enumerate(NTH_LAST_PERIODS))))
const MONTHS = Dict(zip(map(Symbol ∘ uppercase, Dates.ENGLISH.months_abbr),range(1,stop=12)))

"""
    FDOM()
    FDOM(calendars)

Calculate the first day of the month. Optionally can also take calendar names to determine
the first business day of the month.

### Examples
```julia-repl
julia> RDates.FDOM() + Date(2019,1,13)
2019-01-01
julia> rd"FDOM" + Date(2019,1,13)
2019-01-01

julia> cals = SimpleCalendarManager()
julia> setcalendar!(cals, "WEEKEND", WeekendCalendar())
julia> apply(RDates.FDOM("WEEKEND"), Date(2017,1,13), cals)
2017-01-02
julia> apply(rd"FDOM@WEEKEND", Date(2017,1,13), cals)
2017-01-02
```
"""
struct FDOM <: RDate
    FDOM() = new()
    FDOM(calendars::Nothing) = new()
    FDOM(calendars) = CalendarAdj(calendars, new(), HolidayRoundingNBD())
end

apply(::FDOM, d::Dates.Date, ::CalendarManager) = Dates.firstdayofmonth(d)
multiply(x::FDOM, ::Integer) = x
Base.:-(x::FDOM) = x

Base.show(io::IO, ::FDOM) = print(io, "FDOM")
register_grammar!(E"FDOM" > FDOM)
register_grammar!(E"FDOM@" + PCalendarNames() > calendar_name -> FDOM(map(String, split(calendar_name, "|"))))

"""
    LDOM()
    LDOM(calendars)

Calculate the last day of the month. Optionally can also take calendar names to determine
the last business day of the month.

### Examples
```julia-repl
julia> RDates.LDOM() + Date(2019,1,13)
2019-01-31
julia> rd"LDOM" + Date(2019,1,13)
2019-01-31
julia> cals = SimpleCalendarManager()
julia> setcalendar!(cals, "WEEKEND", WeekendCalendar())
julia> apply(RDates.LDOM("WEEKEND"), Date(2021,1,13), cals)
2021-01-29
julia> apply(rd"LDOM@WEEKEND", Date(2021,1,13), cals)
2021-01-29
```
"""
struct LDOM <: RDate
    LDOM() = new()
    LDOM(calendars::Nothing) = new()
    LDOM(calendars) = CalendarAdj(calendars, new(), HolidayRoundingPBD())
end

apply(::LDOM, d::Dates.Date, ::CalendarManager) = Dates.lastdayofmonth(d)
multiply(x::LDOM, ::Integer) = x
Base.:-(x::LDOM) = x

Base.show(io::IO, ::LDOM) = print(io, "LDOM")
register_grammar!(E"LDOM" > LDOM)
register_grammar!(E"LDOM@" + PCalendarNames() > calendar_name -> LDOM(map(String, split(calendar_name, "|"))))

"""
    Easter(yearδ::Int64)

A date that is well known from hunting eggs and pictures of bunnies, it's a rather
tricky calculation to perform. We provide a simple method to allow you to get the
Easter for the given year (plus some delta).

!!! note
    `0E` will get the Easter of the current year, so it could be before or
    after the date you've provided.

### Examples
```julia-repl
julia> RDates.Easter(0) + Date(2019,1,1)
2019-04-21
julia> rd"0E" + Date(2019,1,1)
2019-04-21
julia> RDates.Easter(0) + Date(2019,8,1)
2019-04-21
julia> RDates.Easter(10) + Date(2019,8,1)
2029-04-01
```
"""
struct Easter <: RDate
    yearδ::Int64
end

function apply(rdate::Easter, date::Dates.Date, cal_mgr::CalendarManager)
    y = Dates.year(date) + rdate.yearδ
    a = rem(y, 19)
    b = div(y, 100)
    c = rem(y, 100)
    d = div(b, 4)
    e = rem(b, 4)
    f = div(b + 8, 25)
    g = div(b - f + 1, 3)
    h = rem(19*a + b - d - g + 15, 30)
    i = div(c, 4)
    k = rem(c, 4)
    l = rem(32 + 2*e + 2*i - h - k, 7)
    m = div(a + 11*h + 22*l, 451)
    n = div(h + l - 7*m + 114, 31)
    p = rem(h + l - 7*m + 114, 31)
    return Dates.Date(y, n, p + 1)
end

multiply(x::Easter, count::Integer) = Easter(x.yearδ*count)
Base.:-(rdate::Easter) = Easter(-rdate.yearδ)
Base.show(io::IO, rdate::Easter) = print(io, "$(rdate.yearδ)E")
register_grammar!(PInt64() + E"E" > Easter)

"""
    Day(days::Int64)

Provides us with the ability to add or subtract days from a date. This is
equivalent to the `Dates.Day` struct.

### Examples
```julia-repl
julia> RDates.Day(3) + Date(2019,1,1)
2019-01-04
julia> rd"3d" + Date(2019,1,1)
2019-01-04
julia> RDates.Day(-2) + Date(2019,1,1)
2018-12-30
```
"""
struct Day <: RDate
    days::Int64
end

apply(x::Day, y::Dates.Date, cal_mgr::CalendarManager) = y + Dates.Day(x.days)
multiply(x::Day, count::Integer) = Day(x.days*count)
Base.:-(x::Day) = Day(-x.days)
Base.:+(x::Day, y::Day) = Day(x.days + y.days)

Base.show(io::IO, rdate::Day) = print(io, "$(rdate.days)d")
register_grammar!(PInt64() + E"d" > Day)

"""
    Week(weeks::Int64)

Provides us with the ability to add or subtract weeks from a date. This is
equivalent to the `Dates.Week` struct.

### Examples
```julia-repl
julia> RDates.Week(3) + Date(2019,1,1)
2019-01-22
julia> rd"3w" + Date(2019,1,1)
2019-01-22
julia> RDates.Week(-2) + Date(2019,1,1)
2018-12-18
```
"""
struct Week <: RDate
    weeks::Int64
end

apply(x::Week, y::Dates.Date, cal_mgr::CalendarManager) = y + Dates.Week(x.weeks)
multiply(x::Week, count::Integer) = Week(x.weeks*count)
Base.:-(x::Week) = Week(-x.weeks)
Base.:+(x::Week, y::Week) = Week(x.weeks + y.weeks)
Base.:+(x::Day, y::Week) = Day(x.days + 7*y.weeks)
Base.:+(x::Week, y::Day) = Day(7*x.weeks + y.days)

Base.show(io::IO, rdate::Week) = print(io, "$(rdate.weeks)w")
register_grammar!(PInt64() + E"w" > Week)

"""
    Month(months::Int64)
    Month(months::Int64, idc::InvalidDayConvention, mic::MonthIncrementConvention)

Provides us with the ability to move a specified number of months, with conventions
to handle how we should increment and what to do if we fall on an invalid day.

### Examples
```julia-repl
julia> RDates.Month(1) + Date(2019,1,31)
2019-02-28
julia> rd"1m" + Date(2019,1,31)
2019-02-28
julia> RDates.Month(1, RDates.InvalidDayFDONM(), RDates.MonthIncrementPDOM()) + Date(2019,1,31)
2019-03-01
julia> rd"1m[FDONM;PDOM]" + Date(2019,1,31)
2019-03-01
julia> RDates.Month(1, RDates.InvalidDayNDONM(), RDates.MonthIncrementPDOM()) + Date(2019,1,31)
2019-03-03
julia> rd"1m[NDONM;PDOM]" + Date(2019,1,31)
2019-03-03
julia> RDates.Month(1, RDates.InvalidDayNDONM(), RDates.MonthIncrementPDOMEOM()) + Date(2019,1,31)
2019-02-28
julia> rd"1m[NDONM;PDOMEOM]" + Date(2019,1,31)
2019-02-28
julia> RDates.Month(-1, RDates.InvalidDayNDONM(), RDates.MonthIncrementPDOMEOM()) + Date(2019,2,28)
2019-01-31
julia> rd"-1m[NDONM;PDOMEOM]" + Date(2019,2,28)
2019-01-31
```
"""
struct Month <: RDate
    months::Int64
    idc::InvalidDayConvention
    mic::MonthIncrementConvention

    Month(months::Int64) = new(months, InvalidDayLDOM(), MonthIncrementPDOM())
    Month(months::Int64, idc::InvalidDayConvention, mic::MonthIncrementConvention) = new(months, idc, mic)
end

function apply(rdate::Month, date::Dates.Date, cal_mgr::CalendarManager)
    y, m = Dates.yearmonth(date)
    ny = Dates.yearwrap(y, m, rdate.months)
    nm = Dates.monthwrap(m, rdate.months)
    ay, am, ad = adjust(rdate.mic, date, nm, ny, cal_mgr)
    ld = Dates.daysinmonth(ay, am)
    return ad <= ld ? Dates.Date(ay, am, ad) : adjust(rdate.idc, ad, am, ay)
end
multiply(x::Month, count::Integer) = Month(x.months*count, x.idc, x.mic)
Base.:-(x::Month) = Month(-x.months)

Base.show(io::IO, rdate::Month) = (print(io, "$(rdate.months)m["), show(io, rdate.idc), print(io, ";"), show(io, rdate.mic), print(io,"]"))
register_grammar!(PInt64() + E"m" > Month)
register_grammar!(PInt64() + E"m[" + Alt(map(Pattern, collect(keys(INVALID_DAY_MAPPINGS)))...) + E";" + Alt(map(Pattern, collect(keys(MONTH_INCREMENT_MAPPINGS)))...) + E"]" > (d,idc,mic) -> Month(d, INVALID_DAY_MAPPINGS[idc], MONTH_INCREMENT_MAPPINGS[mic]))
# We also have the more complex PDOMEOM month increment which we handle separately.
register_grammar!(PInt64() + E"m[" + Alt(map(Pattern, collect(keys(INVALID_DAY_MAPPINGS)))...) + E";PDOMEOM@" + PCalendarNames() + E"]" > (d,idc,calendar_name) -> Month(d, INVALID_DAY_MAPPINGS[idc], MonthIncrementPDOMEOM(map(String, split(calendar_name, "|")))))

"""
    Year(years::Int64)
    Year(years::Int64, idc::InvalidDayConvention, mic::MonthIncrementConvention)

Provides us with the ability to move a specified number of months, with conventions
to handle how we should increment and what to do if we fall on an invalid day.

!!! note
    While these conventions are necessary, it's only around the handling of leap years
    and when we're on the last day of the February that it actually matters.

### Examples
```julia-repl
julia> RDates.Year(1) + Date(2019,2,28)
2020-02-28
julia> rd"1y" + Date(2019,2,28)
2020-02-28
julia> RDates.Year(1, RDates.InvalidDayFDONM(), RDates.MonthIncrementPDOMEOM()) + Date(2019,2,28)
2020-02-29
julia> rd"1y[FDONM;PDOMEOM]" + Date(2019,2,28)
2020-02-29
julia> RDates.Year(1, RDates.InvalidDayLDOM(), RDates.MonthIncrementPDOM()) + Date(2020,2,29)
2021-02-28
julia> rd"1y[LDOM;PDOM]" + Date(2020,2,29)
2021-02-28
julia> RDates.Year(1, RDates.InvalidDayFDONM(), RDates.MonthIncrementPDOM()) + Date(2020,2,29)
2021-03-01
julia> rd"1y[FDONM;PDOM]" + Date(2020,2,29)
2021-03-01
```
"""
struct Year <: RDate
    years::Int64
    idc::InvalidDayConvention
    mic::MonthIncrementConvention

    Year(years::Int64) = new(years, InvalidDayLDOM(), MonthIncrementPDOM())
    Year(years::Int64, idc::InvalidDayConvention, mic::MonthIncrementConvention) = new(years, idc, mic)
end

function apply(rdate::Year, date::Dates.Date, cal_mgr::CalendarManager)
    oy, m = Dates.yearmonth(date)
    ny = oy + rdate.years
    (ay, am, ad) = adjust(rdate.mic, date, m, ny, cal_mgr)
    ld = Dates.daysinmonth(ay, am)
    return ad <= ld ? Dates.Date(ay, am, ad) : adjust(rdate.idc, ad, am, ay)
end
multiply(x::Year, count::Integer) = Year(x.years*count, x.idc, x.mic)
Base.:-(x::Year) = Year(-x.years)

Base.show(io::IO, rdate::Year) = (print(io, "$(rdate.years)y["), show(io, rdate.idc), print(io, ";"), show(io, rdate.mic), print(io,"]"))
register_grammar!(PInt64() + E"y" > Year)
register_grammar!(PInt64() + E"y[" + Alt(map(Pattern, collect(keys(INVALID_DAY_MAPPINGS)))...) + E";" + Alt(map(Pattern, collect(keys(MONTH_INCREMENT_MAPPINGS)))...) + E"]" > (d,idc,mic) -> Year(d, INVALID_DAY_MAPPINGS[idc], MONTH_INCREMENT_MAPPINGS[mic]))
# We also have the more complex PDOMEOM month increment which we handle separately.
register_grammar!(PInt64() + E"y[" + Alt(map(Pattern, collect(keys(INVALID_DAY_MAPPINGS)))...) + E";PDOMEOM@" + PCalendarNames() + E"]" > (d,idc,calendar_name) -> Year(d, INVALID_DAY_MAPPINGS[idc], MonthIncrementPDOMEOM(map(String, split(calendar_name, "|")))))

"""
    DayMonth(day::Int64, month::Int64)
    DayMonth(day::Int64, month::Symbol)

Provides us with the ability to move to a specific day and month in the provided
year.
!!! note
    `1MAR` will get the 1st of March of the current year, so it could be before
    or after the date you've provided.

### Examples
```julia-repl
julia> RDates.DayMonth(23, 10) + Date(2019,1,1)
2019-10-23
julia> RDates.DayMonth(23, :OCT) + Date(2019,1,1)
2019-10-23
julia> rd"23OCT" + Date(2019,1,1)
2019-10-23
```
"""
struct DayMonth <: RDate
    day::Int64
    month::Int64

    DayMonth(day::Int64, month::Int64) = new(day, month)
    DayMonth(day::Int64, month::Symbol) = new(day, RDates.MONTHS[month])
end

apply(rdate::DayMonth, date::Dates.Date, cal_mgr::CalendarManager) = Dates.Date(Dates.year(date), rdate.month, rdate.day)
multiply(x::DayMonth, ::Integer) = x
Base.:-(x::DayMonth) = x

Base.show(io::IO, rdate::DayMonth) = print(io, "$(rdate.day)$(uppercase(Dates.ENGLISH.months_abbr[rdate.month]))")
register_grammar!(PPosInt64() + month_short > (d,m) -> DayMonth(d,MONTHS[Symbol(m)]))

"""
    Date(date::Dates.Date)
    Date(year::Int64, month::Int64, day::Int64)

Provides us with the ability to move to a specific date, irrespective of the date
passed in. This is primarily used when you want to provide a pivot point for ranges
which doesn't relate to the start or end.

### Examples
```julia-repl
julia> RDates.Date(Dates.Date(2017,10,23)) + Date(2019,1,1)
2017-10-23
julia> RDates.Date(2017,10,23) + Date(2019,1,1)
2017-10-23
julia> rd"23OCT2017" + Date(2019,1,1)
2017-10-23
```
"""
struct Date <: RDate
    date::Dates.Date

    Date(date::Dates.Date) = new(date)
    Date(y::Int64, m::Int64, d::Int64) = new(Dates.Date(y, m, d))
end

multiply(x::Date, ::Integer) = x
Base.:-(x::Date) = x

apply(rdate::Date, date::Dates.Date, cal_mgr::CalendarManager) = rdate.date
Base.show(io::IO, rdate::Date) = print(io, "$(Dates.day(rdate.date))$(uppercase(Dates.ENGLISH.months_abbr[Dates.month(rdate.date)]))$(Dates.year(rdate.date))")
register_grammar!(PPosInt64() + month_short + PPosInt64() > (d,m,y) -> Date(Dates.Date(y, MONTHS[Symbol(m)], d)))

"""
    NthWeekdays(dayofweek::Int64, period::Int64)
    NthWeekdays(dayofweek::Symbol, period::Int64)

Move to the nth weekday in the given month and year. This is commonly used for holiday
calendars, such as [Thanksgiving](https://en.wikipedia.org/wiki/Thanksgiving) which in
the U.S. falls on the 4th Thursday in November.

!!! note
    It's possible that a given period (such as the 5th weekday) may exist for only
    a subsection of dates. While it's a valid RDate it may not produce valid results
    when applied (and will throw an exception)

### Examples
```julia-repl
julia> RDates.NthWeekdays(:MON, 2) + Date(2019,1,1)
2019-01-14
julia> RDates.NthWeekdays(1, 2) + Date(2019,1,1)
2019-01-14
julia> rd"2nd MON" + Date(2019,1,1)
2019-01-14
julia> RDates.NthWeekdays(:MON, 5) + Date(2019,1,1)
ERROR: ArgumentError: Day: 35 out of range (1:31)
```
"""
struct NthWeekdays <: RDate
    dayofweek::Int64
    period::Int64

    NthWeekdays(dayofweek::Int64, period::Int64) = new(dayofweek, period)
    NthWeekdays(dayofweek::Symbol, period::Int64) = new(RDates.WEEKDAYS[dayofweek], period)
end

function apply(rdate::NthWeekdays, date::Dates.Date, cal_mgr::CalendarManager)
    wd = Dates.dayofweek(date)
    wd1st = mod(wd - mod(Dates.day(date), 7), 7) + 1
    wd1stdiff = wd1st - rdate.dayofweek
    period = wd1stdiff > 0 ? rdate.period : rdate.period - 1
    days = 7*period - wd1stdiff + 1
    return Dates.Date(Dates.year(date), Dates.month(date), days)
end
multiply(x::NthWeekdays, ::Integer) = x
Base.:-(x::NthWeekdays) = x

Base.show(io::IO, rdate::NthWeekdays) = print(io, "$(NTH_PERIODS[rdate.period]) $(uppercase(Dates.ENGLISH.days_of_week_abbr[rdate.dayofweek]))")

register_grammar!(Alt(map(Pattern, NTH_PERIODS)...) + space + weekday_short > (p,wd) -> NthWeekdays(WEEKDAYS[Symbol(wd)], PERIODS[p]))

"""
    NthLastWeekdays(dayofweek::Int64, period::Int64)
    NthLastWeekdays(dayofweek::Symbol, period::Int64)

Move to the nth last weekday in the given month and year. This is commonly used for holiday
calendars, such as the Spring Bank Holiday in the UK, which is the last Monday in May.

!!! note
    It's possible that a given period (such as the 5th Last weekday) may exist for only
    a subsection of dates. While it's a valid RDate it may not produce valid results
    when applied (and will throw an exception)

### Examples
```julia-repl
julia> RDates.NthLastWeekdays(:MON, 2) + Date(2019,1,1)
2019-01-21
julia> RDates.NthLastWeekdays(1, 2) + Date(2019,1,1)
2019-01-21
julia> rd"2nd Last MON" + Date(2019,1,1)
2019-01-21
julia> RDates.NthLastWeekdays(:MON, 5) + Date(2019,1,1)
ERROR: ArgumentError: Day: 0 out of range (1:31)
```
"""
struct NthLastWeekdays <: RDate
    dayofweek::Int64
    period::Int64

    NthLastWeekdays(dayofweek::Int64, period::Int64) = new(dayofweek, period)
    NthLastWeekdays(dayofweek::Symbol, period::Int64) = new(RDates.WEEKDAYS[dayofweek], period)
end

function apply(rdate::NthLastWeekdays, date::Dates.Date, cal_mgr::CalendarManager)
    ldom = LDOM() + date
    ldom_dow = Dates.dayofweek(ldom)
    ldom_dow_diff = ldom_dow - rdate.dayofweek
    period = ldom_dow_diff >= 0 ? rdate.period - 1 : rdate.period
    days_to_sub = 7*period + ldom_dow_diff
    days = Dates.day(ldom) - days_to_sub
    return Dates.Date(Dates.year(date), Dates.month(date), days)
end

multiply(x::NthLastWeekdays, ::Integer) = x
Base.:-(x::NthLastWeekdays) = x

Base.show(io::IO, rdate::NthLastWeekdays) = print(io, "$(NTH_LAST_PERIODS[rdate.period]) $(uppercase(Dates.ENGLISH.days_of_week_abbr[rdate.dayofweek]))")

register_grammar!(Alt(map(Pattern, NTH_LAST_PERIODS)...) + space + weekday_short > (p,wd) -> NthLastWeekdays(WEEKDAYS[Symbol(wd)], PERIODS[p]))

"""
    Weekdays(dayofweek::Int64, count::Int64, inclusive::Bool = false)
    Weekdays(dayofweek::Symbol, count::Int64, inclusive::Bool = false)

Provides a mechanism to ask for the next Saturday or the last Tuesday. The count specifies
what we're looking for. `Weekdays(:MON, 1)` will ask for the next Monday, exclusive of the
date started from. You can make it inclusive by passing the inclusive parameter with
`Weekdays(:MON, 1, true)`.

!!! note
    A count of `0` is not supported as it doesn't specify what you're actually looking for!

Incrementing the count will then be additional weeks (forward or backwards) from the single
count point.

### Examples
```julia-repl
julia> RDates.Weekdays(:WED, 1) + Date(2019,9,24) # Tuesday
2019-09-25
julia> RDates.Weekdays(3, 1) + Date(2019,9,24)
2019-09-25
julia> rd"1WED" + Date(2019,9,24)
2019-09-25
julia> RDates.Weekdays(:WED, 1) + Date(2019,9,25)
2019-10-02
julia> RDates.Weekdays(:WED, 1, true) + Date(2019,9,25)
2019-09-25
julia> rd"1WED!" + Date(2019,9,25)
2019-09-25
julia> RDates.Weekdays(:WED, -1) + Date(2019,9,24)
2019-09-18
```
"""
struct Weekdays <: RDate
    dayofweek::Int64
    count::Int64
    inclusive::Bool

    function Weekdays(dayofweek::Int64, count::Int64, inclusive::Bool = false)
        count != 0 || error("Cannot create 0 Weekdays")
        new(dayofweek, count, inclusive)
    end

    function Weekdays(dayofweek::Symbol, count::Int64, inclusive::Bool = false)
        count != 0 || error("Cannot create 0 Weekdays")
        new(WEEKDAYS[dayofweek], count, inclusive)
    end
end

function apply(rdate::Weekdays, date::Dates.Date, cal_mgr::CalendarManager)
    dayδ = Dates.dayofweek(date) - rdate.dayofweek
    weekδ = rdate.count

    if rdate.count < 0 && (dayδ > 0 || (rdate.inclusive && dayδ == 0))
        weekδ += 1
    elseif rdate.count > 0 && (dayδ < 0 || (rdate.inclusive && dayδ == 0))
        weekδ -= 1
    end

    return date + Dates.Day(weekδ*7 - dayδ)
end

multiply(x::Weekdays, count::Integer) = Weekdays(x.dayofweek, x.count * count, x.inclusive)
Base.:-(rdate::Weekdays) = Weekdays(rdate.dayofweek, -rdate.count)

function Base.show(io::IO, rdate::Weekdays)
    weekday = uppercase(Dates.ENGLISH.days_of_week_abbr[rdate.dayofweek])
    inclusive = rdate.inclusive ? "!" : ""
    print(io, "$(rdate.count)$weekday$inclusive")
end
register_grammar!(PNonZeroInt64() + weekday_short > (i,wd) -> Weekdays(WEEKDAYS[Symbol(wd)], i))
register_grammar!(PNonZeroInt64() + weekday_short + E"!" > (i,wd) -> Weekdays(WEEKDAYS[Symbol(wd)], i, true))

"""
    BizDayZero

Wrapper for a zero day count in business days, which holds the direction. Direction
can either be :next or :prev
"""
struct BizDayZero
    direction::Symbol

    function BizDayZero(direction::Symbol)
        direction in (:next, :prev) || error("unknown direction $direction for BizDayZero")
        new(direction)
    end
end

function Base.:-(x::BizDayZero)
    BizDayZero(x.direction == :next ? :prev : :next)
end

"""
    BizDays(days::Int64, calendars)
    BizDays(days::BizDayZero)

It can be handy to work in business days at times, rather than calendar days. This
allows us to move forwards or backwards `n` days.

```julia-repl
julia> cal_mgr = SimpleCalendarManager()
julia> setcalendar!(cal_mgr, "WEEKEND", WeekendCalendar())
julia> apply(RDates.BizDays(1, "WEEKEND"), Date(2021,7,9), cal_mgr)
2021-07-12
julia> apply(rd"1b@WEEKEND", Date(2021,7,9), cal_mgr)
2021-07-12
julia> apply(RDates.BizDays(-10, "WEEKEND"), Date(2021,7,9), cal_mgr)
2021-06-25
```

If the date falls on a holiday, then it is first moved forward (or backwards) to a valid
business day.

```julia-repl
julia> apply(RDates.BizDays(1, "WEEKEND"), Date(2021,7,10), cal_mgr)
2021-07-13
```

For zero business days, we could either want to move forwards or backwards. As such we
provide `BizDayZero` which can be used to provide each. By default, `0b` will move forward

```julia-repl
julia> apply(RDates.BizDays(RDates.BizDayZero(:next), "WEEKEND"), Date(2021,7,10), cal_mgr)
2021-07-12
julia> apply(RDates.BizDays(RDates.BizDayZero(:prev), "WEEKEND"), Date(2021,7,10), cal_mgr)
2021-07-09
julia> apply(rd"0b@WEEKEND", Date(2021,7,10), cal_mgr)
2021-07-12
julia> apply(rd"-0b@WEEKEND", Date(2021,7,10), cal_mgr)
2021-07-09
```
"""
@auto_hash_equals struct BizDays <: RDate
    days::Union{Int64, BizDayZero}
    calendar_names::Vector{String}

    function BizDays(days::Int64, calendar_names)
        days = days != 0 ? days : BizDayZero(:next)
        new(days, calendar_names)
    end

    BizDays(days::Int64, calendar_names::String) = BizDays(days, split(calendar_names, "|"))

    BizDays(days::BizDayZero, calendar_names) = new(days, calendar_names)
    BizDays(days::BizDayZero, calendar_names::String) = new(days, split(calendar_names, "|"))
end

function apply(rdate::BizDays, date::Dates.Date, cal_mgr::CalendarManager)
    cal = calendar(cal_mgr, rdate.calendar_names)
    if isa(rdate.days, BizDayZero)
        rounding = rdate.days.direction == :next ? HolidayRoundingNBD() : HolidayRoundingPBD()
        apply(rounding, date, cal)
    else
        rounding = rdate.days > 0 ? HolidayRoundingNBD() : HolidayRoundingPBD()
        date = apply(rounding, date, cal)

        count = rdate.days
        if rdate.days > 0
            while count > 0
                date = apply(rounding, date + Dates.Day(1), cal)
                count -= 1
            end
        elseif rdate.days < 0
            while count < 0
                date = apply(rounding, date - Dates.Day(1), cal)
                count += 1
            end
        end

        date
    end
end

function multiply(x::BizDays, count::Integer)
    x = count < 0 ? -x : x
    days = isa(x.days, BizDayZero) ? x.days : x.days * abs(count)
    return BizDays(days, x.calendar_names)
end

Base.:-(x::BizDays) = BizDays(-x.days, x.calendar_names)
register_grammar!(PPosZeroInt64() + E"b@" + PCalendarNames() > (days,calendar_name) -> BizDays(days, map(String, split(calendar_name, "|"))))


text:
module models

using ..gpflow
import ..gpflow: instantiate!, predict_f, predict_f_samples
export
    GPR,
    SGPR,
    VGP,
    SVGP,
    GPMC,
    SGPMC,
    instantiate!,
    predict_f,
    predict_f_samples

abstract type GPModel <: Model end

function predict_f(m::GPModel, Xnew)
    return m.o.predict_f(Xnew)
end
function predict_f_samples(m::GPModel, Xnew, num_samples)
    return m.o.predict_f_samples(Xnew, num_samples)
end

"""
    Gaussian Process Regression.

    This is a vanilla implementation of GP regression with a Gaussian
    likelihood. In this case inference is exact, but costs O(N^3). This means
    that we can compute the predictive distributions (predict_f, predict_y) in
    closed-form, as well as the marginal likelihood, which we use to estimate
    (optimize) the kernel parameters.

    Multiple columns of Y are treated independently, using the same kernel.

    The log likelihood of this model is sometimes referred to as the
    'marginal log likelihood', and is given by

    .. math::

       \\log p(\\mathbf y | \\mathbf f) = \\mathcal N(\\mathbf y | 0, \\mathbf K + \\sigma_n \\mathbf I)
"""
mutable struct GPR{T1,T2} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    mean_function::Union{AbstractMeanFunction, Nothing}
    name::Union{String, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    o::Union{PyObject, Nothing}
end

function GPR(
    X,
    Y,
    kern::Kernel;
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    name::Union{String, Nothing}=nothing
)
    out = GPR(X, Y, kern, mean_function, name, nothing, nothing)
    instantiate!(out)
    return out
end

function instantiate!(o::Union{GPR, Nothing})
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.GPR(o.X, o.Y, kern_, mean_function_, o.name)
    #= TODO: Shoudld we manually initialize likelihood?
    o.likelihood = likelihoods.Gaussian()
    o.likelihood.o = o.o.likelihood =#
    return o.o
end

# TODO: SGPRUpperMixin required?
"""
    Sparse Variational GP regression. The key reference is

    ::

      @inproceedings{titsias2009variational,
        title={Variational learning of inducing variables in
               sparse Gaussian processes},
        author={Titsias, Michalis K},
        booktitle={International Conference on
                   Artificial Intelligence and Statistics},
        pages={567--574},
        year={2009}
      }

"""
mutable struct SGPR{T1,T2,T3} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    feat
    mean_function::Union{AbstractMeanFunction, Nothing}
    Z::T3
    name::Union{String, Nothing}
    o::Union{PyObject, Nothing}
end

function SGPR(
    X,
    Y,
    kern::Union{Kernel, Nothing};
    feat=nothing,
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    Z=nothing,
    name::Union{String, Nothing}=nothing
)
    out = SGPR(X, Y, kern, feat, mean_function, Z, name, nothing)
    instantiate!(out)
    return out
end

function instantiate!(o::Union{SGPR, Nothing})
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.SGPR(
        o.X, o.Y, kern=kern_;
        mean_function=mean_function_,
        name=o.name
    )
    return o.o
end

"""
    This method approximates the Gaussian process posterior using a multivariate Gaussian.

    The idea is that the posterior over the function-value vector F is
    approximated by a Gaussian, and the KL divergence is minimised between
    the approximation and the posterior.

    This implementation is equivalent to svgp with X=Z, but is more efficient.
    The whitened representation is used to aid optimization.

    The posterior approximation is

    .. math::

       q(\\mathbf f) = N(\\mathbf f \\,|\\, \\boldsymbol \\mu, \\boldsymbol \\Sigma)

"""
mutable struct VGP{T1,T2} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    o::Union{PyObject, Nothing}
end

function VGP(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing
)
    out = VGP(X, Y, kern, likelihood, mean_function, num_latent, nothing)
    instantiate!(out)
    return out
end

function instantiate!(o::VGP)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.VGP(
        o.X, o.Y, kern=kern_, likelihood=likelihood_;
        mean_function=mean_function_,
        num_latent=o.num_latent
    )
    return o.o
end

"""
    This is the Sparse Variational GP (SVGP). The key reference is

    ::

    @inproceedings{hensman2014scalable,
        title={Scalable Variational Gaussian Process Classification},
        author={Hensman, James and Matthews,
                Alexander G. de G. and Ghahramani, Zoubin},
        booktitle={Proceedings of AISTATS},
        year={2015}
    }

"""
mutable struct SVGP{T1,T2,T3,T4,T5,T6} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    feat::T3
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    q_diag::Bool
    whiten::Bool
    minibatch_size::Union{Int, Nothing}
    Z::T4
    num_data::Union{Int, Nothing}
    q_mu::T5
    q_sqrt::T6
    o::Union{PyObject, Nothing}
end

function SVGP(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    feat=nothing,
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing,
    q_diag::Bool=false,
    whiten::Bool=true,
    minibatch_size::Union{Int, Nothing}=nothing,
    Z=nothing,
    num_data=nothing,
    q_mu=nothing,
    q_sqrt=nothing
)

    out = SVGP(
        X,
        Y,
        kern,
        likelihood,
        feat,
        mean_function,
        num_latent,
        q_diag,
        whiten,
        minibatch_size,
        Z,
        num_data,
        q_mu,
        q_sqrt,
        nothing
    )
    instantiate!(out)
    return out
end

function instantiate!(o::SVGP)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.SGVP(
        o.X,
        o.Y,
        kern=kern_,
        likelihood=likelihood_;
        feat=o.feat,
        mean_function=mean_function_,
        num_latent=o.num_latent,
        q_diag=o.q_diag,
        whiten=o.whiten,
        minibatch_size=o.minibatch_size,
        Z=o.Z,
        num_data=o.num_data,
        q_mu=o.q_mu,
        q_sqrt=o.q_sqrt,
    )
    return o.o
end

mutable struct GPMC{T1,T2} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    o::Union{PyObject, Nothing}
end

function GPMC(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing,
)
    out = GPMC(
        X,
        Y,
        kern,
        likelihood,
        mean_function,
        num_latent,
        nothing
    )
    instantiate!(out)
    return out
end

function instantiate!(o::GPMC)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.GPMC(
        o.X,
        o.Y,
        kern=kern_,
        likelihood=likelihood_;
        mean_function=mean_function_,
        num_latent=o.num_latent
    )
    return o.o
end

"""
    This is the Sparse Variational GP using MCMC (SGPMC). The key reference is

    ::

      @inproceedings{hensman2015mcmc,
        title={MCMC for Variatinoally Sparse Gaussian Processes},
        author={Hensman, James and Matthews, Alexander G. de G.
                and Filippone, Maurizio and Ghahramani, Zoubin},
        booktitle={Proceedings of NIPS},
        year={2015}
      }

    The latent function values are represented by centered
    (whitened) variables, so

    .. math::
       :nowrap:

       \\begin{align}
       \\mathbf v & \\sim N(0, \\mathbf I) \\\\
       \\mathbf u &= \\mathbf L\\mathbf v
       \\end{align}

    with

    .. math::
        \\mathbf L \\mathbf L^\\top = \\mathbf K

"""
mutable struct SGPMC{T1,T2,T3,T4} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    feat::T3
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    Z::T4
    o::Union{PyObject, Nothing}
end

function SGPMC(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    feat=nothing,
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing,
    Z=nothing
    )
    out = SGPMC(
        X,
        Y,
        kern,
        likelihood,
        feat,
        mean_function,
        num_latent,
        Z,
        nothing
    )
    instantiate!(out)
    return out
end

function instantiate!(o::SGPMC)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.SGPMC(
        o.X,
        o.Y,
        kern=kern_,
        likelihood=likelihood_;
        fwat=o.feat,
        mean_function=mean_function_,
        num_latent=o.num_latent,
        Z=o.Z
    )
    return o.o
end

end # module


text:
function divergence_kernel(
    u, v, w,
    dxi, dyi, dzi,
    is, ie, js, je, ks, ke)

    divmax = 0
    @fast3d begin
        div = @fd (u, v, w) gradx(u) + grady(v) + gradz(w)
        divmax = max(divmax, abs(div))
    end

    return divmax
end

function cfl_kernel(
    u, v, w,
    dxi, dyi, dzi, dt,
    is, ie, js, je, ks, ke)

    cflmax = 0
    @fast3d begin
        cfl = @fd (u, v, w) abs(interpx(u))*dxi + abs(interpy(v))*dyi + abs(interpz(w))*dzi[k]
        cflmax = max(cflmax, cfl)
    end

    return cflmax*dt
end

function calc_divergence(f::Fields, g::Grid)
    div = divergence_kernel(
        f.u, f.v, f.w,
        g.dxi, g.dyi, g.dzi,
        g.is, g.ie, g.js, g.je, g.ks, g.ke)
end

function calc_cfl(f::Fields, g::Grid, t::Timeloop)
    cfl = cfl_kernel(
        f.u, f.v, f.w,
        g.dxi, g.dyi, g.dzi, t.dt,
        g.is, g.ie, g.js, g.je, g.ks, g.ke)
end


text:
"""
    ValidationError

Abstract type for validation errors.
"""
abstract type ValidationError end

errors = (
    :UnknownFragment,
    :UnusedFragment,
    :RepeatedFragmentDefinition,
    :RepeatedOperationDefinition,
    :AnonymousOperationNotAlone,
    :RepeatedVariableDefinition,
    :RepeatedArgumentName,
    :UnknownVariable,
    :UnusedVariable,
    :RepeatedDirectiveName,
    :RepeatedInputObjectField,
)
for err in errors
    @eval begin
        Base.@kwdef struct $err <: ValidationError
            message::String
            locations::Union{Vector{Loc}, Nothing} = nothing
        end
    end
end

Base.show(io::IO, ::MIME"text/plain", err::ValidationError) = print_error(io, err)
function print_error(io, err)
    printstyled(io, typeof(err), color=Base.error_color())
    printstyled(io, "\n      message: ", err.message, color=Base.error_color())
    if !isnothing(err.locations) && !isempty(err.locations)
        if length(err.locations) == 1
            printstyled(io, "\n     location: $(err.locations[1])", color=Base.error_color())
        else
            printstyled(io, "\n  location(s): $(err.locations[1])", color=Base.error_color())
            for i in 2:length(err.locations)
                printstyled(io, "\n               $(err.locations[i])", color=Base.error_color())
            end
        end
    end
end

struct ValidationException <: Exception
    errors::Vector{ValidationError}
end
function Base.showerror(io::IO, ex::ValidationException)
    printstyled(io, "Validation Failed\n", color=Base.error_color(), bold=true)
    for err in ex.errors
        print(io, "\n")
        print_error(io, err)
        print(io, "\n")
    end
end


text:
# Logistic map cobweb and timeseries
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("colorscheme.jl"))
using GLMakie, DynamicalSystems, InteractiveDynamics

lo = Systems.logistic(0.4; r=rrange[1])
fkwargs = [(linewidth = 4.0, color = COLORSCHEME[i+1]) for i in 1:3]

rrange = 1:0.001:4.0
fig = interactive_cobweb(lo, rrange, 2; fkwargs)
record_interaction(fig, projectdir("animations", "4", "logistic_cobweb.mp4"))

# %% the second range is a convenience for intermittency example of logistic
rrange = (rc = 1 + sqrt(8); [rc - 1e-3, rc - 1e-5, rc])
fig = interactive_cobweb(lo, rrange, 3; fkwargs)
record_interaction(
    fig, projectdir("animations", "4", "logistic_intermittency.mp4");
    total_time=15, sleep_time = 2
)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
dates = Date(1960,1,1):Date(2015,9,25)

procIds = addprocs(3)

@everywhere using DataFrames
@everywhere using TimeSeries
@everywhere using EconDatasets

    println("\nData will be downloaded starting ", dates[1], " and ending ",
            dates[end], ".\n\n",
            "In order to change the time period,",
            " simply update the first line of \nfunction getSP500 ",
            "in getDataset/getSP500 according to your needs.\n",
            "Per default, getSP500 makes use of Julia's parallel ",
            "processing \ncapabilities and uses 3 cores to ",
            "download and process the data. \nIn order to change ",
            "this setting, simply modify the command addprocs(3).\n\n",
            "You can also find a detailed description of the ",
            "downloading procedure at\n",
            "http://grollchristian.wordpress.com/2014/09/05/sp500-data-download-julia/")

## load WikiPedia stock ticker symbols
constituents = readcsv(joinpath(Pkg.dir("EconDatasets"), "data",
                                "SP500TickerSymbols.csv"))

tickerSymb = String[ticker for ticker in constituents]

## measure time
t0 = time()

@time vals = readYahooAdjClose(dates, tickerSymb, :d)

t1 = time()
elapsedTime = t1-t0
mins, secs = divrem(elapsedTime, 60)

valsTn = convert(Timenum, vals)
pathToStore = joinpath(Pkg.dir("EconDatasets"), "data", "SP500.csv")
writetimearray(valsTn, pathToStore)

println("elapsed time: ", int(mins), " minutes, ", ceil(secs), " seconds")

rmprocs(procIds)


text:
export SolitaryWaveWhitham, SolitaryWhitham

"""
    SolitaryWaveWhitham(param; kwargs)

Compute the Whitham solitary wave with prescribed velocity.

# Argument
- `param :: NamedTuple`: parameters of the problem containing velocity `c` \
and dimensionless parameters `ϵ` and `μ`, \
and mesh size `L` and number of collocation points `N`;

## Keywords (optional)
- `guess :: Vector{Real}`: initial guess for the surface deformation (if not provided, the exact formula for KdV is used);
- `x₀ :: Real`: center of solitary wave (if guess is not provided);
- `iterative :: Bool`: inverts Jacobian through GMRES if `true`, LU decomposition if `false`;
- `verbose :: Bool`: prints numerical errors at each step if `true`;
- `max_iter :: Int`: maximum number of iterations of the Newton algorithm;
- `tol :: Real`: general tolerance (default is `1e-10`);
- `ktol :: Real`: tolerance of the Krasny filter (default is `0`, i.e. no filtering);
- `gtol :: Real`: relative tolerance of the GMRES algorithm (default is `1e-10`);
- `dealias :: Int`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `q :: Real`: Newton algorithm modified with
`u_{n+1}=q*u_{n+1}+(1-q)*u_n`
(default is `1`);
- `α :: Real`: adds `α` times spectral projection onto the Kernel to the Jacobian;
- `KdV :: Bool`: if `true` computes the KdV (instead of Whitham) solitary wave.
# Return values
`(u,mesh)` with
- `u :: Vector{Float64}`: the solution;
- `mesh :: Mesh`: mesh collocation points.

"""
function SolitaryWaveWhitham(
                param :: NamedTuple;
                guess = zeros(0) :: Vector{Float64},
                x₀ = 0 :: Real,
                iterative = false :: Bool,
                verbose = false :: Bool,
                max_iter = 20 :: Int,
                tol = 1e-10 :: Real,
                gtol = 1e-10 :: Real,
                ktol = 0 :: Real,
                dealias = 0 :: Int,
                q=1 :: Real,
                α=0 :: Real,
                KdV = false :: Bool)

        mesh=Mesh(param)
        c = param.c
        ϵ = param.ϵ
        μ = param.μ
        if abs(c)<1
                @error("The velocity must be greater than 1 (in absolute value).")
        end

        if guess == zeros(0)
                # Using the exact formula for the KdV solitary wave as initial guess
                guess = 2*(c-1)/ϵ*sech.(sqrt(3*2*(c-1)/μ)/2*(mesh.x.-x₀)).^2
        end

        k  = mesh.k
        Dx =  1im * k
        if KdV == true
                F₁ = 1 .-μ/6*k.^2
        else
                F₁ = sqrt.(tanh.(sqrt(μ)*abs.(k))./(sqrt(μ)*abs.(k)))
                F₁[1]=1
        end

        if dealias == 0
                Π = ones(size(k)) # no dealiasing (Π=Id)
        else
                K = (mesh.kmax-mesh.kmin)/(2+dealias)
                Π = abs.(k) .<= K # Dealiasing low-pass filter
        end
        krasny(k) = (abs.(k).> ktol ).*k
        krasny!(k) = k[abs.(k).< ktol ].=0

        function filter( v :: Vector{Float64} )
                real.(ifft(krasny(Π.*fft(v))))
        end
        function filter( v :: Vector{Complex{Float64}} )
                ifft(krasny(Π.*fft(v)))
        end

        function F( u :: Vector{Float64} )
            -c*u+real.(ifft(F₁.*fft(u)))+3*ϵ/4*u.^2
        end
        function Fabs( u :: Vector{Float64} )
            c*abs.(u)+abs.(ifft(F₁.*fft(u)))+3*ϵ/4*u.^2
        end

        if iterative == false
                k = mesh.k
                x = mesh.x
                x₀ = mesh.x[1]
                FFT = exp.(-1im*k*(x.-x₀)');
                IFFT = exp.(1im*k*(x.-x₀)')/length(x);
                M = real.(IFFT*(Diagonal( F₁)*FFT))
                function JacF( v , dxv  )
                    Symmetric(M+Diagonal(3*ϵ/2*v .-c)+ α*dxv*dxv')
                end
        else
                function JacFfast( v ,dxv )
                        dF(φ) = -c*φ+real.(ifft(F₁.*fft(φ)))+3*ϵ/2*v.*φ+ α*dot(dxv,φ)*dxv
                        return LinearMap(dF, length(v); issymmetric=true, ismutating=false)
                end
        end

        flag=0
        iter = 0
        u = filter(guess)
        du = similar(u)
        fu = similar(u)
        dxu = similar(u)
        for i in range(0, stop=max_iter)
                dxu .= real.(ifft(Dx.*fft(u)))
                dxu ./= norm(dxu,2)
                fu .= F(u)
                relerr = norm(fu,Inf)/norm(Fabs(u),Inf)
                abserr = norm(fu,Inf)
                if relerr < tol
    			@info string("Converged : relative error ",relerr," in ",i," steps\n")
    			break
    		elseif verbose == true
                        print(string("absolute error at step ",i,": ",abserr,"\n"))
                        print(string("relative error at step ",i,": ",relerr,"\n"))

    		end
                if i == max_iter
                        @warn string("The algorithm did not converge after ",i," steps: final relative error is ",relerr,"\n")
                        break
                end

                if iterative == false
                        du .= JacF(u,dxu) \ fu
                else
                        du .= gmres( JacFfast(u,dxu) , fu ; verbose = verbose, reltol = gtol )
                end
    		u .-= q*filter(du)
        end
        return (u,mesh)
end

"""
    SolitaryWhitham(param; kwargs)

Build the initial data associated with `SolitaryWaveWhitham(param; kwargs)`, of type `InitialData`,
to be used in initial-value problems `Problem(model, initial::InitialData, param)`.

---
    SolitaryWhitham(c; ϵ=1,μ=1,N=2^12,kwargs)
Build the initial data with velocity `c`, dimensionless parameters `ϵ` and `μ`, and number of collocation points `N`, \
and `kwargs` the other (optional) keyword arguments as above.
"""
struct SolitaryWhitham <: InitialData

	η
	v
	label :: String
	info  :: String

	function SolitaryWhitham(param::NamedTuple;
			guess = zeros(0) :: Vector{Float64},
			x₀ = 0 :: Real,
			iterative = false :: Bool,
			verbose = false :: Bool,
			max_iter = 20 :: Int,
			tol = 1e-10 :: Real,
			gtol = 1e-10 :: Real,
			ktol = 0 :: Real,
			dealias = 0 :: Int,
			q=1 :: Real,
			α=0 :: Real,
			KdV = false :: Bool)
                (u,mesh)=SolitaryWaveWhitham(param;
                                guess=guess,x₀=x₀,KdV=KdV,
                                iterative=iterative,verbose=verbose,max_iter=max_iter,
                                tol=tol,ktol=ktol,gtol=gtol,dealias=dealias,q=q,α=α)
		ϵ=param.ϵ
		η=((ϵ*u/2 .+1).^2 .-1)/ϵ
		init = Init(mesh,η,u)
		if KdV == false
        		model = "Whitham"
		else
			model = "KdV"
		end
                label = "$model solitary wave"
		info = "Solitary travelling wave for the $model model.\n\
		├─velocity c = $(param.c)\n\
		└─maximum h₀ = $(maximum(η)) (from rest state)."

		new( init.η,init.v,label,info  )
	end

	function SolitaryWhitham(
                        c::Real; ϵ=1::Real,μ=1::Real,N=2^12::Int,
			guess = zeros(0) :: Vector{Float64},
	                x₀ = 0 :: Real,
	                iterative = false :: Bool,
	                verbose = false :: Bool,
	                max_iter = 20 :: Int,
	                tol = 1e-10 :: Real,
	                gtol = 1e-10 :: Real,
	                ktol = 0 :: Real,
	                dealias = 0 :: Int,
	                q=1 :: Real,
	                α=0 :: Real,
	                KdV = false :: Bool)
		L=200/sqrt(3*(c^2-1)/(c^2)/μ)/2
		xmin,xmax = x₀-L,x₀+L;
		param=(ϵ=ϵ,μ=μ,c=c,L=L,xmin=xmin,xmax=xmax,N=N)
		sol=SolitaryWhitham(param;
                                guess=guess,x₀=x₀,KdV=KdV,
                                iterative=iterative,verbose=verbose,max_iter=max_iter,
                                tol=tol,ktol=ktol,gtol=gtol,dealias=dealias,q=q,α=α)
		new( sol.η,sol.v,sol.label,sol.info  )
	end

end


text:
struct TsptwTrajectoryState <: GraphTrajectoryState
    fg::GeometricFlux.FeaturedGraph
    variableIdx::Int
    possibleValuesIdx::AbstractVector{Int}
end

Flux.@functor TsptwTrajectoryState