#!/usr/bin/env julia

using Random
using Printf
using ACFlow

# Setup parameters
wmin = -5.0  # Left boundary
wmax = +5.0  # Right boundary
nmesh = 2001 # Number of real-frequency points
niw  = 20    # Number of Matsubara frequencies
beta = 40.0  # Inverse temperature
Ïµâ‚�   = 2.00  # Parameters for gaussian peaks
Ïµâ‚‚   = -2.0
Aâ‚�   = 1.00
Aâ‚‚   = 1.00
Î“â‚�   = 0.50
Î“â‚‚   = 0.50

# Real frequency mesh
rmesh = collect(LinRange(wmin, wmax, nmesh))

# Initial spectral function
image1 = similar(rmesh)
@. image1  = Aâ‚� * exp(-(rmesh - Ïµâ‚�) ^ 2.0 / (2.0 * Î“â‚� ^ 2.0)) / (Î“â‚� * sqrt(2.0 * Ï€))
image1 = image1 ./ trapz(rmesh, image1)
#
image2 = similar(rmesh)
@. image2  = Aâ‚‚ * exp(-(rmesh - Ïµâ‚‚) ^ 2.0 / (2.0 * Î“â‚‚ ^ 2.0)) / (Î“â‚‚ * sqrt(2.0 * Ï€))
image2 = image2 ./ trapz(rmesh, image2)
#
ğ�”¸ = zeros(F64, (2,2,nmesh))
ğ�”¸[1,1,:] .= image1
ğ�”¸[2,2,:] .= image2

# Rotate spectral function to generate non-diagonal element
#
# Set rotation angle
Î¸ = 0.1
#
# Build rotation matrix
â„� = [cos(Î¸) sin(Î¸); -sin(Î¸) cos(Î¸)]
#
# Get final spectral function
ğ�’œ = zeros(F64, (2,2,nmesh))
for w = 1:nmesh
    ğ�’œ[:,:,w] = â„� * ğ�”¸[:,:,w] * â„�'
end

# Matsubara frequency mesh
iw = Ï€ / beta * (2.0 * collect(0:niw-1) .+ 1.0)

# Kernel function
kernel = 1.0 ./ (im * reshape(iw, (niw,1)) .- reshape(rmesh, (1,nmesh)))

# Build green's function
KA = reshape(kernel, (1,1,niw,nmesh)) .* reshape(ğ�’œ, (2,2,1,nmesh))
giw = zeros(C64, (2,2,niw))
for i = 1:2
    for j = 1:2
        for w = 1:niw
            giw[i,j,w] = trapz(rmesh, KA[i,j,w,:])
        end
    end
end

# Build error
err = 1e-5

# Write green's function
# For diagonal part
open("giw.11.data", "w") do fout
    for i = 1:niw
        z = giw[1,1,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iw[i], real(z), imag(z), err)
    end
end
#
# For non-diagonal part
open("giw.12.data", "w") do fout
    for i = 1:niw
        z = giw[1,2,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iw[i], real(z), imag(z), err)
    end
end
#
# For non-diagonal part
open("giw.21.data", "w") do fout
    for i = 1:niw
        z = giw[2,1,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iw[i], real(z), imag(z), err)
    end
end
#
# For diagonal part
open("giw.22.data", "w") do fout
    for i = 1:niw
        z = giw[2,2,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iw[i], real(z), imag(z), err)
    end
end
#
# Be careful, we use two different auxiliary green's functions here
# For auxiliary green's function
open("giw.aux12.data", "w") do fout
    for i = 1:niw
        z = giw[1,1,i] + giw[2,2,i] + 2 * giw[1,2,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iw[i], real(z), imag(z), err)
    end
end
#
# For auxiliary green's function
open("giw.aux21.data", "w") do fout
    for i = 1:niw
        z = giw[1,1,i] + giw[2,2,i] - 2 * giw[2,1,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iw[i], real(z), imag(z), err)
    end
end

# Write spectral function
# For diagonal part
open("image.11.data", "w") do fout
    for i in eachindex(rmesh)
        @printf(fout, "%20.16f %20.16f\n", rmesh[i], ğ�’œ[1,1,i])
    end
end
#
# For non-diagonal part
open("image.12.data", "w") do fout
    for i in eachindex(rmesh)
        @printf(fout, "%20.16f %20.16f\n", rmesh[i], ğ�’œ[1,2,i])
    end
end
#
# For non-diagonal part
open("image.21.data", "w") do fout
    for i in eachindex(rmesh)
        @printf(fout, "%20.16f %20.16f\n", rmesh[i], ğ�’œ[2,1,i])
    end
end
#
# For diagonal part
open("image.22.data", "w") do fout
    for i in eachindex(rmesh)
        @printf(fout, "%20.16f %20.16f\n", rmesh[i], ğ�’œ[2,2,i])
    end
end
