using AxisKeys
using LinearAlgebra
using MAT
using NamedDims
using Plots
using Statistics

"""
    load_besa_av(file_name)

Produces a dictionary of arrays (using AxisKeys and NamedDims), based on the input .mat file
of a single average. The input .mat file is produced from the BESA®-MATLAB® interface and
then saved as a separate .mat file containing the average of a single condition.

Each array has the following format: [time(ms), channels]

"""
function load_besa_av(file_name)
    ## Loading the data of all epochs
    data_file = matread(file_name)
    dict_content = collect(keys(data_file))
    if length(dict_content) > 1
        error("Contents of file contain more than a single struct")
        return
    end

    data = data_file[dict_content[1]]
    averaged_data  = data["data"]["amplitudes"]

    averaged_data_array= wrapdims(
        averaged_data,
        time = dropdims(data["data"]["latencies"],dims=1),
        #WARNING: last channel called REF is being indexed out
        channels = Symbol.(data["channellabels"][1:351])
    )

    return averaged_data_array

end

"""
    load_cont_epochs(file_name)

Produces a dictionary of arrays (using AxisKeys and NamedDims), based on the different
conditions classified from the input .mat file. The input .mat file is produced from the
BESA®-MATLAB® interface when the "Epochs around triggers" option is selected.

Each array has the following format: [time(ms), channels, trials]

"""
function load_cont_epochs(file_name)
## Loading the data of all epochs
cont_epochs = matread(file_name)
cont_epochs = cont_epochs["besa_channels"]
all_epochs  = cont_epochs["data"]["amplitudes"]

## Making sure all epochs are the same length
one_latency = cont_epochs["data"]["latencies"][1]
all_latencies = dropdims(cont_epochs["data"]["latencies"],dims=1)
test_all = map(x -> x==one_latency, all_latencies)
if false ∈ test_all
throw("Input Error: all epochs are not of the same size")
return
end


## Extracting all the conditons via their trigger codes

# Proallocating total number of conditions
stacked_conditions = Array{String,1}(undef, length(all_epochs))

# Going through all epochs to determine the type of event at t=0

for epoch = 1:length(all_epochs)
    # determine index at t=0 in each event
    epoch_latencies  = cont_epochs["data"]["event"][epoch]["latency"]

    if length(epoch_latencies) > 1
        t_zero   = findfirst(isequal(0),epoch_latencies[:])
    else
        t_zero   = findfirst(isequal(0),epoch_latencies)
    end

    # Get the condition label/trigger at t=0
    stacked_conditions[epoch]=string(cont_epochs["data"]["event"][epoch]["label"][t_zero][1])
end

# Getting all the different conditions and their counts so they can be nested (under a subject)
# by name. The count is so that we can preallocate the number of trials
unique_conditions = unique(stacked_conditions)
unique_counts = Dict([(condition,count(x->x==condition, stacked_conditions))
for condition in unique_conditions])

# Determine size of a specific epoch
epoch_dims = size(cont_epochs["data"]["amplitudes"][1])

# Making a dictionary to contain all epochs in each condition
# Preallocating based on known details
# DOC: epoch size of all conditions has to be the same
stacked_epochs = Dict()
for (condition,unique_count) in unique_counts
    stacked_epochs[condition] = Array{Float64,3}(undef,epoch_dims[1],epoch_dims[2],unique_count )
end
for condition in unique_conditions


    for epoch = 1:length(unique_counts[condition])
        # Placing the epoch data based on the condition
        condition_index = condition.==stacked_conditions
        relevant_epochs = all_epochs[condition_index]
        for trial = 1:length(relevant_epochs)
            stacked_epochs[condition][:,:,trial] = relevant_epochs[trial]
        end
        # Adding named dimentions. Getting the timing from from the first epoch, i.e
        # assuming that all epoch sizes are the same
         stacked_epochs[condition] = wrapdims(
            stacked_epochs[condition],
            time = dropdims(cont_epochs["data"]["latencies"][1],dims=1),
            channels = dropdims(Symbol.(cont_epochs["channellabels"]),dims=1),
            trials = 1:(size(stacked_epochs["1"],3)),
            )


    end
