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

#=
### *Input Functions*
=#

"""
    read_real_data(finput::AbstractString, ngrid::I64)

Read input data. This function is used for imaginary time data. The input
file should contain three columns. The first column is the imaginary time
grid, the second column is the value, the third column is the standard
deviation Ïƒ. Here, `ngrid` specifies the number of grid points.

### Arguments
* finput -> Filename for the input data.
* ngrid  -> Number of grid points.

### Returns
* rd -> A RawData struct.

See also: [`read_cmplx_data`](@ref).
"""
function read_real_data(finput::AbstractString, ngrid::I64)
    # Allocate memory
    _grid = zeros(F64, ngrid)
    value = zeros(F64, ngrid)
    error = zeros(F64, ngrid)

    # We have to determine the number of columns and rows at first.
    dlm = readdlm(finput)
    nrows, ncols = size(dlm)
    @assert nrows â‰¥ ngrid
    @assert ncols == 3

    # Read and parse the data
    open(finput, "r") do fin
        for i = 1:ngrid
            arr = parse.(F64, line_to_array(fin)[1:3])
            _grid[i] = arr[1]
            value[i] = arr[2]
            error[i] = arr[3]
        end
    end

    return RawData(_grid, value, error)
end

"""
    read_cmplx_data(finput::AbstractString, ngrid::I64)

Read input data. This function is used for Matsubara frequency data. The
input should contain four columns or five columns. The first column is
the Matsubara freqency grid, the second and third columns are the values
(real part and imaginary part), the four and fifth columns are the standard
deviations Ïƒ for the real and imaginary parts, respectively. If there are
only four columns, it means that the real and imaginary parts share the
same standard deviations.

### Arguments
* finput -> Filename for the input data.
* ngrid  -> Number of grid points.

### Returns
* rd -> A RawData struct.

See also: [`read_real_data`](@ref).
"""
function read_cmplx_data(finput::AbstractString, ngrid::I64)
    # Allocate memory
    _grid = zeros(F64, ngrid)
    value = zeros(C64, ngrid)
    error = zeros(C64, ngrid)

    # We have to determine the number of columns and rows at first.
    dlm = readdlm(finput)
    nrows, ncols = size(dlm)
    @assert nrows â‰¥ ngrid
    @assert ncols in (4, 5)

    # Read and parse the data
    open(finput, "r") do fin
        for i = 1:ngrid
            if ncols == 4
                arr = parse.(F64, line_to_array(fin)[1:4])
                _grid[i] = arr[1]
                value[i] = arr[2] + im * arr[3]
                error[i] = arr[4] + im * arr[4]
            else
                arr = parse.(F64, line_to_array(fin)[1:5])
                _grid[i] = arr[1]
                value[i] = arr[2] + im * arr[3]
                error[i] = arr[4] + im * arr[5]
            end
        end
    end

    return RawData(_grid, value, error)
end

"""
    read_cmplx_data(
        finput::AbstractString,
        ngrid::I64,
        only_real_part::Bool
    )

Read input data. This function is used for Matsubara frequency data. The
input file only contains three columns. The first column is the Matsubara
frequency grid, the second column is the real part or imaginary part of
the data (which is specified by the argument `only_real_part`), and the
third column is the standard deviation Ïƒ. This function is for bosonic
correlation function.

### Arguments
* finput -> Filename for the input data.
* ngrid  -> Number of grid points.
* only_real_part -> See above explanations.

### Returns
* rd -> A RawData struct.

See also: [`read_real_data`](@ref).
"""
function read_cmplx_data(
    finput::AbstractString,
    ngrid::I64,
    only_real_part::Bool
    )
    # Allocate memory
    _grid = zeros(F64, ngrid)
    value = zeros(C64, ngrid)
    error = zeros(C64, ngrid)

    # We have to determine the number of columns and rows at first.
    dlm = readdlm(finput)
    nrows, ncols = size(dlm)
    @assert nrows â‰¥ ngrid
    @assert ncols == 3

    # Read and parse the data
    open(finput, "r") do fin
        for i = 1:ngrid
            arr = parse.(F64, line_to_array(fin)[1:3])
            _grid[i] = arr[1]
            if only_real_part
                value[i] = arr[2]
                error[i] = arr[3]
            else
                value[i] = im * arr[2]
                error[i] = im * arr[3]
            end
        end
    end

    return RawData(_grid, value, error)
