using LatticeGPU
using TOML
using ArgParse
using CUDA
using BDIO
using InteractiveUtils
using MD5


"""
    function read_input()

Stores as global variables 'parsed_args' (info from the command line) and 'params' (info from the input file)
"""
function read_input()
    global parsed_args = parse_commandline()
    global params = TOML.parsefile(parsed_args["i"])

    if parsed_args["R"] && (isfile("./output/"*params["Run"]["name"]*".log") || isfile("./output/"*params["Run"]["name"]*".bdio"))
        error("Output file already exists with this run name. Use flag -R to continue the run.")
    end

    return nothing
end


function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "-i"
        help = "Input parameters file"
        required = true
        arg_type = String

        "-c"
        help = "Gauge configuration list file"
        arg_type = String

        "--cern"
        help = "Config written with the export_cnfg_cern() convention"
        action = :store_true

        "-G"
        help = "GPU for the job. Default 0"
        arg_type = Int
        default = 0

        "-R"
        help = "Continue run. Needed to edit output files already created."
        action = :store_false
    end

    return parse_args(s)
end



"""
    function load_structs()

Stores in global variables the needed structures, i.e. lp, gp, dpar, dws, ymws
"""
function load_structs()

    CUDA.device!(parsed_args["G"])
    global lp = SpaceParm{4}(tuple(params["Space"]["size"]...), tuple(params["Space"]["blocks"]...),params["Space"]["bc"], (0,0,0,0,0,0))
    global gp = GaugeParm{Float64}(SU3{Float64},params["HMC"]["beta"],params["HMC"]["c0"],(params["HMC"]["cG"],0.0),(0.0,0.0),lp.iL);
    global dws = DiracWorkspace(SU3fund,Float64,lp);
    global ymws = YMworkspace(SU3,Float64,lp);
    global int = wfl_rk3(Float64, params["Frontflow"]["epsilon"], params["Backflow"]["tol"])
    global intsch = omf4(Float64,params["HMC"]["eps"], params["HMC"]["ns"]);

    global fernames = [replace(k, "Fermion" => "") for k in keys(params) if startswith(k, "Fermion")]
    global dpar = [DiracParam{Float64}(SU3fund,(1/(2*params["Fermion"*f]["kappa"])) - 4,params["Fermion"*f]["csw"],ntuple(i -> exp(((i!=4)*im*params["Fermion"*f]["theta"]/lp.iL[i]) + ((i==4)*im*params["Fermion"*f]["theta_t"]/lp.iL[i])),4),0.0,params["Fermion"*f]["ct"]) for f in fernames];

    global flow_times = params["Backflow"]["Flow_times"]
    global MCid = 0
    if params["Space"]["bc"] == 0
        global V4 = prod(lp.iL[1:3])*(lp.iL[4]-1)
    else
        global V4 = prod(lp.iL)
    end
    global Nmc_acc = 0
    global RUN_ON = true

    return nothing
end

function write_log()

    println(log_file,"Running Ferflow.jl by ", params["Run"]["user"],". Name of the run: ",params["Run"]["name"])
    println(log_file,"")

    print(log_file,"Calling: ")
    print(log_file,PROGRAM_FILE*" "); for x in ARGS; print(log_file,x*" "); end
    println(log_file,"\n")

    println(log_file,"Version info:")
    versioninfo(log_file)
    println(log_file,"")

    println(log_file,"Reading input file from:", parsed_args["i"], "...\n")

    println(log_file,"Parameters:")
    println(log_file,"Lattice size:        ", lp.iL)
    println(log_file,"Boundary conditions: ", params["Space"]["bc"])
    for i in fernames
        println(log_file,"Fermion"*i*" parameters:")
        println(log_file,"kappa              = ", params["Fermion"*f]["kappa"])
        println(log_file,"theta_t            = ", params["Fermion"*f]["theta_t"])
        println(log_file,"theta              = ", params["Fermion"*f]["theta"])
        println(log_file,"csw                = ", params["Fermion"*f]["csw"])
        println(log_file,"tolerance          = ", params["Fermion"*f]["tolerance"])
        println(log_file,"maxiter            = ", params["Fermion"*f]["maxiter"])
    end
    println(log_file,"N noise frontflow  = ", params["Frontflow"]["N_noise"])
    println(log_file,"t_zero             = ", params["Frontflow"]["t_zero"])
    println(log_file,"eps frontflw       = ", params["Frontflow"]["epsilon"])
    println(log_file,"t_source frontflow = ", params["Frontflow"]["tsource"])
    println(log_file,"Flow times         = ", params["Backflow"]["Flow_times"])
    println(log_file,"N noise backflow   = ", params["Backflow"]["N_noise"])
    println(log_file,"t_source backflow  = ", params["Backflow"]["tsource"])
    println(log_file,"Backflow tolerance = ", params["Backflow"]["tol"])

    flush(log_file)
    return nothing
end