end

return stacked_epochs

end

"""
    load_BSepochs(subject_path::String)

Reads epoch data (trials) from the outputs generated by Brainstorm and bundles them together.
The general output of Brainstorm consists of .mat data files for each different epoch
(trial`number`), in each different condition (data`condition_label`). An example file name
is `data_15_trial045.mat`. The output is a Dict type with labels for each condition, with
data being present as KeyedArrays.

For more information on Brainstorm:
Tadel F, Baillet S, Mosher JC, Pantazis D, Leahy RM (2011). Brainstorm: A User-Friendly
Application for MEG/EEG Analysis Computational Intelligence and Neuroscience, vol. 2011,
ID 879716. Website: https://neuroimage.usc.edu/brainstorm/

# Examples
```julia-repl
julia> load_BSepochs("reg_soi_ext1_ya10_sss")
Dict{Any,Any} with 19 entries:
  "4"  => [showing 3 of 100 slices]…
  "1"  => [showing 3 of 100 slices]…
  "12" => [showing 3 of 100 slices]…
  "2"  => [showing 3 of 100 slices]…
  ⋮    => ⋮
```
"""
function load_BSepochs(subject_path::String)

    # Switch working directory to the location of the subjects data
    # TODO: Correct to actual file structure as designated by Brainstorm
    path = subject_path

    # Collecting all unique conditions present
    dir_contents = readdir(path)

    # Getting channel Data
    channel_labels = matread(joinpath(path, "channel_vectorview306_acc1.mat"))["Channel"]["Name"]

    # Getting names of bad trials
    bad_trials = matread(joinpath(path, "brainstormstudy.mat"))["BadTrials"]

    # Getting epoch files
    epoch_files = filter(x -> x[1:4] == "data", dir_contents)
    epoch_conditions = Vector{String}(undef, length(epoch_files))
    for (idx, file) in enumerate(epoch_files)

        condition = split(file, "_")[2] # Get condition label
        epoch_conditions[idx] = condition
    end

    unique_conditions = unique(epoch_conditions)

    # Going through the conditions (one by one) and collecting the data
    data = Dict()

    for condition in unique_conditions

        # WARNING: This assumes that there are not more than 99 epoch_conditions
        to_read = filter( x-> split(x, "_")[2] == condition, epoch_files)
        trial_full = matread(joinpath(path, to_read[1])) # Contains all sorts of metadata + signal data
        one_trial = trial_full["F"]' # Reading signal data to get expected dimentions
        total_timepoints = size(one_trial, 1)
        n_channels = size(one_trial, 2)
        # Removing bad trials
        n_trials = length(to_read) - length(intersect(bad_trials, to_read))
        stacked_epochs = Array{Float64}(undef, total_timepoints, n_channels, n_trials)
        clean_trials = setdiff(to_read, bad_trials)

        # Bundling together all the trials in a single array
        for (idx, trial) in enumerate(clean_trials)


            trial_data = matread(joinpath(path, trial))["F"]' # F is the label given to the data
            stacked_epochs[:,:,idx] = trial_data .* 10^15 # Converting from tesla to femto tesla

        end

        # Converting the array into an AxisKeys Array for better indexing
        stacked_epochs = wrapdims(
           stacked_epochs,
           # Time is converted to ms (from seconds) to ms
           # This also means that epoch selection should always be integers in proprocessing
           time = round.(Int, trial_full["Time"][:] * 1000),
           channels = dropdims(Symbol.(channel_labels), dims=1),
           trials = 1:n_trials,
           )

        # Assigning the condition data to the central data dictionary in order to keep all
        # conditions accessible from a single subject
        data[condition] = stacked_epochs


    end

    return data

end
