#
# Project : Gardenia
# Source  : san.jl
# Author  : Li Huang (huangli@caep.cn)
# Status  : Unstable
#
# Last modified: 2024/09/30
#

#=
### *Customized Structs* : *StochSK Solver*
=#

"""
    StochSKElement

Mutable struct. It is used to record the field configurations, which will
be sampled by Monte Carlo sweeping procedure.

In the present implementation of StochSK solver, the amplitudes of the Œ¥
functions are fixed. But in principles, they could be sampled in the Monte
Carlo procedure.

### Members
* P -> It means the positions of the Œ¥ functions.
* A -> It means the weights / amplitudes of the Œ¥ functions.
* W -> It denotes the window that is used to shift the Œ¥ functions.
"""
mutable struct StochSKElement
    P :: Vector{I64}
    A :: F64
    W :: I64
end

"""
    StochSKContext

Mutable struct. It is used within the StochSK solver only.

### Members
* G·µ•     -> Input data for correlator.
* G·µß     -> Generated correlator.
* œÉ¬π     -> Actually 1.0 / œÉ¬π.
* allow  -> Allowable indices.
* grid   -> Imaginary axis grid for input data.
* mesh   -> Real frequency mesh for output spectrum.
* kernel -> Default kernel function.
* Aout   -> Calculated spectral function.
* œá¬≤     -> Current goodness-of-fit function.
* œá¬≤min  -> Mininum goodness-of-fit function.
* œá¬≤vec  -> Vector of goodness-of-fit function.
* Œò      -> Current Œò parameter.
* Œòvec   -> Vector of Œò parameter.
"""
mutable struct StochSKContext
    G·µ•     :: Vector{F64}
    G·µß     :: Vector{F64}
    œÉ¬π     :: Vector{F64}
    allow  :: Vector{I64}
    grid   :: AbstractGrid
    mesh   :: AbstractMesh
    kernel :: Array{F64,2}
    Aout   :: Vector{F64}
    œá¬≤     :: F64
    œá¬≤min  :: F64
    œá¬≤vec  :: Vector{F64}
    Œò      :: F64
    Œòvec   :: Vector{F64}
end

#=
### *Global Drivers*
=#

"""
    solve(S::StochSKSolver, rd::RawData)

Solve the analytic continuation problem by using the stochastic analytic
continuation algorithm (A. W. Sandvik's version). It is the driver for the
StochSK solver.

If the input correlators are bosonic, this solver will return A(œâ) / œâ
via `Aout`, instead of A(œâ). At this time, `Aout` is not compatible with
`Gout`. If the input correlators are fermionic, this solver will return
A(œâ) in `Aout`. Now it is compatible with `Gout`. These behaviors are just
similar to the MaxEnt, StochAC, and StochOM solvers.

Now the StochSK solver supports both continuous and Œ¥-like spectra.

### Arguments
* S -> A StochSKSolver struct.
* rd -> A RawData struct, containing raw data for input correlator.

### Returns
* mesh -> Real frequency mesh, œâ.
* Aout -> Spectral function, A(œâ).
* Gout -> Retarded Green's function, G(œâ).
"""
function solve(S::StochSKSolver, rd::RawData)
    nmesh = get_b("nmesh")
    nwarm = get_k("nwarm")

    println("[ StochSK ]")
    MC, SE, SC = init(S, rd)

    # Parallel version
    if nworkers() > 1
        #
        println("Using $(nworkers()) workers")
        #
        # Copy configuration dicts
        p1 = deepcopy(PBASE)
        p2 = deepcopy(PStochSK)
        #
        # Launch the tasks one by one
        ùêπ = Future[]
        for i = 1:nworkers()
            ùëì = @spawnat i + 1 prun(S, p1, p2, MC, SE, SC)
            push!(ùêπ, ùëì)
        end
        #
        # Wait and collect the solutions
        sol = []
        for i = 1:nworkers()
            wait(ùêπ[i])
            push!(sol, fetch(ùêπ[i]))
        end
        #
        # Average the solutions
        Aout = zeros(F64, nmesh)
        œá¬≤out = zeros(F64, nwarm)
        Œòout = zeros(F64, nwarm)
        for i in eachindex(sol)
            a, b, c = sol[i]
            @. Aout = Aout + a / nworkers()
            @. œá¬≤out = œá¬≤out + b / nworkers()
            @. Œòout = Œòout + c / nworkers()
        end
        #
        # Postprocess the solutions
        Gout = last(SC, Aout, œá¬≤out, Œòout)
        #
    # Sequential version
    else
        #
        Aout, œá¬≤out, Œòout = run(MC, SE, SC)
        Gout = last(SC, Aout, œá¬≤out, Œòout)
        #
    end

    return SC.mesh.mesh, Aout, Gout
end

