module ElPhDynamics

using LinearAlgebra
using FFTW
using TOML

#######################
## INCLUDING MODULES ##
#######################

include("Utilities.jl")

include("UnitCells.jl")

include("Lattices.jl")

include("Checkerboard.jl")

include("IterativeSolvers.jl")

include("TimeFreqFFTs.jl")

include("Models.jl")

include("KPMPreconditioners.jl")

include("InitializePhonons.jl")

include("PhononAction.jl")

include("FourierAcceleration.jl")

include("LangevinDynamics.jl")

include("HMC.jl")

include("SpecialUpdates.jl")

include("GreensFunctions.jl")

include("SimulationParams.jl")

include("Measurements.jl")

include("MuFinder.jl")

include("RunSimulation.jl")

include("SimulationSummary.jl")

include("ProcessInputFile.jl")

####################################
## DEFINING HIGHET LEVEL FUNCTION ##
##     TO RUN A SIMULATION        ##
####################################

using ..RunSimulation: run_simulation!
using ..ProcessInputFile: process_input_file, process_checkpoint, initialize_model
using ..Models: read_phonons!
using ..MuFinder: MuTuner, update_μ!
using ..SimulationSummary: write_simulation_summary!

export simulate, load_model

"""
Highest level function used to run a QMC simulation.
To run a simulation execute the following command:
`julia -O3 -e "using ElPhDynamics; simulate(ARGS)" -- input.toml`
"""
function simulate(args)

    # set the number of threads used by BLAS and FFTW libraries
    BLAS.set_num_threads(1)
    FFTW.set_num_threads(1)

    ######################
    ## READ CONFIG FILE ##
    ######################

    # getting iput filename
    config_file = args[1]

    # read input file
    input = TOML.parsefile(config_file)

    # construct data folder name
    filepath   = input["simulation"]["filepath"]
    foldername = input["simulation"]["foldername"]
    if length(args)==1
        datafolder = name_datafolder(filepath,foldername)
    else
        id = parse(Int,args[2])
        datafolder = name_datafolder(filepath,foldername,id)
    end
    input["simulation"]["datafolder"] = datafolder

    ###########################
    ## INITIALIZE SIMULATION ##
    ###########################

    if isdir(datafolder) # resume previous simulation

        # extract state from checkpoint
        (model, Gr, μ_tuner, sim_params, simulation_dynamics, burnin_dynamics,
         burnin_reflect_update, sim_reflect_update, burnin_swap_update, sim_swap_update,
         fourier_accelerator, preconditioner, container, burnin_start, sim_start, sim_stats) = process_checkpoint(input)

    else # start new simulation

        # initialize new simulation
        (model, Gr, μ_tuner, sim_params, simulation_dynamics, burnin_dynamics,
         burnin_reflect_update, sim_reflect_update, burnin_swap_update, sim_swap_update,
         fourier_accelerator, preconditioner, container, burnin_start, sim_start, sim_stats) = process_input_file(config_file,input)
    end

    ####################
    ## RUN SIMULATION ##
    ####################

    if sim_start < sim_params.nsteps
        sim_stats = run_simulation!(model, Gr, μ_tuner, sim_params, simulation_dynamics, burnin_dynamics,
                                    sim_reflect_update, burnin_reflect_update, sim_swap_update, burnin_swap_update,
                                    fourier_accelerator, container, preconditioner, sim_stats, burnin_start, sim_start)
    end

    #################################
    ## SUMARIZE SIMULATION RESULTS ##
    #################################

    if sim_start < sim_params.nsteps
        write_simulation_summary!(datafolder)
    end

    return nothing
end


"""
Pass this function a directory name generated by a simulation and it returns a HolsteinModel object
with the phonons intialized to the final phonon configuration sampled in the simulation.
"""
function load_model(dir::String)

    files  = readdir(dir)
    config = findall(f -> endswith(f, r"\.toml|\.TOML"), files)
    phonon = findall(f -> endswith(f, "_config.out"), files)
    @assert length(config) == length(phonon) == 1
    config_file = joinpath(dir, files[config[1]])
    phonon_file = joinpath(dir, files[phonon[1]])
    input = TOML.parsefile(config_file)
    input["simulation"]["datafolder"] = dir
    model = initialize_model(input)
    read_phonons!(model, phonon_file)
    
    return model
end

###########################################
## PRIVATE METHODS THAT ARE NOT EXPORTED ##
###########################################

"""
Name the data folder.
"""
function name_datafolder(filepath::String,foldername::String,id::Int=0)

    # data folder name without id at end
    key = joinpath(filepath,foldername)

    if id>0 # if ID given then attach to end of key
        datafolder = "$key-$id"
    else # determine the ID
        id = 1
        while true
            datafolder = "$key-$id"
            if !isdir(datafolder)
                break
            else
                id += 1
            end
        end
    end

    return datafolder
end

end
