text:
export TFTS, TAAFT

"""
    TFTS(fϵ::Real)

A truncated Fourier transform surrogate[^Nakamura2006] (TFTS).

TFTS surrogates are generated by leaving some frequencies untouched when performing the
phase shuffling step (as opposed to randomizing all frequencies, like for
[`RandomFourier`](@ref) surrogates).

These surrogates were designed to deal with data with irregular fluctuations superimposed
over long term trends (by preserving low frequencies)[^Nakamura2006]. Hence, TFTS surrogates
can be used to test the null hypothesis that the signal is a stationary linear system
generated the irregular fluctuations part of the signal[^Nakamura2006].

## Controlling the truncation of the spectrum

The truncation parameter `fϵ ∈ [-1, 0) ∪ (0, 1]` controls which parts of the spectrum are preserved.

- If `fϵ > 0`, then `fϵ` indicates the ratio of high frequency domain to the entire frequency domain.
    For example, `fϵ = 0.5` preserves 50% of the frequency domain (randomizing the higher
    frequencies, leaving low frequencies intact).
- If `fϵ < 0`, then `fϵ` indicates ratio of low frequency domain to the entire frequency domain.
    For example, `fϵ = -0.2` preserves 20% of the frequency domain (leaving higher frequencies intact,
    randomizing the lower frequencies).
- If `fϵ ± 1`, then all frequencies are randomized. The method is then equivalent to
    [`RandomFourier`](@ref).

The appropriate value of `fϵ` strongly depends on the data and time series length, and must be
manually determined[^Nakamura2006], for example by comparing periodograms for the time series and
the surrogates.

[^Nakamura2006]: Nakamura, Tomomichi, Michael Small, and Yoshito Hirata. "Testing for nonlinearity in irregular fluctuations with long-term trends." Physical Review E 74.2 (2006): 026205.
"""
struct TFTS <: Surrogate
    fϵ::Real

    function TFTS(fϵ::Real)
        if !(0 < fϵ ≤ 1) && !(-1 ≤ fϵ < 0)
            throw(ArgumentError("`fϵ` must be on the interval [-1, 0) ∪ (0, 1] (positive if preserving high frequencies, negative if preserving low frequencies)"))
        end
        new(fϵ)
    end
end

function surrogenerator(x, method::TFTS, rng = Random.default_rng())
    # Pre-plan Fourier transforms
    forward = plan_rfft(x)
    inverse = plan_irfft(forward*x, length(x))

    # Pre-compute 𝓕
    𝓕 = forward * x

    # Polar coordinate representation of the Fourier transform
    rx = abs.(𝓕)
    ϕx = angle.(𝓕)
    n = length(𝓕)

    # These are updated during iteration procedure
    𝓕new = Vector{Complex{Float64}}(undef, length(𝓕))
    𝓕s = Vector{Complex{Float64}}(undef, length(𝓕))
    ϕs = Vector{Complex{Float64}}(undef, length(𝓕))

    init = (forward = forward, inverse = inverse,
        rx = rx, ϕx = ϕx, n = n,
        𝓕new = 𝓕new, 𝓕s = 𝓕s, ϕs = ϕs)

    return SurrogateGenerator(method, x, similar(x), init, rng)
end

function (sg::SurrogateGenerator{<:TFTS})()
    x, s = sg.x, sg.s
    fϵ = sg.method.fϵ
    L = length(x)

    init_fields = (:forward, :inverse,
        :rx, :ϕx, :n,
        :𝓕new, :𝓕s, :ϕs)

    forward, inverse,
        rx, ϕx, n,
        𝓕new, 𝓕s, ϕs = getfield.(Ref(sg.init), init_fields)

    # Surrogate starts out as a random permutation of x
    s .= x[StatsBase.sample(sg.rng, 1:L, L; replace = false)]
    𝓕s .= forward * s
    ϕs .= angle.(𝓕s)

    # Updated spectrum is the old amplitudes with the mixed phases.
    if fϵ > 0
        # Frequencies are ordered from lowest when taking the Fourier
        # transform, so by keeping the 1:n_ni first phases intact,
        # we are only randomizing the high-frequency components of the
        # signal.
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[1:n_preserve] .= ϕx[1:n_preserve]
    elseif fϵ < 0
        # Do the exact opposite to preserve high-frequencies
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[end-n_preserve+1:end] .= ϕx[end-n_preserve+1:end]
    end

    𝓕new .= rx .* exp.(ϕs .* 1im)
    s .= inverse * 𝓕new

    return s