end

#=
### *Output Functions*
=#

"""
    write_spectrum(am::AbstractMesh, Aout::Vector{F64})

Write spectrum A(Ï‰) or A(Ï‰) / Ï‰ to `Aout.data`. The grid is defined in
`am`, and the spectral data are contained in `Aout`.

Note that for the MaxEnt, StochAC, StochSK, and StochOM solvers, `Aout`
is actually A(Ï‰) / Ï‰, instead of A(Ï‰). However, for the BarRat, NevanAC,
and StochPX solvers, `Aout` is just A(Ï‰).

### Arguments
* am   -> Real frequency mesh.
* Aout -> Spectral function. See above explanations.

### Returns
N/A
"""
function write_spectrum(am::AbstractMesh, Aout::Vector{F64})
    @assert length(am) == length(Aout)

    open("Aout.data", "w") do fout
        for i in eachindex(am)
            @printf(fout, "%16.12f %16.12f\n", am[i], Aout[i])
        end
    end
end

"""
    write_spectrum(am::AbstractMesh, Î±â‚—::Vector{F64}, Aout::Array{F64,2})

Write Î±-resolved spectrum A(Ï‰) to `Aout.data.alpha`. The grid is defined
in `am`, the Î±-resolved spectrum is contained in `Aout`, `Î±â‚—` is the list
for the Î± parameters. This function is called by the `StochAC` solver.

Note that `Aout` is actually A(Ï‰) / Ï‰, instead of A(Ï‰).

### Arguments
* am   -> Real frequency mesh.
* Î±â‚—   -> List for Î± parameters.
* Aout -> Î±-dependent spectral function.

### Returns
N/A
"""
function write_spectrum(am::AbstractMesh, Î±â‚—::Vector{F64}, Aout::Array{F64,2})
    nmesh, nalph = size(Aout)
    @assert nmesh == length(am)
    @assert nalph == length(Î±â‚—)

    for i in eachindex(Î±â‚—)
        open("Aout.data.alpha_$i", "w") do fout
            println(fout, "# $i : Î± = ", Î±â‚—[i])
            for j in eachindex(am)
                @printf(fout, "%16.12f %16.12f\n", am[j], Aout[j,i])
            end
        end
    end
end

"""
    write_backward(ag::AbstractGrid, G::Vector{F64})

We can use the calculated spectrum in real axis to reproduce the input
data in imaginary axis. This function will write the reproduced data to
`repr.data`, which can be compared with the original data. Here, `G` is
the reproduced data.

### Arguments
* ag -> Grid for input data, Ï„ or iÏ‰â‚™.
* G  -> Reconstructed Green's function, G(Ï„) or G(iÏ‰â‚™).

### Returns
N/A

See also: [`reprod`](@ref).
"""
function write_backward(ag::AbstractGrid, G::Vector{F64})
    ngrid = length(ag)
    ng = length(G)
    @assert ngrid == ng || ngrid * 2 == ng

    # The reproduced data are defined in imaginary time axis.
    if ngrid == ng
        open("repr.data", "w") do fout
            for i in eachindex(ag)
                @printf(fout, "%16.12f %16.12f\n", ag[i], G[i])
            end
        end
    # The reproduced data are defined in Matsubara frequency axis.
    else
        open("repr.data", "w") do fout
            for i in eachindex(ag)
                @printf(fout, "%16.12f %16.12f %16.12f\n", ag[i], G[i], G[i+ngrid])
            end
        end
    end
end

