using FileIO 
using Parameters
using Statistics
include("LS_calculations.jl")

#=
Change these parameters to inform the analysis which eta values to check,
which directory to search in, how many network realizations to expect,
and also information about timesteps and dt.
=#
etas = -1.0:0.1:1.0
dir = "results1"
outdir = "results1-analysis"
trials = 20
N = 3000
dt = 0.01
maxlag_steps = 100000


# Allocate space to store the autocorrelation and Lyapunov spectrum data
avg_autocorrs = Array{Float64,3}(undef, maxlag_steps + 1, trials, length(etas))
PCA_dims = Array{Float64,2}(undef, trials, length(etas))
Lspecs = Array{Float64,3}(undef, N, trials, length(etas))
max_LEs = Array{Float64,2}(undef, trials, length(etas))

avg_autocorrs_trial_counts = zeros(Int64, length(etas))
Lspecs_trial_counts = zeros(Int64, length(etas))

#=
Finds and loads all jld2 files, then checks their keys to see if
they contain the data that we need to plot stuff
=#
for (path, dirname, filenames) in walkdir(dir)
    for fn in filenames
        if endswith(fn, ".jld2") && startswith(fn, "AC_PCAdim")
            # Next experiment, save the eta parameter. Regexes are annoying
            eta = parse(Float64, match(r",eta=[0-9.-]+,", fn).match[6:end-1])
            match_indices = isapprox.(etas, eta)
            sum(match_indices) != 1 && @warn "Multiple eta values for file $fn found"
            if any(match_indices)
                etaind = findlast(match_indices)
                # Increment the trial counts 
                avg_autocorrs_trial_counts[etaind] += 1
                trial_count = avg_autocorrs_trial_counts[etaind]
                # Fill the appropriate vector in the data array
                avg_autocorrs[:,trial_count,etaind] = load(joinpath(path, fn), 
                        "avg_C")
                PCA_dims[trial_count,etaind] = load(joinpath(path, fn), "PR")
            else 
                @warn "File $fn with invalid eta value found"
            end
        elseif endswith(fn, ".jld2") && startswith(fn, "Lspec")
            eta = load(joinpath(path, fn), "eta")
            match_indices = isapprox.(etas, eta)
            sum(match_indices) != 1 && @warn "Multiple eta values for file $fn found"
            if any(match_indices)
                etaind = findlast(match_indices)
                # Increment the trial counts 
                Lspecs_trial_counts[etaind] += 1
                trial_count = Lspecs_trial_counts[etaind]
                # Fill the appropriate vector in the data array
                Lspecs[:,trial_count,etaind] = load(joinpath(path, fn), 
                        "Lspec")
                max_LEs[trial_count,etaind] = load(joinpath(path, fn), 
                        "max_LE")
            else 
                @warn "File $fn with invalid eta value found"
            end
        end
    end 
end 

# Validate that we got all the data 
for (index, count1, count2) in zip(1:length(etas), avg_autocorrs_trial_counts, Lspecs_trial_counts)
    if count1 != trials 
        @warn "Found only $count1 of $trials autocorr data files for eta=$(etas[index])"
    end 
    if count2 != trials 
        @warn "Found only $count2 of $trials Lspec data files for eta=$(etas[index])"
    end
end 


# Autocorrelation stats
ensbl_avg_Cs = reshape(mean(avg_autocorrs, dims=2), 
        (maxlag_steps+1, length(etas)))
ensbl_std_Cs = reshape(std(avg_autocorrs, dims=2), 
        (maxlag_steps+1, length(etas)))
timelags = dt .* collect(0:maxlag_steps)

# PCA dimension stats 
avg_PCA_dims = vec(mean(PCA_dims, dims=1))
std_PCA_dims = vec(std(PCA_dims, dims=1))

# Lyapunov spectrum & max Lyapunov exponent stats 
ensbl_avg_Lspecs = reshape(mean(Lspecs, dims=2), (N, length(etas)))
ensbl_std_Lspecs = reshape(std(Lspecs, dims=2), (N, length(etas)))
ensbl_avg_max_LE = vec(mean(max_LEs, dims=1))
ensbl_std_max_LE = vec(std(max_LEs, dims=1))

# Dimensionality and entropy rate stats 
entropy_rates = reshape(mapslices(KS_entropy_rate, Lspecs, dims=1), 
        (trials, length(etas)))
attractor_dims = reshape(mapslices(KY_attractor_dim, Lspecs, dims=1),
        (trials, length(etas)))
avg_entropy_rates = vec(mean(entropy_rates, dims=1))
std_entropy_rates = vec(std(entropy_rates, dims=1))
avg_attractor_dims = vec(mean(attractor_dims, dims=1))
std_attractor_dims = vec(std(attractor_dims, dims=1))


# Save all the statistics that we calculated 
out = Dict{String,Any}()
@pack! out = ensbl_avg_Cs, ensbl_std_Cs, timelags, ensbl_avg_Lspecs, 
        ensbl_std_Lspecs, ensbl_avg_max_LE, ensbl_std_max_LE, avg_PCA_dims,
        std_PCA_dims, avg_entropy_rates, std_entropy_rates, 
        avg_attractor_dims, std_attractor_dims
output_fn = "summary_stats_for_plotting_results1.jld2"
save(joinpath(outdir, output_fn), out)