end

"""
    TAAFT(fϵ)

An truncated version of the amplitude-adjusted-fourier-transform surrogate[^Theiler1991][^Nakamura2006].

The truncation parameter and phase randomization procedure is identical to [`TFTS`](@ref), but here an
additional step of rescaling back to the original data is performed. This preserves the
amplitude distribution of the original data.

[^Theiler1991]: J. Theiler, S. Eubank, A. Longtin, B. Galdrikian, J. Farmer, Testing for nonlinearity in time series: The method of surrogate data, Physica D 58 (1–4) (1992) 77–94.
[^Nakamura2006]: Nakamura, Tomomichi, Michael Small, and Yoshito Hirata. "Testing for nonlinearity in irregular fluctuations with long-term trends." Physical Review E 74.2 (2006): 026205.
"""
struct TAAFT <: Surrogate
    fϵ::Real

    function TAAFT(fϵ::Real)
        fϵ != 0 || throw(ArgumentError("`fϵ` must be on the interval [-1, 0) ∪ (0, 1] (positive if preserving high frequencies, negative if preserving low frequencies)"))
        new(fϵ)
    end
end

function surrogenerator(x, method::TAAFT, rng = Random.default_rng())
    init = (
        gen = surrogenerator(x, TFTS(method.fϵ), rng),
        x_sorted = sort(x),
        idxs = collect(1:length(x)),
        perm = zeros(Int, length(x)),
    )

    s = similar(x)
    return SurrogateGenerator(method, x, s, init, rng)
end

function (taaft::SurrogateGenerator{<:TAAFT})()
    sg = taaft.init.gen
    x_sorted, idxs, perm = taaft.init.x_sorted, taaft.init.idxs, taaft.init.perm

    x, s = sg.x, sg.s
    fϵ = sg.method.fϵ
    L = length(x)

    init_fields = (:forward, :inverse,
        :rx, :ϕx, :n,
        :𝓕new, :𝓕s, :ϕs)

    forward, inverse,
        rx, ϕx, n,
        𝓕new, 𝓕s, ϕs = getfield.(Ref(sg.init), init_fields)

    # Surrogate starts out as a random permutation of x
    StatsBase.sample!(sg.rng, idxs, perm, replace = false)
    permuted_x_into_s!(s, x, perm)
    𝓕s .= forward * s
    ϕs .= angle.(𝓕s)

    # Updated spectrum is the old amplitudes with the mixed phases.
    if fϵ > 0
        # Frequencies are ordered from lowest when taking the Fourier
        # transform, so by keeping the 1:n_ni first phases intact,
        # we are only randomizing the high-frequency components of the
        # signal.
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[1:n_preserve] .= @view ϕx[1:n_preserve]
    elseif fϵ < 0
        # Do the exact opposite to preserve high-frequencies
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[end-n_preserve+1:end] .= @view ϕx[end-n_preserve+1:end]
    end

    𝓕new .= rx .* exp.(ϕs .* 1im)
    s .= inverse * 𝓕new

    s[sortperm(s)] .= x_sorted
    return s
end

function permuted_x_into_s!(s, x, perm)
    k = 1
    for i in perm
        s[k] = x[i]
        k += 1
    end
end


text:
module JuLIPMaterials

include("common.jl")

include("CLE.jl")

include("LGF.jl")

include("FCC.jl")

include("BCC.jl")

include("PoleExpansion.jl")

include("Testing.jl")

# include("CauchyBorn.jl")

include("Si.jl")

# include("TriangularLattice.jl")

# include("PointDefects.jl")

include("fracture.jl")

end # module

ext: jl


text:
using JLD2

