#!/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 = 20.0  # Inverse temperature
Ïµâ‚�   = 2.00  # Parameters for gaussian peaks
Ïµâ‚‚   = -2.0
Ïµâ‚ƒ   = 1.00
Ïµâ‚„   = -1.0
Aâ‚�   = 0.50
Aâ‚‚   = 0.50
Aâ‚ƒ   = 0.50
Aâ‚„   = 0.50
Î·â‚�   = 1e-2
Î·â‚‚   = 1e-2
Î·â‚ƒ   = 1e-2
Î·â‚„   = 1e-2

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

# Matsubara frequency mesh
iÏ‰â‚™ = Ï€ / beta * (2.0 * collect(0:niw-1) .+ 1.0)

# Noise
seed = rand(1:100000000)
rng = MersenneTwister(seed)
noise_ampl = 1.0e-4
noise_abs = randn(rng, F64, niw) * noise_ampl
noise_phase = rand(rng, niw) * 2.0 * Ï€
noise = noise_abs .* exp.(noise_phase * im)

# Initial green's function (in Matsubara axis)
giw1 = zeros(C64, niw)
for i in eachindex(giw1)
    giw1[i] = (
        Aâ‚� / (iÏ‰â‚™[i] * im - Ïµâ‚�) + Aâ‚ƒ / (iÏ‰â‚™[i] * im - Ïµâ‚ƒ) + noise[i]
    )
end
#
giw2 = zeros(C64, niw)
for i in eachindex(giw2)
    giw2[i] = (
        Aâ‚‚ / (iÏ‰â‚™[i] * im - Ïµâ‚‚) + Aâ‚„ / (iÏ‰â‚™[i] * im - Ïµâ‚„) + noise[i]
    )
end
#
ğ�”¾iw = zeros(C64, (2,2,niw))
ğ�”¾iw[1,1,:] .= giw1
ğ�”¾iw[2,2,:] .= giw2

# Initial green's function (in real axis)
gre1 = zeros(C64, nmesh)
for i in eachindex(gre1)
    gre1[i] = (
        Aâ‚� / (Ï‰[i] + Î·â‚� * im - Ïµâ‚�) + Aâ‚ƒ / (Ï‰[i] + Î·â‚ƒ * im - Ïµâ‚ƒ)
    )
end
#
gre2 = zeros(C64, nmesh)
for i in eachindex(gre2)
    gre2[i] = (
        Aâ‚‚ / (Ï‰[i] + Î·â‚‚ * im - Ïµâ‚‚) + Aâ‚„ / (Ï‰[i] + Î·â‚„ * im - Ïµâ‚„)
    )
end
#
ğ�”¾re = zeros(C64, (2,2,nmesh))
ğ�”¾re[1,1,:] .= gre1
ğ�”¾re[2,2,:] .= gre2

# Rotate green's function to generate non-diagonal element
#
# Set rotation angle
Î¸ = 0.1
#
# Build rotation matrix
â„� = [cos(Î¸) sin(Î¸); -sin(Î¸) cos(Î¸)]
#
# Get final green's function (in Matsubara axis)
ğ�’¢iw = zeros(C64, (2,2,niw))
for w = 1:niw
    ğ�’¢iw[:,:,w] = â„� * ğ�”¾iw[:,:,w] * â„�'
end
#
# Get final green's function (in real axis)
ğ�’¢re = zeros(C64, (2,2,nmesh))
for w = 1:nmesh
    ğ�’¢re[:,:,w] = â„� * ğ�”¾re[:,:,w] * â„�'
end

# Evaluate final spectral function
ğ�’œ = -imag.(ğ�’¢re) / Ï€

# Build error
err = 1e-5

# Write green's function
# For diagonal part
open("giw.11.data", "w") do fout
    for i = 1:niw
        z = ğ�’¢iw[1,1,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iÏ‰â‚™[i], real(z), imag(z), err)
    end
end
#
# For non-diagonal part
open("giw.12.data", "w") do fout
    for i = 1:niw
        z = ğ�’¢iw[1,2,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iÏ‰â‚™[i], real(z), imag(z), err)
    end
end
#
# For non-diagonal part
open("giw.21.data", "w") do fout
    for i = 1:niw
        z = ğ�’¢iw[2,1,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iÏ‰â‚™[i], real(z), imag(z), err)
    end
end
#
# For diagonal part
open("giw.22.data", "w") do fout
    for i = 1:niw
        z = ğ�’¢iw[2,2,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iÏ‰â‚™[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 = ğ�’¢iw[1,1,i] + ğ�’¢iw[2,2,i] + 2 * ğ�’¢iw[1,2,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iÏ‰â‚™[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 = ğ�’¢iw[1,1,i] + ğ�’¢iw[2,2,i] - 2 * ğ�’¢iw[2,1,i]
        @printf(fout, "%20.16f %20.16f %20.16f %20.16f\n", iÏ‰â‚™[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(Ï‰)
        @printf(fout, "%20.16f %20.16f\n", Ï‰[i], ğ�’œ[1,1,i])
    end
end
#
# For non-diagonal part
open("image.12.data", "w") do fout
    for i in eachindex(Ï‰)
        @printf(fout, "%20.16f %20.16f\n", Ï‰[i], ğ�’œ[1,2,i])
    end
end
#
# For non-diagonal part
open("image.21.data", "w") do fout
    for i in eachindex(Ï‰)
        @printf(fout, "%20.16f %20.16f\n", Ï‰[i], ğ�’œ[2,1,i])
    end
end
#
# For diagonal part
open("image.22.data", "w") do fout
    for i in eachindex(Ï‰)
        @printf(fout, "%20.16f %20.16f\n", Ï‰[i], ğ�’œ[2,2,i])
    end
end