"""
    init(S::StochSKSolver, rd::RawData)

Initialize the StochSK solver and return the StochSKMC, StochSKElement,
and StochSKContext structs. Please don't call this function directly.

### Arguments
* S -> A StochSKSolver struct.
* rd -> A RawData struct, containing raw data for input correlator.

### Returns
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.
"""
function init(S::StochSKSolver, rd::RawData)
    # Initialize possible constraints.
    # The array allow contains all the possible indices for Œ¥ functions.
    fmesh = calc_fmesh(S)
    allow = constraints(S, fmesh)

    # Prepare input data
    G·µ•, œÉ¬π = init_iodata(S, rd)
    println("Postprocess input data: ", length(œÉ¬π), " points")

    # Prepare grid for input data
    grid = make_grid(rd)
    println("Build grid for input data: ", length(grid), " points")

    # Prepare mesh for output spectrum
    mesh = make_mesh()
    println("Build mesh for spectrum: ", length(mesh), " points")

    # Initialize counters for Monte Carlo engine
    MC = init_mc(S)
    println("Create infrastructure for Monte Carlo sampling")

    # Initialize Monte Carlo configurations
    SE = init_element(S, MC.rng, allow)
    println("Randomize Monte Carlo configurations")

    # Prepare some key variables
    SC = init_context(SE, G·µ•, œÉ¬π, allow, grid, mesh, fmesh)
    println("Initialize context for the StochSK solver")

    return MC, SE, SC
end

