using LatticeGPU
using TOML
using TimerOutputs
using ArgParse
using CUDA

function load_fields()
    global U = vector_field(SU3{Float64}, lp)
    fill!(U,one(SU3{Float64}));

    global U_CPU = Array(U)
    global psi = scalar_field(Spinor{4,SU3fund{Float64}},lp)
    global psi_CPU = Array(psi)

    global pp_density = Array{Float64}(undef,lp.bsz,lp.rsz);
    global ap_density = Array{Complex{Float64}}(undef,lp.bsz,lp.rsz);

    global pp_corr_t0 = fill(zero(Float64),(lp.iL[4],N_noise));
    global ap_corr_t0 = fill(zero(ComplexF64),(lp.iL[4],N_noise));

    global pp_corr_t = fill(zero(Float64),(lp.iL[4],N_noise,length(flow_times)));
    global ap_corr_t = fill(zero(ComplexF64),(lp.iL[4],N_noise,length(flow_times)));

    global Eoft = Array{Complex{Float64}}(undef,1+length(flow_times),lp.iL[4]);
    global Eofpla = Array{Complex{Float64}}(undef,lp.iL[4],lp.npls);

    global phat_t = Array{Float64}(undef,N_noise,length(flow_times));
    global Quark_cond = Array{Complex{Float64}}(undef,lp.iL[4],N_noise,length(flow_times));
    global Quark_cond_cfl = Array{Complex{Float64}}(undef,lp.iL[4],N_noise,length(flow_times));

    return nothing
end


"""
function two_pt()

Stores the two point function in the variables 'pp_corr_t0', 'ap_corr_t0', 'pp_corr_t' and 'ap_corr_t'
"""
function two_pt()

    println(log_file,"Measuring 2pt...")
    flush(log_file)

    Eoft_plaq(Eofpla,U, gp, lp, ymws)
    Eoft[1,:] .= sum(Eofpla,dims = 2)

    for noi in 1:N_noise

        niter = propagator!(psi,U, dpar, dws, lp,params["Measurements"]["maxiter"], params["Measurements"]["tolerance"],params["Fermion"]["tsource"])

        println(log_file,"CG converged after ",niter," iterations with residue ",CUDA.mapreduce(x -> norm2(x), +, dws.sr))
        flush(log_file)

        pp_density .= Array(norm2.(psi))
        ap_density .= Array(dot.(psi,dmul.(Gamma{4},psi)))

        pp_corr_t0[:,noi] .= zero(Float64)
        ap_corr_t0[:,noi] .= zero(ComplexF64)
        for t in 1:lp.iL[4]
            for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                pp_corr_t0[t,noi] += pp_density[b,r]
                ap_corr_t0[t,noi] += ap_density[b,r]
            end end end
        end

        delta_t = flow_times .- [0.0,flow_times[1:end-1]...]
        neps = int.eps_ini
        for fl in 1:length(flow_times)
            _,epslist = flw_adapt(U, psi, int, delta_t[fl], neps, gp, dpar, lp, ymws, dws)
            neps = epslist[end]

            if noi ==1
                Eoft_plaq(Eofpla,U, gp, lp, ymws)
                Eoft[fl+1,:] .= sum(Eofpla,dims = 2)
            end

            pp_density .= Array(norm2.(psi))
            ap_density .= Array(dot.(psi,dmul.(Gamma{4},psi)))
            pp_corr_t[:,noi,fl] .= zero(Float64)
            ap_corr_t[:,noi,fl] .= zero(ComplexF64)
            for t in 1:lp.iL[4]
                for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                    b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                    pp_corr_t[t,noi,fl] += pp_density[b,r]
                    ap_corr_t[t,noi,fl] += ap_density[b,r]
                end end end
            end
        end

        @timeit "CPU to GPU" copyto!(U,U_CPU)
    end

    pp_corr_t0 .= pp_corr_t0./prod(lp.iL[1:3])
    ap_corr_t0 .= ap_corr_t0./prod(lp.iL[1:3])
    pp_corr_t .= pp_corr_t./prod(lp.iL[1:3])
    ap_corr_t .= ap_corr_t./prod(lp.iL[1:3])

    println(log_file,"Finished measuring 2pt")
    flush(log_file)

    return nothing
end



"""
function one_pt()

Stores the one point function in the variable 'Quark_cond' and 'Quark_cond_cfl'.
The final result is 'Quark_cond' + C_fl*'Quark_cond_cfl'
"""
function one_pt()

    @timeit "CPU to GPU" copyto!(U,U_CPU)

    for fl in 1:length(flow_times)
        println(log_file,"Measuring 1pt at t = "*string(flow_times[fl])*"...")
        flush(log_file)

        for noi in 1:N_noise

            ## The change in the field before randomize should be moved to the package
            psi .= 0.0*psi
            pfrandomize!(psi,lp,params["Fermion"]["tsource"])

            backflow(psi, U, flow_times[fl], params["Measurements"]["Nsaves"], gp, dpar, lp, ymws, dws)

            @timeit "GPU to CPU" psi_CPU .= Array(psi)

            @timeit "Inversion" begin
                dws.sp .= dmul.(Gamma{5},psi)
                g5Dw!(psi,U,dws.sp,dpar,dws,lp)

                niter = CG!(psi,U,DwdagDw!,dpar,lp,dws,params["Measurements"]["maxiter"], params["Measurements"]["tolerance"])
            end

            println(log_file,"CG converged after ",niter," iterations with residue ",CUDA.mapreduce(x -> norm2(x), +, dws.sr))
            flush(log_file)

            @timeit "CPU to GPU" copyto!(dws.sp,psi_CPU)
            Quark_cond[:,noi,fl] .= zero(ComplexF64)
            Quark_cond_cfl[:,noi,fl] .= zero(ComplexF64)
            for t in 1:lp.iL[4]
                for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                    b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                    CUDA.@allowscalar Quark_cond[t,noi,fl] += dot(dws.sp[b,r],psi[b,r])
                    CUDA.@allowscalar Quark_cond_cfl[t,noi,fl] += -norm2(dws.sp[b,r])
                end end end
            end
            @timeit "CPU to GPU" copyto!(U,U_CPU)

        end
    end

    Quark_cond .= Quark_cond./prod(lp.iL[1:3])
    Quark_cond_cfl .= Quark_cond_cfl./prod(lp.iL[1:3])
    println(log_file,"Finished measuring 1pt")

    return nothing
end



"""
function two_pt_rnd()

Stores the random field evolved in flow time in the variable 'phat_t'
"""
function two_pt_rnd()

    copyto!(U,U_CPU)

    println(log_file,"Measuring PhatPt correlator...")
    flush(log_file)
    for noi in 1:N_noise

        pfrandomize!(psi,lp)

        delta_t = flow_times .- [0.0,flow_times[1:end-1]...]
        neps = int.eps_ini
        for fl in 1:length(flow_times)
            _,epslist = flw_adapt(U, psi, int, delta_t[fl], neps, gp, dpar, lp, ymws, dws)
            neps = epslist[end]

            phat_t[noi,fl] = sum(norm2.(psi))
        end

        copyto!(U,U_CPU)
    end

    phat_t .= phat_t./V4

    println(log_file,"Finished measuring PhatPt")
    flush(log_file)

    return nothing
end