"""
    write_complete(am::AbstractMesh, G::Vector{C64})

Write the full data at real axis to `Gout.data`. `am` denotes the real
axis, `G` is the calculated Green's function data. Note that its real
part is obtained via the so-called Kramers-Kronig transformation.

### Arguments
* am -> Real frequency mesh, Ï‰.
* G  -> Retarded Green's function, G(Ï‰).

### Returns
N/A

See also: [`kramers`](@ref).
"""
function write_complete(am::AbstractMesh, G::Vector{C64})
    @assert length(am) == length(G)

    open("Gout.data", "w") do fout
        for i in eachindex(am)
            z = G[i]
            @printf(fout, "%16.12f %16.12f %16.12f\n", am[i], real(z), imag(z))
        end
    end
end

"""
    write_misfit(Î±_vec::Vector{F64}, Ï‡Â²_vec::Vector{F64})

Write `log10(Î±)-log10(Ï‡Â²)` data to `chi2.data`, which could be used to
judge whether the obtained optimal Î± parameter is reasonable. It is used
by the `MaxEnt` solver only.

### Arguments
* Î±_vec  -> List for Î± parameters.
* Ï‡Â²_vec -> Î±-dependent goodness-of-fit functional.

### Returns
N/A

See also: [`write_goodness`](@ref).
"""
function write_misfit(Î±_vec::Vector{F64}, Ï‡Â²_vec::Vector{F64})
    @assert length(Î±_vec) == length(Ï‡Â²_vec)

    open("chi2.data", "w") do fout
        _Î± = log10.(Î±_vec)
        _Ï‡Â² = log10.(Ï‡Â²_vec)
        for i in eachindex(Î±_vec)
            @printf(fout, "%16.12f %16.12f\n", _Î±[i], _Ï‡Â²[i])
        end
    end
end

"""
    write_goodness(Î˜_vec::Vector{F64}, Ï‡Â²_vec::Vector{F64})

Write `log10(Î˜)-log10(Ï‡Â²)` data to `goodness.data`, which could be used
to judge whether the obtained optimal Î˜ parameter is reasonable. This
function is only useful for the `StochSK` solver.

### Arguments
* Î˜_vec  -> List for Î˜ parameters.
* Ï‡Â²_vec -> Î˜-dependent goodness-of-fit functional.

### Returns
N/A

See also: [`write_misfit`](@ref).
"""
function write_goodness(Î˜_vec::Vector{F64}, Ï‡Â²_vec::Vector{F64})
    @assert length(Î˜_vec) == length(Ï‡Â²_vec)

    open("goodness.data", "w") do fout
        _Î˜ = log10.(Î˜_vec)
        _Ï‡Â² = log10.(Ï‡Â²_vec)
        for i in eachindex(Î˜_vec)
            if !isinf(_Î˜[i]) && !isinf(_Ï‡Â²[i])
                @printf(fout, "%16.12f %16.12f\n", _Î˜[i], _Ï‡Â²[i])
            end
        end
    end
end

"""
    write_model(am::AbstractMesh, D::Vector{F64})

Write the default model function to `model.data`. This function is usually
for the `MaxEnt` solver.

### Arguments
* am -> Real frequency mesh, Ï‰.
* D  -> Default model, m(Ï‰).

### Returns
N/A
"""
function write_model(am::AbstractMesh, D::Vector{F64})
    @assert length(am) == length(D)

    open("model.data", "w") do fout
        for i in eachindex(am)
            @printf(fout, "%16.12f %16.12f\n", am[i], D[i])
        end
    end
end