"""
    run(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Perform stochastic analytic continuation simulation, sequential version.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
* Aout -> Spectral function, A(œâ).
* œá¬≤vec -> Œò-dependent œá¬≤, œá¬≤(Œò).
* Œòvec -> List of Œò parameters.
"""
function run(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    # By default, we should write the analytic continuation results
    # into the external files.
    _fwrite = get_b("fwrite")
    fwrite = isa(_fwrite, Missing) || _fwrite ? true : false

    # Setup essential parameters
    nstep = get_k("nstep")
    retry = get_k("retry")
    output_per_steps = get_k("ndump")
    measure_per_steps = 10

    # Warmup the Monte Carlo engine
    println("Start thermalization...")
    warmup(MC, SE, SC)

    # Shuffle the Monte Carlo configuration again
    shuffle(MC, SE, SC)

    # Sample and collect data
    step = 0.0
    println("Start stochastic sampling...")
    for iter = 1:nstep
        if iter % retry == 0
            SC.œá¬≤ = calc_goodness(SC.G·µß, SC.G·µ•)
        end

        sample(MC, SE, SC)

        if iter % measure_per_steps == 0
            step = step + 1.0
            measure(SE, SC)
        end

        if iter % output_per_steps == 0
            prog = round(I64, iter / nstep * 100)
            @printf("step = %6i ", iter)
            @printf("[progress = %3i]\n", prog)
            flush(stdout)
            fwrite && write_statistics(MC)
        end
    end

    return average(step, SC)
end

"""
    prun(
        S::StochSKSolver,
        p1::Dict{String,Vector{Any}},
        p2::Dict{String,Vector{Any}},
        MC::StochSKMC,
        SE::StochSKElement,
        SC::StochSKContext
    )

Perform stochastic analytic continuation simulation, parallel version.
The arguments `p1` and `p2` are copies of PBASE and PStochSK, respectively.

### Arguments
* S -> A StochSKSolver struct.
* p1 -> A copy of PBASE.
* p2 -> A copy of PStochSK.
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
* Aout -> Spectral function, A(œâ).
* œá¬≤vec -> Œò-dependent œá¬≤, œá¬≤(Œò).
* Œòvec -> List of Œò parameters.
"""
function prun(
    S::StochSKSolver,
    p1::Dict{String,Vector{Any}},
    p2::Dict{String,Vector{Any}},
    MC::StochSKMC,
    SE::StochSKElement,
    SC::StochSKContext
    )
    # Revise parameteric dicts
    rev_dict_b(p1)
    rev_dict_k(S, p2)

    # Initialize random number generator again
    MC.rng = MersenneTwister(rand(1:10000) * myid() + 1981)

    # By default, we should write the analytic continuation results
    # into the external files.
    _fwrite = get_b("fwrite")
    fwrite = isa(_fwrite, Missing) || _fwrite ? true : false

    # Setup essential parameters
    nstep = get_k("nstep")
    retry = get_k("retry")
    output_per_steps = get_k("ndump")
    measure_per_steps = 10

    # Warmup the Monte Carlo engine
    println("Start thermalization...")
    warmup(MC, SE, SC)

    # Shuffle the Monte Carlo configuration again
    shuffle(MC, SE, SC)

    # Sample and collect data
    step = 0.0
    println("Start stochastic sampling...")
    for iter = 1:nstep
        if iter % retry == 0
            SC.œá¬≤ = calc_goodness(SC.G·µß, SC.G·µ•)
        end

        sample(MC, SE, SC)

        if iter % measure_per_steps == 0
            step = step + 1.0
            measure(SE, SC)
        end

        if iter % output_per_steps == 0
            prog = round(I64, iter / nstep * 100)
            @printf("step = %6i ", iter)
            @printf("[progress = %3i]\n", prog)
            flush(stdout)
            myid() == 2 && fwrite && write_statistics(MC)
        end
    end

    return average(step, SC)
end

"""
    average(step::F64, SC::StochSKContext)

Postprocess the results generated during the stochastic analytic
continuation simulations. It will generate the spectral functions.

### Arguments
* step -> Number of Monte Carlo samplings.
* SC   -> A StochSKContext struct.

### Returns
* Aout -> Spectral function, A(œâ).
* œá¬≤vec -> Œò-dependent œá¬≤, œá¬≤(Œò).
* Œòvec -> List of Œò parameters.
"""
function average(step::F64, SC::StochSKContext)
    #
    # Here, the factor SC.mesh.weight in denominator is used to make sure
    # that the sum-rule
    #
    # ‚à´ A(œâ) dœâ = 1
    #
    # is obeyed by the obtained spectral function A(œâ).
    #
    SC.Aout = SC.Aout ./ (step * SC.mesh.weight)

    return SC.Aout, SC.œá¬≤vec, SC.Œòvec
end

"""
    last(
        SC::StochSKContext,
        Asum::Vector{F64},
        œá¬≤vec::Vector{F64},
        Œòvec::Vector{F64}
    )

It will process and write the calculated results by the StochSK solver,
including the final spectral function and reproduced correlator.

### Arguments
* SC    -> A StochSKContext struct.
* Asum  -> Spectral function, A(œâ).
* œá¬≤vec -> Œò-dependent œá¬≤.
* Œòvec  -> List of Œò parameters.

### Returns
* G -> Retarded Green's function, G(œâ).
"""
function last(
    SC::StochSKContext,
    Asum::Vector{F64},
    œá¬≤vec::Vector{F64},
    Œòvec::Vector{F64}
    )
    # By default, we should write the analytic continuation results
    # into the external files.
    _fwrite = get_b("fwrite")
    fwrite = isa(_fwrite, Missing) || _fwrite ? true : false

    # Write Œò-dependent goodness-of-fit function
    fwrite && write_goodness(Œòvec, œá¬≤vec)

    # Write final spectral function
    fwrite && write_spectrum(SC.mesh, Asum)

    # Reproduce input data and write them
    kernel = make_kernel(SC.mesh, SC.grid)
    G = reprod(SC.mesh, kernel, Asum)
    fwrite && write_backward(SC.grid, G)

    # Calculate full response function on real axis and write them
    if get_b("ktype") == "fermi"
        _G = kramers(SC.mesh, Asum)
    else
        _G = kramers(SC.mesh, Asum .* SC.mesh)
    end
    fwrite && write_complete(SC.mesh, _G)

    return _G
end

#=
### *Core Algorithms*
=#

"""
    warmup(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Warmup the Monte Carlo engine to acheieve thermalized equilibrium. Then
it will try to figure out the optimized Œò and the corresponding Monte
Carlo field configuration.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A
"""
function warmup(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    # Get essential parameters
    nwarm = get_k("nwarm")
    ratio = get_k("ratio")
    threshold = 1e-3

    # To store the historic Monte Carlo field configurations
    ùíû·µß = StochSKElement[]

    # Change the Œò parameter and approch the equilibrium state
    for i = 1:nwarm
        # Shuffle the Monte Carlo configurations
        shuffle(MC, SE, SC)

        # Backup key parameters and Monte Carlo field configurations
        SC.œá¬≤vec[i] = SC.œá¬≤
        SC.Œòvec[i] = SC.Œò
        push!(ùíû·µß, deepcopy(SE))

        # Check whether the equilibrium state is reached
        Œ¥œá¬≤ = SC.œá¬≤ - SC.œá¬≤min
        @printf("step : %5i ", i)
        @printf("œá¬≤ - œá¬≤min -> %12.6e\n", Œ¥œá¬≤)
        if Œ¥œá¬≤ < threshold
            println("Reach equilibrium state")
            break
        else
            if i == nwarm
                error("Fail to reach equilibrium state")
            end
        end

        # Adjust the Œò parameter
        SC.Œò = SC.Œò * ratio
    end

    # Well, we have vectors for Œò and œá¬≤. We have to figure out the
    # optimized Œò and œá¬≤, and then extract the corresponding Monte
    # Carlo field configuration.
    c = calc_theta(length(ùíû·µß), SC)
    @assert 1 ‚â§ c ‚â§ length(ùíû·µß)

    # Retrieve the Monte Carlo field configuration
    @. SE.P = ùíû·µß[c].P
    SE.A = ùíû·µß[c].A
    SE.W = ùíû·µß[c].W

    # Reset Œò
    SC.Œò = SC.Œòvec[c]

    # Update G·µß and œá¬≤
    SC.G·µß = calc_correlator(SE, SC.kernel)
    SC.œá¬≤ = calc_goodness(SC.G·µß, SC.G·µ•)
    println("Œò = ", SC.Œò, " œá¬≤ = ", SC.œá¬≤, " (step = $c)")
end

"""
    sample(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Perform Monte Carlo sweeps and sample the field configurations.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A
"""
function sample(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    if rand(MC.rng) < 0.80
        try_move_s(MC, SE, SC)
    else
        if rand(MC.rng) < 0.50
            try_move_p(MC, SE, SC)
        else
            try_move_q(MC, SE, SC)
        end
    end
end

"""
    measure(SE::StochSKElement, SC::StochSKContext)

Accumulate the final spectral functions A(œâ).

### Arguments
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A

See also: [`nearest`](@ref).
"""
function measure(SE::StochSKElement, SC::StochSKContext)
    nfine = get_k("nfine")
    ngamm = get_k("ngamm")

    for j = 1:ngamm
        d_pos = SE.P[j]
        # d_pos / nfine denotes the position of the selected Œ¥-like peak
        # in the fine linear mesh.
        #
        # The nearest() function is used to extract the approximated
        # position (index) of the selected Œ¥ function in the spectral
        # mesh, which could be linear or non-linear.
        #
        # Note that nearest() is defined in mesh.jl.
        s_pos = nearest(SC.mesh, d_pos / nfine)
        SC.Aout[s_pos] = SC.Aout[s_pos] + SE.A
    end
end

"""
    shuffle(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Try to shuffle the Monte Carlo field configuration via the Metropolis
algorithm. Then the window for shifting the Œ¥ functions is adjusted.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A
"""
function shuffle(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    # Get/set essential parameters
    nfine = get_k("nfine")
    retry = get_k("retry")
    max_bin_size = 100 # You can increase it to improve the accuracy

    # Announce counters
    bin_œá¬≤  = 0.0
    bin_acc = 0.0
    bin_try = 0.0

    # Perform Monte Carlo sweeping
    for s = 1:max_bin_size
        # Recalculate the goodness-of-fit function
        if s % retry == 0
            SC.œá¬≤ = calc_goodness(SC.G·µß, SC.G·µ•)
        end

        sample(MC, SE, SC)

        # Update the counters
        bin_œá¬≤ = bin_œá¬≤ + SC.œá¬≤
        bin_acc = bin_acc + (MC.Sacc + MC.Pacc)
        bin_try = bin_try + (MC.Stry + MC.Ptry)
    end

    # Calculate the transition probability, and then adjust the window,
    # which restricts the movement of the Œ¥ functions.
    #
    # The transition probability will be kept around 0.5.
    ùëù = bin_acc / bin_try
    #
    if ùëù > 0.5
        r = SE.W * 1.5
        if ceil(I64, r) < nfine
            SE.W = ceil(I64, r)
        else
            SE.W = nfine
        end
    end
    #
    if ùëù < 0.4
        SE.W = ceil(I64, SE.W / 1.5)
    end

    # Update œá¬≤ with averaged œá¬≤
    SC.œá¬≤ = bin_œá¬≤ / max_bin_size
end

#=
### *Service Functions*
=#

"""
    init_iodata(S::StochSKSolver, rd::RawData)

Preprocess the input data (`rd`).

### Arguments
* S -> A StochSKSolver struct.
* rd -> A RawData struct, which contains essential input data.

### Returns
* G·µ• -> Input correlator.
* œÉ¬π -> 1.0 / œÉ¬π.

See also: [`RawData`](@ref).
"""
function init_iodata(S::StochSKSolver, rd::RawData)
    G = make_data(rd)
    G·µ• = G.value # G·µ• = abs.(G.value)
    œÉ¬π = 1.0 ./ sqrt.(G.covar)

    return G·µ•, œÉ¬π
end

"""
    init_mc(S::StochSKSolver)

Try to create a StochSKMC struct. Some counters for Monte Carlo updates
are initialized here.

### Arguments
* S -> A StochSKSolver struct.

### Returns
* MC -> A StochSKMC struct.

See also: [`StochSKMC`](@ref).
"""
function init_mc(S::StochSKSolver)
    seed = rand(1:100000000)
    rng = MersenneTwister(seed)
    #
    Sacc = 0
    Stry = 0
    Pacc = 0
    Ptry = 0
    Qacc = 0
    Qtry = 0
    #
    MC = StochSKMC(rng, Sacc, Stry, Pacc, Ptry, Qacc, Qtry)

    return MC
end

"""
    init_element(
        S::StochSKSolver,
        rng::AbstractRNG,
        allow::Vector{I64}
    )

Randomize the configurations for future Monte Carlo sampling. It will
return a StochSKElement struct.

### Arguments
* S     -> A StochSKSolver struct.
* rng   -> Random number generator.
* allow -> Allowed positions for the Œ¥ peaks.

### Returns
* SE -> A StochSKElement struct.

See also: [`StochSKElement`](@ref).
"""
function init_element(
    S::StochSKSolver,
    rng::AbstractRNG,
    allow::Vector{I64}
    )
    Œ≤ = get_b("beta")
    wmax = get_b("wmax")
    wmin = get_b("wmin")
    nfine = get_k("nfine")
    ngamm = get_k("ngamm")

    position = rand(rng, allow, ngamm)
    #
    amplitude = 1.0 / ngamm
    #
    Œ¥f = (wmax - wmin) / (nfine - 1)
    average_freq = abs(log(2.0) / Œ≤)
    window_width = ceil(I64, 0.1 * average_freq / Œ¥f)

    return StochSKElement(position, amplitude, window_width)
end

"""
    init_context(
        SE::StochSKElement,
        G·µ•::Vector{F64},
        œÉ¬π::Vector{F64},
        allow::Vector{I64},
        grid::AbstractGrid,
        mesh::AbstractMesh,
        fmesh::AbstractMesh
    )

Try to create a StochSKContext struct, which contains some key variables,
including grid, mesh, input correlator and the corresponding standard
deviation, kernel matrix, spectral function, and goodness-of-fit function.

### Arguments
* SE -> A StochSKElement struct.
* G·µ• -> Input correlator. It will be changed in this function.
* œÉ¬π -> Standard deviation for input correlator.
* allow -> Allowable indices for Œ¥-like peaks.
* grid -> Imaginary axis grid for input data.
* mesh -> Real frequency mesh for output spectrum.
* fmesh -> Very fine mesh in [wmin, wmax].

### Returns
* SC -> A StochSKContext struct.
"""
function init_context(
    SE::StochSKElement,
    G·µ•::Vector{F64},
    œÉ¬π::Vector{F64},
    allow::Vector{I64},
    grid::AbstractGrid,
    mesh::AbstractMesh,
    fmesh::AbstractMesh
    )
    # Get parameters
    nmesh = get_b("nmesh")
    nwarm = get_k("nwarm")
    Œò = get_k("theta")

    # Allocate memory for spectral function, A(œâ)
    Aout = zeros(F64, nmesh)

    # Allocate memory for œá¬≤ and Œò
    œá¬≤vec = zeros(F64, nwarm)
    Œòvec = zeros(F64, nwarm)

    # Build kernel matrix
    kernel = make_kernel(fmesh, grid)

    # In order to accelerate the calculations, the singular space of the
    # kernel function is used. At first, we preform singular value
    # decomposition for K/œÉ:
    #     K/œÉ = U S V·µÄ
    # Then
    #     (G - KA)/œÉ = G/œÉ - K/œÉA
    #                = UU'(G/œÉ - USV·µÄA)
    #                = U(U'G/œÉ - U'USV·µÄA)
    #                = U(U'G/œÉ - SV·µÄA)
    #                = U(G' - K'A)
    # In the StochAC solver, let G·µ• ‚Üí G', kernel ‚Üí K'. Then new œá¬≤ is
    # calculated by
    #     |G' - K'A|¬≤
    # instead of
    #     |G - KA|¬≤/œÉ¬≤

    # Singular value decomposition of K/œÉ
    U, V, S = make_singular_space(Diagonal(œÉ¬π) * kernel)

    # Get new kernel matrix
    kernel = Diagonal(S) * V'

    # Get new (input) correlator
    G·µ• = U' *  (G·µ• .* œÉ¬π)

    # Calculate reconstructed correlator using current field configuration
    G·µß = calc_correlator(SE, kernel)

    # Calculate goodness-of-fit functional œá¬≤
    ùöæ = calc_goodness(G·µß, G·µ•)
    œá¬≤, œá¬≤min = ùöæ, ùöæ

    return StochSKContext(G·µ•, G·µß, œÉ¬π, allow, grid, mesh, kernel, Aout,
                        œá¬≤, œá¬≤min, œá¬≤vec, Œò, Œòvec)
end

"""
    calc_fmesh(S::StochSKSolver)

Try to calculate very fine (dense) linear mesh in [wmin, wmax], which
is used internally to build the kernel function. Note that the stochastic
analytic continuation method (A. W. Sandvik's version) does not support
the self-adaptive mesh.

### Arguments
* S -> A StochSKSolver struct.

### Returns
* fmesh -> A very fine, uniform mesh in [wmin, wmax].

See also: [`LinearMesh`](@ref), [`DynamicMesh`](@ref).
"""
function calc_fmesh(S::StochSKSolver)
    wmin = get_b("wmin")
    wmax = get_b("wmax")
    nfine = get_k("nfine")

    fmesh = LinearMesh(nfine, wmin, wmax)

    return fmesh
end

"""
    calc_correlator(SE::StochSKElement, kernel::Array{F64,2})

Try to calculate correlator with the kernel function and the Monte Carlo
field configuration. This correlator will then be used to evaluate the
goodness-of-fit function œá¬≤.

### Arguments
* SE     -> A StochSKElement struct.
* kernel -> The fermionic or bosonic kernel.

### Returns
* G -> Reconstructed correlator.

See also: [`calc_goodness`](@ref).
"""
function calc_correlator(SE::StochSKElement, kernel::Array{F64,2})
    ngamm = length(SE.P)
    ùê¥ = fill(SE.A, ngamm)
    ùêæ = kernel[:, SE.P]
    return ùêæ * ùê¥
end

"""
    calc_goodness(G‚Çô::Vector{F64}, G·µ•::Vector{F64})

Try to calculate the goodness-of-fit function (i.e, œá¬≤), which measures
the distance between input and regenerated correlators.

### Arguments
* G‚Çô -> Reconstructed correlators.
* G·µ• -> Input (original) correlators.

### Returns
* œá¬≤ -> Goodness-of-fit function.

See also: [`calc_correlator`](@ref).
"""
function calc_goodness(G‚Çô::Vector{F64}, G·µ•::Vector{F64})
    ŒîG = G‚Çô - G·µ•
    return dot(ŒîG, ŒîG)
end

"""
    calc_theta(len::I64, SC::StochSKContext)

Try to locate the optimal Œò and œá¬≤. This function implements the `chi2min`
and `chi2kink` algorithms. Note that the `chi2min` algorithm is preferred.

### Arguments
* len -> Length of vector Œò.
* SC -> A StochSKContext struct.

### Returns
* c -> Selected index for optimal Œò.
"""
function calc_theta(len::I64, SC::StochSKContext)
    function fitfun(x, p)
        return @. p[1] + p[2] / (1.0 + exp(-p[4] * (x - p[3])))
    end

    # Which algorithm is preferred ?
    method = get_k("method")

    # Get length of Œò and œá¬≤ vectors
    c = len

    # `chi2min` algorithm, proposed by Shao and Sandvik
    if method == "chi2min"
        while c ‚â• 1
            if SC.œá¬≤vec[c] > SC.œá¬≤min + 2.0 * sqrt(SC.œá¬≤min)
                break
            end
            c = c - 1
        end
    end

    # `chi2kink` algorithm, inspired by the `chi2kink` algorithm
    # used in MaxEnt solver
    if method == "chi2kink"
        guess = [0.0, 5.0, 2.0, 0.0]
        fit = curve_fit(fitfun, log10.(SC.Œòvec[1:c]), log10.(SC.œá¬≤vec[1:c]), guess)
        _, _, a, b = fit.param
        #
        fit_pos = 2.5
        Œò_opt = a - fit_pos / b
        c = argmin( abs.( log10.(SC.Œòvec[1:c]) .- Œò_opt ) )
    end

    return c
end

"""
    constraints(S::StochSKSolver, fmesh::AbstractMesh)

Try to implement the constrained stochastic analytic continuation
method. This function will return a collection. It contains all the
allowable indices. Be careful, `fmesh` should be a fine linear mesh.

### Arguments
* S     -> A StochSKSolver struct.
* fmesh -> Very dense mesh for the Œ¥ peaks.

### Returns
* allow -> Allowable indices.

See also: [`StochSKSolver`](@ref).
"""
function constraints(S::StochSKSolver, fmesh::AbstractMesh)
    exclude = get_b("exclude")
    nfine = get_k("nfine")
    @assert nfine == length(fmesh)

    allow = I64[]

    # Go through the fine mesh and check every mesh point.
    # Is is excluded?
    for i in eachindex(fmesh)
        is_excluded = false
        #
        if !isa(exclude, Missing)
            for j in eachindex(exclude)
                if exclude[j][1] ‚â§ fmesh[i] ‚â§ exclude[j][2]
                    is_excluded = true
                    break
                end
            end
        end
        #
        if !is_excluded
            push!(allow, i)
        end
    end

    return allow
end

"""
    try_move_s(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Try to update the Monte Carlo field configurations via the Metropolis
algorithm. In each update, only single Œ¥ function is shifted.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A

See also: [`try_move_p`](@ref).
"""
function try_move_s(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    # Get parameters
    nfine = get_k("nfine")
    ngamm = get_k("ngamm")

    # Reset counters
    MC.Sacc = 0
    MC.Stry = ngamm
    @assert 1 < SE.W ‚â§ nfine

    # Allocate memory for new correlator
    G‚Çô = zeros(F64, size(SC.G·µß))
    ŒîG = zeros(F64, size(SC.G·µß))

    for _ = 1:ngamm
        # Choose single Œ¥ function
        s = rand(MC.rng, 1:ngamm)

        # Evaluate new position for the Œ¥ function
        pcurr = SE.P[s]
        #
        if 1 < SE.W < nfine
            Œ¥W = rand(MC.rng, 1:SE.W)
            #
            if rand(MC.rng) > 0.5
                pnext = pcurr + Œ¥W
            else
                pnext = pcurr - Œ¥W
            end
            #
            pnext < 1     && (pnext = pnext + nfine)
            pnext > nfine && (pnext = pnext - nfine)
        else
            pnext = rand(MC.rng, 1:nfine)
        end

        # Apply the constraints
        !(pnext in SC.allow) && continue

        # Calculate the transition probability
        Knext = view(SC.kernel, :, pnext)
        Kcurr = view(SC.kernel, :, pcurr)
        #
        @. G‚Çô = SC.G·µß + SE.A * (Knext - Kcurr)
        @. ŒîG = G‚Çô - SC.G·µ•
        œá¬≤new = dot(ŒîG, ŒîG)
        #
        prob = exp( 0.5 * (SC.œá¬≤ - œá¬≤new) / SC.Œò )

        # Important sampling, if true, the Œ¥ function is shifted and the
        # corresponding objects are updated.
        if rand(MC.rng) < min(prob, 1.0)
            SE.P[s] = pnext
            @. SC.G·µß = G‚Çô
            #
            SC.œá¬≤ = œá¬≤new
            if œá¬≤new < SC.œá¬≤min
                SC.œá¬≤min = œá¬≤new
            end
            #
            MC.Sacc = MC.Sacc + 1
        end
    end
end

"""
    try_move_p(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Try to update the Monte Carlo field configurations via the Metropolis
algorithm. In each update, only a pair of Œ¥ functions are shifted.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A

See also: [`try_move_s`](@ref).
"""
function try_move_p(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    # Get parameters
    nfine = get_k("nfine")
    ngamm = get_k("ngamm")

    # We have to make sure that there are at least two Œ¥ functions here.
    ngamm < 2 && return

    # Reset counters
    MC.Pacc = 0
    MC.Ptry = ngamm
    @assert 1 < SE.W ‚â§ nfine

    # Allocate memory for new correlator
    G‚Çô = zeros(F64, size(SC.G·µß))
    ŒîG = zeros(F64, size(SC.G·µß))

    for _ = 1:ngamm
        # Choose a pair of Œ¥ functions
        s‚ÇÅ = rand(MC.rng, 1:ngamm)
        s‚ÇÇ = s‚ÇÅ
        while s‚ÇÅ == s‚ÇÇ
            s‚ÇÇ = rand(MC.rng, 1:ngamm)
        end

        # Evaluate new positions for the two Œ¥ functions
        pcurr‚ÇÅ = SE.P[s‚ÇÅ]
        pcurr‚ÇÇ = SE.P[s‚ÇÇ]
        #
        if 1 < SE.W < nfine
            Œ¥W‚ÇÅ = rand(MC.rng, 1:SE.W)
            Œ¥W‚ÇÇ = rand(MC.rng, 1:SE.W)
            #
            if rand(MC.rng) > 0.5
                pnext‚ÇÅ = pcurr‚ÇÅ + Œ¥W‚ÇÅ
                pnext‚ÇÇ = pcurr‚ÇÇ - Œ¥W‚ÇÇ
            else
                pnext‚ÇÅ = pcurr‚ÇÅ - Œ¥W‚ÇÅ
                pnext‚ÇÇ = pcurr‚ÇÇ + Œ¥W‚ÇÇ
            end
            #
            pnext‚ÇÅ < 1     && (pnext‚ÇÅ = pnext‚ÇÅ + nfine)
            pnext‚ÇÅ > nfine && (pnext‚ÇÅ = pnext‚ÇÅ - nfine)
            pnext‚ÇÇ < 1     && (pnext‚ÇÇ = pnext‚ÇÇ + nfine)
            pnext‚ÇÇ > nfine && (pnext‚ÇÇ = pnext‚ÇÇ - nfine)
        else
            pnext‚ÇÅ = rand(MC.rng, 1:nfine)
            pnext‚ÇÇ = rand(MC.rng, 1:nfine)
        end

        # Apply the constraints
        !(pnext‚ÇÅ in SC.allow) && continue
        !(pnext‚ÇÇ in SC.allow) && continue

        # Calculate the transition probability
        Knext‚ÇÅ = view(SC.kernel, :, pnext‚ÇÅ)
        Kcurr‚ÇÅ = view(SC.kernel, :, pcurr‚ÇÅ)
        Knext‚ÇÇ = view(SC.kernel, :, pnext‚ÇÇ)
        Kcurr‚ÇÇ = view(SC.kernel, :, pcurr‚ÇÇ)
        #
        @. G‚Çô = SC.G·µß + SE.A * (Knext‚ÇÅ - Kcurr‚ÇÅ + Knext‚ÇÇ - Kcurr‚ÇÇ)
        @. ŒîG = G‚Çô - SC.G·µ•
        œá¬≤new = dot(ŒîG, ŒîG)
        #
        prob = exp( 0.5 * (SC.œá¬≤ - œá¬≤new) / SC.Œò )

        # Important sampling, if true, the Œ¥ functions are shifted and the
        # corresponding objects are updated.
        if rand(MC.rng) < min(prob, 1.0)
            SE.P[s‚ÇÅ] = pnext‚ÇÅ
            SE.P[s‚ÇÇ] = pnext‚ÇÇ
            @. SC.G·µß = G‚Çô
            #
            SC.œá¬≤ = œá¬≤new
            if œá¬≤new < SC.œá¬≤min
                SC.œá¬≤min = œá¬≤new
            end
            #
            MC.Pacc = MC.Pacc + 1
        end
    end
end

"""
    try_move_q(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)

Try to update the Monte Carlo field configurations via the Metropolis
algorithm. In each update, four different Œ¥ functions are shifted.

### Arguments
* MC -> A StochSKMC struct.
* SE -> A StochSKElement struct.
* SC -> A StochSKContext struct.

### Returns
N/A

See also: [`try_move_s`](@ref).
"""
function try_move_q(MC::StochSKMC, SE::StochSKElement, SC::StochSKContext)
    # Get parameters
    nfine = get_k("nfine")
    ngamm = get_k("ngamm")

    # We have to make sure that there are at least four Œ¥ functions here.
    ngamm < 4 && return

    # Reset counters
    MC.Qacc = 0
    MC.Qtry = ngamm
    @assert 1 < SE.W ‚â§ nfine

    # Allocate memory for new correlator
    G‚Çô = zeros(F64, size(SC.G·µß))
    ŒîG = zeros(F64, size(SC.G·µß))

    for _ = 1:ngamm
        # Choose four different Œ¥ functions
        ùëÜ = nothing
        while true
            ùëÜ = rand(MC.rng, 1:ngamm, 4)
            ùíÆ = unique(ùëÜ)
            if length(ùëÜ) == length(ùíÆ)
                break
            end
        end
        s‚ÇÅ, s‚ÇÇ, s‚ÇÉ, s‚ÇÑ = ùëÜ

        # Evaluate new positions for the four Œ¥ functions
        pcurr‚ÇÅ = SE.P[s‚ÇÅ]
        pcurr‚ÇÇ = SE.P[s‚ÇÇ]
        pcurr‚ÇÉ = SE.P[s‚ÇÉ]
        pcurr‚ÇÑ = SE.P[s‚ÇÑ]
        #
        if 1 < SE.W < nfine
            Œ¥W‚ÇÅ = rand(MC.rng, 1:SE.W)
            Œ¥W‚ÇÇ = rand(MC.rng, 1:SE.W)
            Œ¥W‚ÇÉ = rand(MC.rng, 1:SE.W)
            Œ¥W‚ÇÑ = rand(MC.rng, 1:SE.W)
            #
            if rand(MC.rng) > 0.5
                pnext‚ÇÅ = pcurr‚ÇÅ + Œ¥W‚ÇÅ
                pnext‚ÇÇ = pcurr‚ÇÇ - Œ¥W‚ÇÇ
                pnext‚ÇÉ = pcurr‚ÇÉ + Œ¥W‚ÇÉ
                pnext‚ÇÑ = pcurr‚ÇÑ - Œ¥W‚ÇÑ
            else
                pnext‚ÇÅ = pcurr‚ÇÅ - Œ¥W‚ÇÅ
                pnext‚ÇÇ = pcurr‚ÇÇ + Œ¥W‚ÇÇ
                pnext‚ÇÉ = pcurr‚ÇÉ - Œ¥W‚ÇÉ
                pnext‚ÇÑ = pcurr‚ÇÑ + Œ¥W‚ÇÑ
            end
            #
            pnext‚ÇÅ < 1     && (pnext‚ÇÅ = pnext‚ÇÅ + nfine)
            pnext‚ÇÅ > nfine && (pnext‚ÇÅ = pnext‚ÇÅ - nfine)
            pnext‚ÇÇ < 1     && (pnext‚ÇÇ = pnext‚ÇÇ + nfine)
            pnext‚ÇÇ > nfine && (pnext‚ÇÇ = pnext‚ÇÇ - nfine)
            pnext‚ÇÉ < 1     && (pnext‚ÇÉ = pnext‚ÇÉ + nfine)
            pnext‚ÇÉ > nfine && (pnext‚ÇÉ = pnext‚ÇÉ - nfine)
            pnext‚ÇÑ < 1     && (pnext‚ÇÑ = pnext‚ÇÑ + nfine)
            pnext‚ÇÑ > nfine && (pnext‚ÇÑ = pnext‚ÇÑ - nfine)
        else
            pnext‚ÇÅ = rand(MC.rng, 1:nfine)
            pnext‚ÇÇ = rand(MC.rng, 1:nfine)
            pnext‚ÇÉ = rand(MC.rng, 1:nfine)
            pnext‚ÇÑ = rand(MC.rng, 1:nfine)
        end

        # Apply the constraints
        !(pnext‚ÇÅ in SC.allow) && continue
        !(pnext‚ÇÇ in SC.allow) && continue
        !(pnext‚ÇÉ in SC.allow) && continue
        !(pnext‚ÇÑ in SC.allow) && continue

        # Calculate the transition probability
        Knext‚ÇÅ = view(SC.kernel, :, pnext‚ÇÅ)
        Kcurr‚ÇÅ = view(SC.kernel, :, pcurr‚ÇÅ)
        Knext‚ÇÇ = view(SC.kernel, :, pnext‚ÇÇ)
        Kcurr‚ÇÇ = view(SC.kernel, :, pcurr‚ÇÇ)
        Knext‚ÇÉ = view(SC.kernel, :, pnext‚ÇÉ)
        Kcurr‚ÇÉ = view(SC.kernel, :, pcurr‚ÇÉ)
        Knext‚ÇÑ = view(SC.kernel, :, pnext‚ÇÑ)
        Kcurr‚ÇÑ = view(SC.kernel, :, pcurr‚ÇÑ)
        #
        @. G‚Çô = SC.G·µß + SE.A * ( Knext‚ÇÅ - Kcurr‚ÇÅ +
                                 Knext‚ÇÇ - Kcurr‚ÇÇ +
                                 Knext‚ÇÉ - Kcurr‚ÇÉ +
                                 Knext‚ÇÑ - Kcurr‚ÇÑ )
        @. ŒîG = G‚Çô - SC.G·µ•
        œá¬≤new = dot(ŒîG, ŒîG)
        #
        prob = exp( 0.5 * (SC.œá¬≤ - œá¬≤new) / SC.Œò )

        # Important sampling, if true, the Œ¥ functions are shifted and the
        # corresponding objects are updated.
        if rand(MC.rng) < min(prob, 1.0)
            SE.P[s‚ÇÅ] = pnext‚ÇÅ
            SE.P[s‚ÇÇ] = pnext‚ÇÇ
            SE.P[s‚ÇÉ] = pnext‚ÇÉ
            SE.P[s‚ÇÑ] = pnext‚ÇÑ
            @. SC.G·µß = G‚Çô
            #
            SC.œá¬≤ = œá¬≤new
            if œá¬≤new < SC.œá¬≤min
                SC.œá¬≤min = œá¬≤new
            end
            #
            MC.Qacc = MC.Qacc + 1
        end
    end
end