# Read in from JLD
tpf_main_input = load("reference/tpf_main_inputs.jld2")
data = tpf_main_input["data"]
TTT = tpf_main_input["TTT"]
RRR = tpf_main_input["RRR"]
CCC = tpf_main_input["CCC"]
ZZ = tpf_main_input["ZZ"]
DD = tpf_main_input["DD"]
F_ϵ = tpf_main_input["F_epsilon"]
F_u = tpf_main_input["F_u"]
s_init = tpf_main_input["s_init"]

# Tune algorithm
tuning = Dict(:r_star => 2., :c_init => 0.3, :target_accept_rate => 0.4,
              :resampling_method => :systematic, :n_mh_steps => 1,
              :n_particles => 1000, :n_presample_periods => 0,
              :allout => true)

# Define Φ and Ψ (can't be saved to JLD)
Φ(s_t::AbstractVector{Float64}, ϵ_t::AbstractVector{Float64}) = TTT*s_t + RRR*ϵ_t + CCC
Ψt(s_t::AbstractVector{Float64}, t) = t < 100 ? ZZ*s_t + DD : DD
Ψ(x) = Ψt(x, 47)

# Load in test inputs and outputs
test_file_inputs = load("reference/tpf_aux_inputs.jld2")
test_file_outputs = load("reference/tpf_aux_outputs.jld2")

φ_old = test_file_inputs["phi_old"]
norm_weights = test_file_inputs["norm_weights"]
coeff_terms = test_file_inputs["coeff_terms"]
log_e_1_terms = test_file_inputs["log_e_1_terms"]
log_e_2_terms = test_file_inputs["log_e_2_terms"]
inc_weights = test_file_inputs["inc_weights"]
HH = cov(F_u)
s_t_nontemp = test_file_inputs["s_t_nontemp"]

weight_kernel!(coeff_terms, log_e_1_terms, log_e_2_terms, φ_old, Ψ, data[:, 47], s_t_nontemp, det(HH), inv(HH);
               initialize = false, parallel = false)
φ_new = next_φ(φ_old, coeff_terms, log_e_1_terms, log_e_2_terms, length(data[:,47]), tuning[:r_star], 2)

correction!(inc_weights, norm_weights, φ_new, coeff_terms, log_e_1_terms, log_e_2_terms, length(data[:,47]))

@testset "Correction and Auxiliary Tests" begin
    @test coeff_terms[1] ≈ test_file_outputs["coeff_terms"][1]
    @test log_e_1_terms[1] ≈ test_file_outputs["log_e_1_terms"][1]
    @test log_e_2_terms[1] ≈ test_file_outputs["log_e_2_terms"][1]
    @test φ_new ≈ test_file_outputs["phi_new"]
    @test inc_weights[1] ≈ test_file_outputs["inc_weights"][1]
end

# Incorrect measurement equation
Ψtinc(s_t::AbstractVector{Float64}, t) = t < 46 ? ZZ*s_t + DD : ZZ*s_t + DD + [0.; 1; 0.];
Ψinc(x) = Ψtinc(x, 47)

φ_old = test_file_inputs["phi_old"]
norm_weights = test_file_inputs["norm_weights"]
coeff_terms = test_file_inputs["coeff_terms"]
log_e_1_terms = test_file_inputs["log_e_1_terms"]
log_e_2_terms = test_file_inputs["log_e_2_terms"]
inc_weights = test_file_inputs["inc_weights"]
HH = cov(F_u)
s_t_nontemp = test_file_inputs["s_t_nontemp"]

weight_kernel!(coeff_terms, log_e_1_terms, log_e_2_terms, φ_old, Ψinc, data[:, 47], s_t_nontemp, det(HH), inv(HH);
               initialize = false, parallel = false)
φ_new = next_φ(φ_old, coeff_terms, log_e_1_terms, log_e_2_terms, length(data[:,47]), tuning[:r_star], 2)