"""
    write_prony(ð�‘�â‚š::I64, Î“â‚š::Vector{C64}, Î©â‚š::Vector{C64})

Write Prony approximation to the input correlator. This information can
be used to reconstruct or interpolate the correlator. This function is
only useful for the `BarRat` solver.

### Arguments
* ð�‘�â‚š -> Number of nodes for Prony approximation.
* Î“â‚š -> Nodes for Prony approximation, ``Î³_i``.
* Î©â‚š -> Weights for Prony approximation, ``w_i``.

### Returns
N/A
"""
function write_prony(ð�‘�â‚š::I64, Î“â‚š::Vector{C64}, Î©â‚š::Vector{C64})
    open("prony.data", "w") do fout
        println(fout, "# Prony Approximation")
        #
        println(fout, "# ð�‘�â‚š :")
        @printf(fout, "%4i\n", ð�‘�â‚š)
        #
        println(fout, "# Î“â‚š :")
        for i in eachindex(Î“â‚š)
            z = Î“â‚š[i]
            @printf(fout, "%4i %16.12f %16.12f\n", i, real(z), imag(z))
        end
        #
        println(fout, "# Î©â‚š :")
        for i in eachindex(Î©â‚š)
            z = Î©â‚š[i]
            @printf(fout, "%4i %16.12f %16.12f\n", i, real(z), imag(z))
        end
    end
end

"""
    write_prony(ag::AbstractGrid, G::Vector{C64})

Write Matsubara Green's function approximated by the Prony approximation.
This function is only useful for the `BarRat` solver.

### Arguments
* ag -> Grid for input data, iÏ‰â‚™.
* G  -> Approximated Green's function, G(iÏ‰â‚™).

### Returns
N/A
"""
function write_prony(ag::AbstractGrid, G::Vector{C64})
    ngrid = length(ag)
    ng = length(G)
    @assert ngrid == ng

    open("Gprony.data", "w") do fout
        for i in eachindex(ag)
            z = G[i]
            @printf(fout, "%16.12f %16.12f %16.12f\n", ag[i], real(z), imag(z))
        end
    end
end

"""
    write_barycentric(
        nodes::Vector{C64},
        values::Vector{C64},
        weights::Vector{C64}
    )

Write barycentric rational function approximation to the input correlator.
This information can be used to reconstruct or interpolate the correlator.
This function is only useful for the `BarRat` solver.

### Arguments
* nodes   -> Nodes of the rational function, ``z_i``.
* values  -> Values of the rational function, ``r(z_i)``.
* weights -> Weights of the rational function, ``w_i``.

### Returns
N/A
"""
function write_barycentric(
    nodes   :: Vector{C64},
    values  :: Vector{C64},
    weights :: Vector{C64}
    )
    w_times_f = values .* weights

    open("barycentric.data", "w") do fout
        println(fout, "# Barycentric Rational Function Approximation")
        #
        println(fout, "# nodes :")
        for i in eachindex(nodes)
            z = nodes[i]
            @printf(fout, "%4i %16.12f %16.12f\n", i, real(z), imag(z))
        end
        #
        println(fout, "# values :")
        for i in eachindex(values)
            z = values[i]
            @printf(fout, "%4i %16.12f %16.12f\n", i, real(z), imag(z))
        end
        #
        println(fout, "# weights :")
        for i in eachindex(weights)
            z = weights[i]
            @printf(fout, "%4i %16.12f %16.12f\n", i, real(z), imag(z))
        end
        #
        println(fout, "# w_times_f :")
        for i in eachindex(w_times_f)
            z = w_times_f[i]
            @printf(fout, "%4i %16.12f %16.12f\n", i, real(z), imag(z))
        end
    end
end

"""
    write_hamiltonian(Î±_vec::Vector{F64}, UÎ±::Vector{F64})

Write `Î±-U(Î±)` data to `hamil.data`, which could be used to judge whether
the obtained optimal Î± parameter is reasonable. This function is only
useful for the `StochAC` solver.

### Arguments
* Î±_vec -> List for Î± parameters.
* UÎ±    -> Î±-dependent Hamiltonian.

### Returns
N/A
"""
function write_hamiltonian(Î±_vec::Vector{F64}, UÎ±::Vector{F64})
    @assert length(Î±_vec) == length(UÎ±)

    open("hamil.data", "w") do fout
        for i in eachindex(Î±_vec)
            @printf(fout, "%16.8f %16.12f\n", Î±_vec[i], UÎ±[i])
        end
    end
end