function save_data()

    ihdr = [convert(Int32,1730280201)]
    fname = "./output/"*params["Run"]["name"]*".bdio"

    if isfile(fname)
        fb = BDIO_open(fname, "a")
        println(log_file,"\nAppending output to "*fname*"\n")
    else
        fb = BDIO_open(fname, "w", "BDIO output from ferflow.jl")
        println(log_file,"Creating new BDIO output file "*fname)

        BDIO_start_record!(fb, BDIO_BIN_GENERIC, 14)
        BDIO_write!(fb, ihdr)
        BDIO_write_hash!(fb)

        BDIO_start_record!(fb, BDIO_BIN_GENERIC, 1)
        BDIO_write!(fb, [convert(Int32, 4)])
        BDIO_write!(fb, [convert(Int32, lp.iL[i]) for i in 1:4])
        BDIO_write!(fb, [convert(Int32, length(fernames))])
        for k in 1:length(fernames)
            BDIO_write!(fb, [dpar[k].m0, dpar[k].csw])
            BDIO_write!(fb, [dpar[k].th[i] for i in 1:4])
        end
        BDIO_write!(fb, [convert(Int32,params["Frontflow"]["N_noise"])])
        BDIO_write!(fb, [convert(Int32,params["Frontflow"]["nsteps"])])
        BDIO_write!(fb, [convert(Int32,params["Backflow"]["N_noise"])])
        BDIO_write!(fb, [convert(Int32,length(flow_times))])
        BDIO_write!(fb, [params["Frontflow"]["t_zero"],])
        BDIO_write!(fb, [params["Frontflow"]["epsilon"],])
        BDIO_write!(fb, [convert(Int32,params["Frontflow"]["tsource"])])
        BDIO_write!(fb, [convert(Int32,params["Backflow"]["tsource"])])
        length(flow_times) > 0 ? BDIO_write!(fb, flow_times) : nothing

        BDIO_write_hash!(fb)
    end


    BDIO_start_record!(fb, BDIO_BIN_GENERIC, 8, true)

    for f in 1:length(dpar)
        for noi in 1:params["Frontflow"]["N_noise"]
            BDIO_write!(fb,pp_corr_t0[:,noi,f])
            BDIO_write!(fb,ap_corr_t0[:,noi,f])
            BDIO_write!(fb,pphat_t0[:,noi,f])
            BDIO_write!(fb,pptilde_t0[:,noi,f])
            for fl in 1:params["Frontflow"]["nsteps"]+1
                BDIO_write!(fb,pp_corr_t[:,noi,fl,f])
                BDIO_write!(fb,ap_corr_t[:,noi,fl,f])

                BDIO_write!(fb,pphat_t[:,noi,fl,f])
                BDIO_write!(fb,pptilde_t[:,noi,fl,f])
            end
        end
    end


    BDIO_write!(fb,Eoft0)
    BDIO_write!(fb,Qt0)
    for fl in 1:params["Frontflow"]["nsteps"]+1
        BDIO_write!(fb,Eoft[fl,:])
        BDIO_write!(fb,Qt[fl,:])
    end

    for f in 1:length(dpar)
        for noi in 1:params["Backflow"]["N_noise"]
            for fl in 1:length(flow_times)
                BDIO_write!(fb,pp_corr_tfl[:,noi,fl,f])
                BDIO_write!(fb,ap_corr_tfl[:,noi,fl,f])

                BDIO_write!(fb,Quark_cond[:,noi,fl,f])
                BDIO_write!(fb,Quark_cond_cfl[:,noi,fl,f])

                BDIO_write!(fb,Quark_cond2[:,noi,fl,f])
                BDIO_write!(fb,Quark_cond2_cfl[:,noi,fl,f])

                BDIO_write!(fb,ChiDchi[:,noi,fl,f])
                BDIO_write!(fb,ChiDchi_cfl[:,noi,fl,f])
            end
        end
    end

    BDIO_write_hash!(fb)

    BDIO_close!(fb)

    return nothing
end



function read_cnfg_cern(path::String,lp::SpaceParm)
    Ucpu = Array{SU3{Float64}, 3}(undef, lp.bsz, lp.ndim, lp.rsz)
    file = open(path)
    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]

                    for d in [4,1,2,3]

                        f,r = point_index(CartesianIndex((i,j,k,t)),lp)

                        #a11 !!

                        re11 = read(file,Float64)
                        co11 = read(file,Float64)

                        #a12 !!

                        re12 = read(file,Float64)
                        co12 = read(file,Float64)

                        #a13 !!

                        re13 = read(file,Float64)
                        co13 = read(file,Float64)

                        #a21 !!

                        re21 = read(file,Float64)
                        co21 = read(file,Float64)

                        #a22 !!

                        re22 = read(file,Float64)
                        co22 = read(file,Float64)

                        #a23 !!

                        re23 = read(file,Float64)
                        co23 = read(file,Float64)

                        #a31

                        re31 = read(file,Float64)
                        co31 = read(file,Float64)

                        #a32

                        re32 = read(file,Float64)
                        co32 = read(file,Float64)

                        #a33

                        re33 = read(file,Float64)
                        co33 = read(file,Float64)

                        (Ucpu[f,d,r] = SU3{Float64}(re11 + im*co11, re12 + im*co12, re13 + im*co13,
                                                    re21 + im*co21, re22 + im*co22, re23 + im*co23))
                    end
                end
            end
        end
    end
    length(read(file)) == (prod(lp.iL[1:3])*4*8*9*2) ? nothing : error("File not fully read")
    close(file)

    return CuArray(Ucpu)
end
