__precompile__()

module Stats


########################
# Load external modules
########################

using CSV, DataFrames
using Statistics

export calculate_statistics,
       extract_parameters

function extract_parameters(filename :: String)
        # Extracting parameters
        params_line = readlines(filename)[1]

        # Split the line by commas and trim each part
        params_parts = map(strip, split(params_line, ','))

        # Initialize variables to store parameter values
        α_value, τ₀_value, σ₀_value = missing, missing, missing

        # Look for specific substrings to identify the values of α, τ₀, and σ₀
        for param_part in params_parts
            if contains(param_part, "α = ")
                α_value = parse(Float64, split(param_part, '=')[2])
            elseif contains(param_part, "τ₀ = ")
                τ₀_value = parse(Float64, split(param_part, '=')[2])
            elseif contains(param_part, "σ₀ = ")
                σ₀_value = parse(Float64, split(param_part, '=')[2])
            end
        end

        # Assign the values to α, τ₀, and σ₀
        α = α_value
        τ₀ = τ₀_value
        σ₀ = σ₀_value
        return α, τ₀, σ₀
end

const default_imnames = ("lighthouse200x300", "shepplogan256x256", "brainphantom256x256")
const default_algnames = ("dualscaling", "greedy","noprediction","primalonly","proximal","rotation","zerodual","activateddual")

function calculate_statistics(;imnames=default_imnames,
                               algnames=default_algnames,
                               csv_path = "./img/summarystats.csv")
    mystart = 41 # Corresponds to the 500th iterate
    
    # Define an array to store results
    results = DataFrame(experiment = String[], α = Float64[], algorithm = String[], psnr_mean1 = Float64[], psnr_mean500 = Float64[], psnr_ci = String[], ssim_mean1 = Float64[], ssim_mean500 = Float64[], ssim_ci = String[])

    for imname in imnames
        for algname in algnames
            directory_path = "./img/"
            files = readdir(directory_path)
            filtered_files = filter(file -> startswith(file, "$(imname)_pdps_known_$(algname)") && endswith(file, ".txt"), files)

            for file in filtered_files
                filename = directory_path * file
                data = CSV.File(filename, delim='\t', header=2) |> DataFrame

                # Extract α from filename
                α, _, _ = extract_parameters(filename)


                # Extract SSIM and PSNR columns starting from 1st iteration
                ssim_values1 = Float64.(data[:, :ssim])
                psnr_values1 = Float64.(data[:, :psnr])

                # Extract SSIM and PSNR columns starting from 500th iteration
                ssim_values500 = Float64.(data[mystart:end, :ssim])
                psnr_values500 = Float64.(data[mystart:end, :psnr])

                # Calculate mean and confidence intervals
                ssim_mean1 = round(mean(ssim_values1), digits=4)
                psnr_mean1 = round(mean(psnr_values1), digits=4)

                ssim_mean500 = round(mean(ssim_values500), digits=4)
                psnr_mean500 = round(mean(psnr_values500), digits=4)
                ssim_std500 = round(std(ssim_values500), digits=4)
                psnr_std500 = round(std(psnr_values500), digits=4)
                n = length(ssim_values500)

                ssim_ci_lower = round(ssim_mean500 - 1.96 * ssim_std500 / sqrt(n), digits=4)
                ssim_ci_upper = round(ssim_mean500 + 1.96 * ssim_std500 / sqrt(n), digits=4)
                psnr_ci_lower = round(psnr_mean500 - 1.96 * psnr_std500 / sqrt(n), digits=4)
                psnr_ci_upper = round(psnr_mean500 + 1.96 * psnr_std500 / sqrt(n), digits=4)

                ssim_ci = "$(ssim_ci_lower) - $(ssim_ci_upper)"
                psnr_ci = "$(psnr_ci_lower) - $(psnr_ci_upper)"
                experiment = "$(imname)"
                algorithm = "$(algname)"

                # Append results to DataFrame
                push!(results, (experiment, α, algorithm, psnr_mean1, psnr_mean500, psnr_ci, ssim_mean1, ssim_mean500, ssim_ci))
            end
        end
    end
    sort!(results, [:experiment, :α, :algorithm])
    if isfile(csv_path)
    	rm(csv_path)
    end
        CSV.write(csv_path, results)
end

end