"""
    write_passed(passed::Vector{I64}, med::F64, Î±good::F64)

Write indices of selected solutions which should be used to calculate
the averaged spectrum. Here, `passed` means the indices, `med` is the
median value of Ï‡Â², and `Î±good` is the factor that is used to filter
the solutions. This function is only useful for the `StochOM` and the
`StochPX` solvers.

### Arguments
* passed -> Indices for selected solutions.
* med    -> Median value of Ï‡Â².
* Î±good  -> Predefined parameter used to filter the solutions (spectra).

### Returns
N/A
"""
function write_passed(passed::Vector{I64}, med::F64, Î±good::F64)
    open("passed.data", "w") do fout
        npass = length(passed)
        println(fout, "# Count: ", npass, "  Median: ", med, "  Î±good: ", Î±good)
        for i in eachindex(passed)
            @printf(fout, "%4i %8i\n", i, passed[i])
        end
    end
end

"""
    write_pole(
        Páµ¥::Vector{Vector{I64}},
        Aáµ¥::Vector{Vector{F64}},
        ð�•Šáµ¥::Vector{Vector{F64}},
        Ï‡Â²áµ¥::Vector{F64},
        fmesh::AbstractMesh
    )

Write positions, amplitudes, and signs of poles to `pole.data`. This
function is only useful for the `StochPX` solver.

Note that the `util/ppole.jl` script can read poles' information from
the `pole.data` file, and construct new spectral functions with different
`eta` parameters.

### Arguments
* Páµ¥    -> Positions of the poles.
* Aáµ¥    -> Amplitudes of the poles.
* ð�•Šáµ¥    -> Signs of the poles.
* Ï‡Â²áµ¥   -> Goodness-of-fit functionals for all the solutions.
* fmesh -> A dense mesh for the poles.

### Returns
N/A
"""
function write_pole(
    Páµ¥::Vector{Vector{I64}},
    Aáµ¥::Vector{Vector{F64}},
    ð�•Šáµ¥::Vector{Vector{F64}},
    Ï‡Â²áµ¥::Vector{F64},
    fmesh::AbstractMesh
    )
    ntry = length(Páµ¥)

    open("pole.data", "w") do fout
        for i = 1:ntry
            println(fout, "# Try: ", i, "  Ï‡Â²: ", Ï‡Â²áµ¥[i])
            #
            for j in eachindex(Páµ¥[i])
                â„™ = Páµ¥[i][j]
                ð�”¸ = Aáµ¥[i][j]
                ð�•Š = ð�•Šáµ¥[i][j]
                Ï‰ = fmesh[â„™]
                @printf(fout, "%4i %8i %16.12f %16.12f %6.2f\n", j, â„™, Ï‰, ð�”¸, ð�•Š)
            end
            #
            println(fout)
            println(fout)
        end
    end
end

"""
    write_probability(Î±_vec::Vector{F64}, p_vec::Vector{F64})

Write `p(Î±)` data to `prob.data`. This function is only useful for the
`MaxEnt` solver (`bryan` algorithm).

### Arguments
* Î±_vec -> List for Î± parameters.
* p_vec -> Î±-dependent probabilities.

### Returns
N/A
"""
function write_probability(Î±_vec::Vector{F64}, p_vec::Vector{F64})
    @assert length(Î±_vec) == length(p_vec)

    open("prob.data", "w") do fout
        _p = -p_vec ./ trapz(Î±_vec, p_vec)
        _Î± = log10.(Î±_vec)
        for i in eachindex(Î±_vec)
            @printf(fout, "%16.12f %16.12f\n", _Î±[i], _p[i])
        end
    end
end

"""
    write_statistics(MC::StochACMC)

Write Monte Carlo statistical information for the `StochAC` solver. Note
that the `StochAC` solver is based on a stochastic approach.

### Arguments
* MC -> A StochACMC struct.

### Returns
N/A

See also: [`PStochAC`](@ref), [`StochACMC`](@ref).
"""
function write_statistics(MC::StochACMC)
    nalph = get_a("nalph")

    open("stat.data", "w") do fout
        println(fout, "# Move statistics:")
        for i = 1:nalph
            @printf(fout, "Î± %3i -> %16.12f\n", i, MC.Macc[i] / MC.Mtry[i])
        end
        println(fout)
        println(fout, "# Swap statistics:")
        for i = 1:nalph
            @printf(fout, "Î± %3i -> %16.12f\n", i, MC.Sacc[i] / MC.Stry[i])
        end
    end