correction!(inc_weights, norm_weights, φ_new, coeff_terms, log_e_1_terms, log_e_2_terms, length(data[:,47]))
@testset "Correction and Auxiliary Tests with Dynamically Wrong Measurement Equation" begin
    @test !(log_e_1_terms[1] ≈ test_file_outputs["log_e_1_terms"][1])
    @test !(log_e_2_terms[1] ≈ test_file_outputs["log_e_2_terms"][1])
    @test !(φ_new ≈ test_file_outputs["phi_new"])
    @test !(inc_weights[1] ≈ test_file_outputs["inc_weights"][1])
end

## Selection Tests
Ψt(s_t::AbstractVector{Float64}, t) = t < 100 ? ZZ*s_t + DD : DD
Ψ(x) = Ψt(x, 47)
s_t1_temp = test_file_inputs["s_t1_temp"]
ϵ_t = test_file_inputs["eps_t"]

Random.seed!(47)
selection!(norm_weights, s_t1_temp, s_t_nontemp,ϵ_t, resampling_method = tuning[:resampling_method])
@testset "Selection Tests" begin
    @test s_t1_temp[1] ≈ test_file_outputs["s_t1_temp"][1]
    @test s_t_nontemp[1] ≈ test_file_outputs["s_t_nontemp"][1]
    @test ϵ_t[1] ≈ test_file_outputs["eps_t"][1]
end

## Mutation Tests
QQ = cov(F_ϵ)
accept_rate = test_file_inputs["accept_rate"]
c = test_file_inputs["c"]

c = update_c(c, accept_rate, tuning[:target_accept_rate])
Random.seed!(47)
StateSpaceRoutines.mutation!(Φ, Ψ, QQ, det(HH), inv(HH), φ_new, data[:,47], s_t_nontemp, s_t1_temp, ϵ_t, c, tuning[:n_mh_steps])

@testset "Mutation Tests" begin
    @test s_t_nontemp[1] ≈ test_file_outputs["s_t_nontemp_mutation"][1]
    @test ϵ_t[1] ≈ test_file_outputs["eps_t_mutation"][1]
end

## Whole TPF Tests
Random.seed!(47)
out_no_parallel = tempered_particle_filter(data, Φ, Ψt, F_ϵ, F_u, s_init; tuning..., verbose = :none, parallel = false, dynamic_measurement = true)
Random.seed!(47)
out_parallel_one_worker = tempered_particle_filter(data, Φ, Ψt, F_ϵ, F_u, s_init; tuning..., verbose = :none, parallel = true, dynamic_measurement = true)
@testset "TPF tests" begin
    @test out_no_parallel[1] ≈ -302.99967306704133
    # This is different than in Julia6 because @distributed seeds differently than @parallel
    @test out_parallel_one_worker[1] ≈ -303.64727963725073 # Julia 6 was: -306.8211172094595
end

nothing


text:
type MEXCLPDeployment <: DeploymentModel
    m::JuMP.Model
    x::Vector{JuMP.Variable}
end
deployment(m::MEXCLPDeployment) = Int[round(Int,x) for x in JuMP.getvalue(m.x)]

function MEXCLPDeployment(p::DeploymentProblem,
                          q::Float64; # busy fraction
                          max_amb::Int = 0,
                          tol = params.δ,
                          solver = GurobiSolver(OutputFlag=0))
    (max_amb == 0) && (max_amb = p.nambulances)
    @assert max_amb > 0
    demand = vec(mean(p.demand[p.train,:],1))
    @assert length(demand) == p.nregions

    I = 1:p.nlocations
    J = 1:p.nregions
    K = 1:max_amb

    m = JuMP.Model(solver=solver)
    JuMP.@variable(m, x[1:p.nlocations] >= 0, Int)
    JuMP.@variable(m, z[1:p.nregions, 1:max_amb], Bin)

    JuMP.@objective(m, Max, sum(demand[j]*(1-q)*(q^k)*z[j,k] for j in J, k in K))

    JuMP.@constraint(m, sum(x[i] for i in I) <= p.nambulances)

    for j in J # coverage over all regions
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >= 1)
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >=
                            sum(z[j,k] for k in K))
    end

    MEXCLPDeployment(m, x)
end

solve(model::MEXCLPDeployment) = JuMP.solve(model.m)
solve(model::MEXCLPDeployment, p::DeploymentProblem) = JuMP.solve(model.m)