end

"""
    write_statistics(MC::StochSKMC)

Write Monte Carlo statistical information for the `StochSK` solver. Note
that the `StochSK` solver is based on a stochastic approach.

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

### Returns
N/A

See also: [`PStochSK`](@ref), [`StochSKMC`](@ref).
"""
function write_statistics(MC::StochSKMC)
    open("stat.data", "w") do fout
        println(fout, "# Move S statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Sacc)
        @printf(fout, "try    -> %16i    \n", MC.Stry)
        @printf(fout, "prob   -> %16.12f \n", MC.Sacc / MC.Stry)
        println(fout)
        println(fout, "# Move P statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Pacc)
        @printf(fout, "try    -> %16i    \n", MC.Ptry)
        @printf(fout, "prob   -> %16.12f \n", MC.Pacc / MC.Ptry)
        println(fout)
        println(fout, "# Move Q statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Qacc)
        @printf(fout, "try    -> %16i    \n", MC.Qtry)
        @printf(fout, "prob   -> %16.12f \n", MC.Qacc / MC.Qtry)
    end
end

"""
    write_statistics(MC::StochOMMC)

Write Monte Carlo statistical information for the `StochOM` solver. Note
that the `StochOM` solver is based on a stochastic approach.

### Arguments
* MC -> A StochOMMC struct.

### Returns
N/A

See also: [`PStochOM`](@ref), [`StochOMMC`](@ref).
"""
function write_statistics(MC::StochOMMC)
    open("stat.data", "w") do fout
        println(fout, "# Update statistics:")
        @printf(fout, "insert box -> %16.12f\n", MC.Macc[1] / MC.Mtry[1])
        @printf(fout, "remove box -> %16.12f\n", MC.Macc[2] / MC.Mtry[2])
        @printf(fout, "shift  box -> %16.12f\n", MC.Macc[3] / MC.Mtry[3])
        @printf(fout, "width  box -> %16.12f\n", MC.Macc[4] / MC.Mtry[4])
        @printf(fout, "height box -> %16.12f\n", MC.Macc[5] / MC.Mtry[5])
        @printf(fout, "split  box -> %16.12f\n", MC.Macc[6] / MC.Mtry[6])
        @printf(fout, "merge  box -> %16.12f\n", MC.Macc[7] / MC.Mtry[7])
    end
end

"""
    write_statistics(MC::StochPXMC)

Write Monte Carlo statistical information for the `StochPX` solver. Note
that the `StochPX` solver is based on a stochastic approach.

### Arguments
* MC -> A StochPXMC struct.

### Returns
N/A

See also: [`PStochPX`](@ref), [`StochPXMC`](@ref).
"""
function write_statistics(MC::StochPXMC)
    open("stat.data", "w") do fout
        println(fout, "# Move S statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Sacc)
        @printf(fout, "try    -> %16i    \n", MC.Stry)
        @printf(fout, "prob   -> %16.12f \n", MC.Sacc / MC.Stry)
        println(fout)
        println(fout, "# Move P statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Pacc)
        @printf(fout, "try    -> %16i    \n", MC.Ptry)
        @printf(fout, "prob   -> %16.12f \n", MC.Pacc / MC.Ptry)
        println(fout)
        println(fout, "# Move A statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Aacc)
        @printf(fout, "try    -> %16i    \n", MC.Atry)
        @printf(fout, "prob   -> %16.12f \n", MC.Aacc / MC.Atry)
        println(fout)
        println(fout, "# Move X statistics:")
        @printf(fout, "accept -> %16i    \n", MC.Xacc)
        @printf(fout, "try    -> %16i    \n", MC.Xtry)
        @printf(fout, "prob   -> %16.12f \n", MC.Xacc / MC.Xtry)
    end
end
