text:
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# Functions for training each model using its fitted parameters
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MB_Leaky_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "Leaky"

    Param_Path =
    if Part_ind==-1
        Path = string(Param_Path, "MB_Leaky", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MB_Leaky", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end
    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form = "fix"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                  surprise_form=surprise_form,
                                  novelty_form=novelty_form,
                                  update_form=update_form,
                                  update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                        α = 0, Δα = 0,
                        λ_elig_R = 0, λ_elig_N = 0,
                        count_decay=1, Q_R_0 = 0, Q_N_0 = 0,
                        w_MF_to_MB_1=0, w_MF_to_MB_2=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5

    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.N_prior_sweep = Param_Data[Fold,2+8]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.β_R = Param_Data[Fold,2+9]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                  Sub_Set = Sub_Set,
                                                  Epi_Set=Epi_Set,
                                                  Block_Set=Block_Set,
                                                  Param=deepcopy(Param),
                                                  HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MBQ0S_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "MBQ0S", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MBQ0S", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end
    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form = "satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free,
                                   First_Epi_reward_update=true)                           # WARNING WARNING WARNING WARNING

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                        α = 0, Δα = 0,
                        λ_elig_R = 0, λ_elig_N = 0,
                        β_ratio_N2R=0, γ_RL_N=0,
                        count_decay=1, Q_R_0 = 0, Q_N_0 = 0,
                        w_MF_to_MB_1=0, w_MF_to_MB_2=0, w_MF_to_MB_Scaler=1)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Fold_Set = [[9,11,12,13],
               [5,6,7,8],
               [1,2,3,4]]

    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.γ_RL_R = Param_Data[Fold,2+4]
        Param1.N_prior_sweep = Param_Data[Fold,2+5]

        Param2 = deepcopy(Param1)
        Param2.β_R = Param_Data[Fold,2+6]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MBS_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "MBS", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MBS", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end
    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form = "satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                          α = 0, Δα = 0,
                          λ_elig_R = 0, λ_elig_N = 0,
                          count_decay=1, Q_R_0 = 0, Q_N_0 = 0,
                          w_MF_to_MB_1=0, w_MF_to_MB_2=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.N_prior_sweep = Param_Data[Fold,2+8]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.β_R = Param_Data[Fold,2+9]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MBSExp_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "MBSExp", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MBSExp", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end
    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form = "satur"

    exploration_form = "MB_Curiosity"       #### WARNING ####

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                surprise_form=surprise_form,
                                novelty_form=novelty_form,
                                exploration_form=exploration_form,
                                update_form=update_form,
                                update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                          α = 0, Δα = 0,
                          λ_elig_R = 0, λ_elig_N = 0,
                          count_decay=1, Q_R_0 = 0, Q_N_0 = 0,
                          w_MF_to_MB_1=0, w_MF_to_MB_2=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.N_prior_sweep = Param_Data[Fold,2+8]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.β_R = Param_Data[Fold,2+9]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MFN_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "Leaky"

    if Part_ind==-1
        Path = string(Param_Path, "MFN", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MFN", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="fix"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                          count_decay=1, Q_R_0=0, Δα=0,
                          p_change=0, ϵ=1, N_prior_sweep=0,
                          w_MF_to_MB_1 = 1, w_MF_to_MB_2 = 1)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.β_R = Param_Data[Fold,2+1]
        Param1.β_ratio_N2R = Param_Data[Fold,2+2]
        Param1.γ_RL_R = Param_Data[Fold,2+4]
        Param1.γ_RL_N = Param_Data[Fold,2+5]
        Param1.λ_elig_R = Param_Data[Fold,2+6]
        Param1.λ_elig_N = Param_Data[Fold,2+7]
        Param1.α = Param_Data[Fold,2+8]
        Param1.Q_N_0 = Param_Data[Fold,2+9]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+3]
        Param2.β_R = Param_Data[Fold,2+10]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MFQ0_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "Leaky"

    if Part_ind==-1
        Path = string(Param_Path, "MFQ0", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MFQ0", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="fix"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                        count_decay=1, Q_N_0=0, Δα=0,
                        β_ratio_N2R=0, γ_RL_N=0, λ_elig_N=0,
                        p_change=0, ϵ=1, N_prior_sweep=0,
                        w_MF_to_MB_1 = 1, w_MF_to_MB_2 = 1)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.β_R = Param_Data[Fold,2+1]
        Param1.γ_RL_R = Param_Data[Fold,2+2]
        Param1.λ_elig_R = Param_Data[Fold,2+3]
        Param1.α = Param_Data[Fold,2+4]
        Param1.Q_R_0 = Param_Data[Fold,2+5]

        Param2 = deepcopy(Param1)
        Param2.β_R = Param_Data[Fold,2+6]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MFNS_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "MFNS", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MFNS", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                        count_decay=1, Q_R_0 = 0,
                        w_MF_to_MB_1=1, w_MF_to_MB_2=1, N_prior_sweep=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.β_R = Param_Data[Fold,2+13]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_MFSExp_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "MFSExp", "_BestParams.csv" )
    else
        Path = string(Param_Path, "MFSExp", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    exploration_form = "MB_Curiosity"       #### WARNING ####

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                surprise_form=surprise_form,
                                novelty_form=novelty_form,
                                exploration_form=exploration_form,
                                update_form=update_form,
                                update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                        count_decay=1, Q_R_0 = 0,
                        w_MF_to_MB_1=1, w_MF_to_MB_2=1, N_prior_sweep=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.β_R = Param_Data[Fold,2+13]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_HybN_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "Leaky"

    if Part_ind==-1
        Path = string(Param_Path, "HybN", "_BestParams.csv" )
    else
        Path = string(Param_Path, "HybN", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="fix"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                               surprise_form=surprise_form,
                               novelty_form=novelty_form,
                               update_form=update_form,
                               update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                                count_decay=1, Q_R_0 = 0, Δα=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Q_N_0 = Param_Data[Fold,2+11]
        Param1.N_prior_sweep = Param_Data[Fold,2+12]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+13]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+14]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+16]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+15]
        Param2.β_R = Param_Data[Fold,2+17]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_HybQ0S_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "HybQ0S", "_BestParams.csv" )
    else
        Path = string(Param_Path, "HybQ0S", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                               surprise_form=surprise_form,
                               novelty_form=novelty_form,
                               update_form=update_form,
                               update_form_model_free=update_form_model_free,
                               First_Epi_reward_update=true)                           # WARNING WARNING WARNING WARNING

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                        count_decay=1, Q_N_0=0,
                        β_ratio_N2R=0, γ_RL_N=0, λ_elig_N=0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.γ_RL_R = Param_Data[Fold,2+4]
        Param1.λ_elig_R = Param_Data[Fold,2+5]
        Param1.α = Param_Data[Fold,2+6]
        Param1.Δα = Param_Data[Fold,2+7]
        Param1.Q_R_0 = Param_Data[Fold,2+8]
        Param1.N_prior_sweep = Param_Data[Fold,2+9]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+10]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+11]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+13]

        Param2 = deepcopy(Param1)
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+12]
        Param2.β_R = Param_Data[Fold,2+14]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_HybSExp_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "HybSExp", "_BestParams.csv" )
    else
        Path = string(Param_Path, "HybSExp", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    exploration_form = "MB_Curiosity"       #### WARNING ####

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                surprise_form=surprise_form,
                                novelty_form=novelty_form,
                                exploration_form=exploration_form,
                                update_form=update_form,
                                update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                                  count_decay=1, Q_R_0 = 0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]
        Param1.N_prior_sweep = Param_Data[Fold,2+13]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+14]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+15]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+17]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+16]
        Param2.β_R = Param_Data[Fold,2+18]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_HybSTrapDet_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "HybSTrapDet", "_BestParams.csv" )
    else
        Path = string(Param_Path, "HybSTrapDet", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form = "trap-det"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                           surprise_form=surprise_form,
                                           novelty_form=novelty_form,
                                           update_form=update_form,
                                           update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                                  count_decay=1, Q_R_0 = 0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]
        Param1.N_prior_sweep = Param_Data[Fold,2+13]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+14]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+15]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+17]
        Param1.C_thresh = Param_Data[Fold,2+19]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+16]
        Param2.β_R = Param_Data[Fold,2+18]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_HybSTrapDet2_CV_train(;Part_ind=-1,
                                 BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                 Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                                 Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "HybSTrapDet2", "_BestParams.csv" )
    else
        Path = string(Param_Path, "HybSTrapDet2", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form = "trap-det-2"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                           surprise_form=surprise_form,
                                           novelty_form=novelty_form,
                                           update_form=update_form,
                                           update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                                  count_decay=1, Q_R_0 = 0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]
        Param1.N_prior_sweep = Param_Data[Fold,2+13]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+14]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+15]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+17]
        Param1.C_thresh = Param_Data[Fold,2+19]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+16]
        Param2.β_R = Param_Data[Fold,2+18]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_SurNoR_NovInExploit_CV_train(;Part_ind=-1,
                               BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                               Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                               Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]])

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "SurNoR_NovInExploit", "_BestParams.csv" )
    else
        Path = string(Param_Path, "SurNoR_NovInExploit", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                                  count_decay=1, Q_R_0 = 0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]
        Param1.N_prior_sweep = Param_Data[Fold,2+13]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+14]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+15]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+17]
        Param1.β_ratio_N2R_exploit = Param_Data[Fold,2+19]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+16]
        Param2.β_R = Param_Data[Fold,2+18]

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_SurNoR_CV_train(;Part_ind=-1,
                               BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                               Param_Path = "./data/Fitted_parameters/CV_3Fold/",
                               Fold_Set = [[9,11,12,13],[5,6,7,8],[1,2,3,4]],
                               w_MF_to_MB = -1)

    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    if Part_ind==-1
        Path = string(Param_Path, "SurNoR", "_BestParams.csv" )
    else
        Path = string(Param_Path, "SurNoR", "_BestParams","_Part", string(Part_ind), ".csv" )
    end
    Param_Data = CSV.read(Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                                  count_decay=1, Q_R_0 = 0)

    Param2 = deepcopy(Param1)

    Param = [Param1, Param2]

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5
    N_sub = 12
    N_blo = length(Block_Set)
    N_epi = length(Epi_Set)
    Agents_Tot = Array{Str_Agent,3}(undef,N_sub,N_blo,N_epi)
    for Fold = 1:3
        Sub_Set = Fold_Set[Fold]

        Param1.p_change = Param_Data[Fold,2+1]
        Param1.ϵ = Param_Data[Fold,2+2]
        Param1.β_R = Param_Data[Fold,2+3]
        Param1.β_ratio_N2R = Param_Data[Fold,2+4]
        Param1.γ_RL_R = Param_Data[Fold,2+6]
        Param1.γ_RL_N = Param_Data[Fold,2+7]
        Param1.λ_elig_R = Param_Data[Fold,2+8]
        Param1.λ_elig_N = Param_Data[Fold,2+9]
        Param1.α = Param_Data[Fold,2+10]
        Param1.Δα = Param_Data[Fold,2+11]
        Param1.Q_N_0 = Param_Data[Fold,2+12]
        Param1.N_prior_sweep = Param_Data[Fold,2+13]
        Param1.w_MF_to_MB_2 = Param_Data[Fold,2+14]
        Param1.w_MF_to_MB_1 = Param_Data[Fold,2+15]
        Param1.w_MF_to_MB_Scaler = Param_Data[Fold,2+17]

        Param2 = deepcopy(Param1)
        Param2.β_ratio_N2R = Param_Data[Fold,2+5]
        Param2.w_MF_to_MB_1 = Param_Data[Fold,2+16]
        Param2.β_R = Param_Data[Fold,2+18]

        if w_MF_to_MB != -1
            Param1.w_MF_to_MB_1 = w_MF_to_MB
            Param1.w_MF_to_MB_2 = w_MF_to_MB
            Param2.w_MF_to_MB_1 = w_MF_to_MB
            Param2.w_MF_to_MB_2 = w_MF_to_MB
        end

        Param = [Param1, Param2]

        (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                                   Sub_Set = Sub_Set,
                                                   Epi_Set=Epi_Set,
                                                   Block_Set=Block_Set,
                                                   Param=deepcopy(Param),
                                                   HyperParam=HyperParam)

        Agents_Tot[ ((Fold-1)*4) .+ (1:4) ,:,:] .= Agents
    end

    return deepcopy(Agents_Tot)
end

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
function Func_SurNoR_Overall_train_opt_ind_i(i;BehavData = CSV.read("./data/BehavData_SwitchState.csv"),
                                               Sub_Set = [1,2,3,4,5,6,7,8,9,11,12,13],
                                               Param_Path = "./data/Fitted_parameters/Overall/",
                                               w_MF_to_MB = -1)
    # --------------------------------------------------------------------------
    # Model Based Algorithms
    update_form = "VarSMiLe"

    Param_Path = string(Param_Path, "OverAll_SurNoR_BestParams_diff5.csv" )
    Param_Data = CSV.read(Param_Path)

    if update_form == "SMiLe"
        surprise_form="CC"
    else
        surprise_form="BF"
    end

    novelty_form="log-count"
    update_form_model_free = "Q_learning"
    gamma_form="satur"

    # --------------------------------------------------------------------------
    # Defining Hyper Parameters
    HyperParam = Str_HyperParam(; gamma_form=gamma_form,
                                   surprise_form=surprise_form,
                                   novelty_form=novelty_form,
                                   update_form=update_form,
                                   update_form_model_free=update_form_model_free)

    # --------------------------------------------------------------------------
    # Fitting the models
    Block_Set = 1:2
    Epi_Set = 1:5

    # --------------------------------------------------------------------------
    # Defining Parameters
    Param1 = Str_Param(; num_act=4, num_obs=11,
                         count_decay=1, Q_R_0 = 0)

    Param1.p_change = Param_Data[i,1+1]
    Param1.ϵ = Param_Data[i,1+2]
    Param1.β_R = Param_Data[i,1+3]
    Param1.β_ratio_N2R = Param_Data[i,1+4]
    Param1.γ_RL_R = Param_Data[i,1+6]
    Param1.γ_RL_N = Param_Data[i,1+7]
    Param1.λ_elig_R = Param_Data[i,1+8]
    Param1.λ_elig_N = Param_Data[i,1+9]
    Param1.α = Param_Data[i,1+10]
    Param1.Δα = Param_Data[i,1+11]
    Param1.Q_N_0 = Param_Data[i,1+12]
    Param1.N_prior_sweep = Param_Data[i,1+13]
    Param1.w_MF_to_MB_2 = Param_Data[i,1+14]
    Param1.w_MF_to_MB_1 = Param_Data[i,1+15]
    Param1.w_MF_to_MB_Scaler = Param_Data[i,1+17]

    Param2 = deepcopy(Param1)
    Param2.β_ratio_N2R = Param_Data[i,1+5]
    Param2.w_MF_to_MB_1 = Param_Data[i,1+16]
    Param2.β_R = Param_Data[i,1+18]

    Param = [Param1, Param2]

    if w_MF_to_MB != -1
        Param1.w_MF_to_MB_1 = w_MF_to_MB
        Param1.w_MF_to_MB_2 = w_MF_to_MB
        Param2.w_MF_to_MB_1 = w_MF_to_MB
        Param2.w_MF_to_MB_2 = w_MF_to_MB
    end

    (Agents, LogL) = Func_train_set_of_Agents(BehavData;
                                               Sub_Set = Sub_Set,
                                               Epi_Set=Epi_Set,
                                               Block_Set=Block_Set,
                                               Param=deepcopy(Param),
                                               HyperParam=HyperParam)

    Agents_Tot = deepcopy(Agents)

    return deepcopy(Agents_Tot)
end


text:
module TestBeforeAfterHooks

using Test
using TestFunctionRunner.Internal: runtests
using ..XFails: XFailTestSet

module BeforeAfterHooks
using Test
const RUN_BEFORE = Ref(true)
const VALUE = Ref(0)
before_test_module() = RUN_BEFORE[] && (VALUE[] = 1)
after_test_module() = VALUE[] = 0
test() = @test VALUE[] == 1
end  # module BeforeAfterHooks

function test_run_before()
    @test BeforeAfterHooks.VALUE[] == 0
    runtests(BeforeAfterHooks)
    @test BeforeAfterHooks.VALUE[] == 0
end

function test_no_run_before()
    BeforeAfterHooks.RUN_BEFORE[] = false
    try
        @testset XFailTestSet "xfail" begin
            runtests(BeforeAfterHooks)
        end
    finally
        BeforeAfterHooks.RUN_BEFORE[] = true
    end
end

end  # module

ext: jl


text:
module ProtoGrad

_training_mode = false

is_training_mode() = _training_mode

function within_training_mode(f)
    global _training_mode
    _training_mode = true
    ret = f()
    _training_mode = false
    return ret
end

include("itertools.jl")
include("losses.jl")
include("model.jl")
include("layers.jl")
include("objectives.jl")
include("gradient.jl")
include("optimizers.jl")

end # module


text:
using GoogleCodeSearch
using Test
using HTTP
using JSON
using Sockets

const test_data = [
(path="path1", file="file1", content="""
path1, file1, line1
path1, file1, line2
"""),
(path="path1", file="file2", content="""
path1, file2, line1
path1, file2, line2
"""),
(path="path2", file="file1", content="""
path2, file1, line1
path2, file1, line2
"""),
(path="path2", file="file2", content="""
path2, file2, line1
path2, file2, line2
"""),
(path="path22", file="file1", content="""
path22, file1, line1
path22, file1, line2
"""),
(path="path22", file="file2", content="""
path22, file2, line1
path22, file2, line2
""")
]

# split index on the last alphabet of path component of test_data
test_index_resolver(ctx::Ctx, inpath::String) = joinpath(ctx.store, "index" * rsplit(inpath, '/'; limit=2)[end][end])
test_default_resolver(ctx::Ctx, inpath::String) = joinpath(ctx.store, "index")

function test_noindices(ctx::Ctx)
    @test isempty(paths_indexed(ctx))
    @test isempty(indices(ctx))
    @test isempty(search(ctx, ".*"; ignorecase=true))

    # clear_indices should not error out if there are no indices
    @test clear_indices(ctx) === nothing
    nothing
end

function test_indexing(ctx::Ctx, datadir::String)
    paths = Set(joinpath(datadir, item.path) for item in test_data)
    # paths must be in sorted order during indexing, otherwise indexer will remove a longer path when it encounters a path that is a substring of that
    for path in sort(collect(paths))
        @debug("indexing $path")
        @test index(ctx, path)
        @test !isempty(indices(ctx))
        @test any(occursin(x, path) for x in paths_indexed(ctx))
    end
    all_paths = paths_indexed(ctx)
    @test !("/nosuchpath" in all_paths)
    # test that `path2` subsumes `path22`
    @test !any(occursin("path22", x) for x in all_paths)
    nothing
end

function test_search(ctx::Ctx, datadir::String)
    # test case sensitiveness
    @test isempty(search(ctx, "File"))
    @test !isempty(search(ctx, "File"; ignorecase=true))

    # test pathfilter
    for ignorecase in (true, false)
        res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path1/.*")
        if ignorecase
            @test length(res) == 2
            for item in res
                @test item.line == 1
                @test endswith(strip(item.text), "line1")
                @test occursin("path1", item.file)
            end
        else
            @test isempty(res)
        end
    end

    # test split indices
    for ignorecase in (true, false)
        res = search(ctx, "Line2"; ignorecase=ignorecase)
        if ignorecase
            @test length(res) == length(test_data)
            for item in res
                @test item.line == 2
                @test endswith(strip(item.text), "line2")
                @test occursin("path", item.file)
            end
        else
            @test isempty(res)
        end
    end
    nothing
end

function create_test_data(datadir::String)
    for item in test_data
        dir = joinpath(datadir, item.path)
        mkpath(dir)
        open(joinpath(dir, item.file), "w") do f
            println(f, item.content)
        end
    end
end

function test_apis(testdir, resolver)
    storedir = joinpath(testdir, "store")
    datadir = joinpath(testdir, "data")
    create_test_data(datadir)
    mkpath(storedir)

    ctx = Ctx(store=storedir, resolver=resolver)
    test_noindices(ctx)
    test_indexing(ctx, datadir)
    test_search(ctx, datadir)

    @test clear_indices(ctx) === nothing
    test_noindices(ctx)

    nothing
end

function test_http_service(testdir)
    storedir = joinpath(testdir, "store")
    datadir = joinpath(testdir, "data")
    create_test_data(datadir)
    mkpath(storedir)

    ctx = Ctx(store=storedir)
    srvr = @async run_http(ctx)
    wait_for_httpsrvr()

    paths = Set(joinpath(datadir, item.path) for item in test_data)
    # paths must be in sorted order during indexing, otherwise indexer will remove a longer path when it encounters a path that is a substring of that
    for path in sort(collect(paths))
        resp = HTTP.request("POST", "http://127.0.0.1:5555/index"; body="{\"path\":\"$path\"}")
        @test resp.status == 200
        json_resp = JSON.parse(String(resp.body))
        @test json_resp["success"]
    end

    resp = HTTP.request("POST", "http://127.0.0.1:5555/search"; body="{\"pattern\":\"line1\"}")
    @test resp.status == 200
    json_resp = JSON.parse(String(resp.body))
    @test json_resp["success"]
    data = json_resp["data"]
    @test length(data) == 6
    for item in data
        @test item["line"] == 1
        @test endswith(strip(item["text"]), "line1")
    end
end

function wait_for_httpsrvr()
    while true
        try
            sock = connect("127.0.0.1", 5555)
            close(sock)
            return
        catch
            @debug("waiting for httpserver to come up at port 5555...")
            sleep(5)
        end
    end
end

@testset "GoogleCodeSearch" begin
    @testset "APIs" begin
        for (resolver_name,resolver) in ("split_index"=>test_index_resolver, "single_index"=>test_default_resolver)
            @debug("testing with $resolver_name")
            mktempdir() do testdir
                test_apis(testdir, resolver)
            end
        end
    end

    @testset "Index read write" begin
        mktempdir() do testdir
            idx = GoogleCodeSearch.Index()
            input_file = joinpath(@__DIR__, "simple.index")
            output_file = joinpath(testdir, "simple.index")
            open(input_file, "r") do inp
                read(inp, idx)
            end
            @test length(idx.paths.entries) == 0
            @test length(idx.names.entries) == 6
            @test length(idx.postings) == 12
            @test length(idx.nameindex.entries) == 7
            @test length(idx.postingindex.entries) == 12
            @test idx.offsets.path_list == 0x00000010
            @test idx.offsets.name_list == 0x00000011
            @test idx.offsets.posting_list == 0x00000037
            @test idx.offsets.name_index == 0x00000075
            @test idx.offsets.posting_list_index == 0x00000091
            open(output_file, "w") do out
                write(out, idx)
            end
            @test read(input_file) == read(output_file)
        end
    end

    @testset "Index prune" begin
        mktempdir() do testdir
            # create index and test all indexed as expected
            storedir = joinpath(testdir, "store")
            datadir = joinpath(testdir, "data")
            create_test_data(datadir)
            mkpath(storedir)

            ctx = Ctx(store=storedir)
            test_noindices(ctx)
            test_indexing(ctx, datadir)
            test_search(ctx, datadir)

            idx = GoogleCodeSearch.Index()
            input_file = joinpath(storedir, "index")
            open(input_file, "r") do inp
                read(inp, idx)
            end

            # recreating indices without any changes should leave things unchanged
            GoogleCodeSearch.recreate_indices!(idx)
            open(input_file, "w") do out
                write(out, idx)
            end
            ctx = Ctx(store=storedir)
            test_search(ctx, datadir)
            @test length(paths_indexed(ctx)) == 2
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path1/.*")
                if ignorecase
                    @test !isempty(res)
                else
                    @test isempty(res)
                end
            end

            # prune one of the paths
            @test length(idx.paths.entries) == 2
            @test length(idx.names.entries) == 6
            @test length(idx.postings) == 29
            @test length(idx.postings[1].deltas) == 7

            GoogleCodeSearch.prune_paths!(idx, joinpath(datadir, "path1"))

            @test length(idx.paths.entries) == 1
            @test length(idx.names.entries) == 4
            @test length(idx.postings) == 27
            @test length(idx.postings[1].deltas) == 5

            open(input_file, "w") do out
                write(out, idx)
            end
            ctx = Ctx(store=storedir)
            @test length(paths_indexed(ctx)) == 1

            # test pathfilter
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path1/.*")
                @test isempty(res)
            end
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path2/.*")
                if ignorecase
                    @test length(res) == 2
                else
                    @test isempty(res)
                end
            end

            # prune a single file
            GoogleCodeSearch.prune_files!(idx, joinpath(datadir, "path2", "file1"))
            @test length(idx.paths.entries) == 1
            @test length(idx.names.entries) == 3
            @test length(idx.postings) == 27
            @test length(idx.postings[1].deltas) == 4

            open(input_file, "w") do out
                write(out, idx)
            end
            ctx = Ctx(store=storedir)
            @test length(paths_indexed(ctx)) == 1

            # test pathfilter
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path1/.*")
                @test isempty(res)
            end
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path2/.*")
                if ignorecase
                    @test length(res) == 1
                else
                    @test isempty(res)
                end
            end
        end

        mktempdir() do testdir
            # create index and test all indexed as expected
            storedir = joinpath(testdir, "store")
            datadir = joinpath(testdir, "data")
            create_test_data(datadir)
            mkpath(storedir)

            ctx = Ctx(store=storedir)
            test_noindices(ctx)
            test_indexing(ctx, datadir)
            test_search(ctx, datadir)

            GoogleCodeSearch.prune_paths!(ctx, joinpath(datadir, "path1"))
            @test length(paths_indexed(ctx)) == 1

            # test pathfilter
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path1/.*")
                @test isempty(res)
            end
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path2/.*")
                if ignorecase
                    @test length(res) == 2
                else
                    @test isempty(res)
                end
            end

            # prune a single file
            GoogleCodeSearch.prune_files!(ctx, joinpath(datadir, "path2", "file1"))
            @test length(paths_indexed(ctx)) == 1

            # test pathfilter
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path1/.*")
                @test isempty(res)
            end
            for ignorecase in (true, false)
                res = search(ctx, "Line1"; ignorecase=ignorecase, pathfilter=".*/path2/.*")
                if ignorecase
                    @test length(res) == 1
                else
                    @test isempty(res)
                end
            end
        end
    end

    @testset "HTTP service" begin
        mktempdir() do testdir
            test_http_service(testdir)
        end
    end
end


text:
using Pkg; Pkg.activate("./")
using nevanlinna_ac
using DelimitedFiles
using Printf

println("2021-10-14: scan g and beta")
η = 0.001
beta = [10.0, 20.0, 30.0, 40.0]
gamma = [0.8, 0.9, 1.0,  1.1, 1.2]
omega = [i for i in range(-5,5,length=1000)]

for j = 1:length(gamma), b = 1:length(beta)
    g = gamma[j]; β = beta[b]
    path1 = @sprintf "./data/gt/gt8_g_%.1f_b_%i.txt" g β
    path2 = @sprintf "./data/gt/gt28_g_%.1f_b_%i.txt" g β
    opath1 = @sprintf "./data/gt/a_g_%.1f_b_%i.txt" g β
    opath2 = @sprintf "./data/gt/a5_g_%.1f_b_%i.txt" g β

    g8 = make_input(path1); g8s = make_input(path1, 5)
    g28 = make_input(path2); g28s = make_input(path2, 5)

    chi8 = [spectrum_density(i,η,g8) for i in omega]
    chi28 = [spectrum_density(i,η,g28) for i in omega]

    chi8s = [spectrum_density(i,η,g8s) for i in omega]
    chi28s = [spectrum_density(i,η,g28s) for i in omega]

    open(opath1,"w") do file
        for i = 1:length(omega)
            writedlm(file,[omega[i] chi8[i] chi28[i]])
        end
    end

    open(opath2,"w") do file
        for i = 1:length(omega)
            writedlm(file,[omega[i] chi8s[i] chi28s[i]])
        end
    end
end


text:
module EarthOrientation

using Dates: datetime2julian, julian2datetime, Date, DateTime, today, days
using DelimitedFiles
using LeapSeconds
using OptionalData
using RemoteFiles

export EOParams, EOP_DATA, interpolate
export polarmotion, getxp, getxp_err, getyp, getyp_err
export precession_nutation80, getdψ, getdψ_err, getdϵ, getdϵ_err
export precession_nutation00, getdx, getdx_err, getdy, getdy_err
export getΔUT1, getΔUT1_err, getlod, getlod_err

function __init__()
    if isfile(data)
        eop = EOParams(paths(data, :iau1980, :iau2000)...)
        push!(EOP_DATA, eop)
    end
end

@RemoteFileSet data "IERS Data" begin
    iau1980 = @RemoteFile(
        "https://datacenter.iers.org/eop/-/somos/5Rgv/latestXL/7/finals.all/csv",
        file="finals.csv",
        updates=:thursdays,
    )
    iau2000 = @RemoteFile(
        "https://datacenter.iers.org/eop/-/somos/6Rgv/latestXL/9/finals2000A.all/csv",
        file="finals2000A.csv",
        updates=:thursdays,
    )
end

include("akima.jl")

const MJD_EPOCH = 2400000.5
date_from_mjd(mjd) = Date(julian2datetime(mjd + MJD_EPOCH))

struct OutOfRangeError <: Base.Exception
    mjd::Float64
    when::String
end
Base.showerror(io::IO, err::OutOfRangeError) = print(io, "No data available ",
    err.when, " ", date_from_mjd(err.mjd), ".")
warn_range(mjd, when) = @warn("No data available $when " *
    "$(date_from_mjd(mjd)). The last valid value will be returned.")

"""
    update()

Download weekly EOP data from the IERS servers if newer files are available or
no data has been downloaded previously.
"""
function update()
    download(data)
    push!(EOP_DATA, paths(data, :iau1980, :iau2000)...)
    nothing
end

"""
    getdate(data)

Determine the creation date of an IERS table by finding the last entry which is marked as "final".
"""
function getdate(data)
    idx = findlast(data[:,5] .== "final")
    Date(data[idx,2], data[idx,3], data[idx,4])
end

mutable struct EOParams
    "Creation date of the contained IERS tables."
    date::Date
    "All modified Julian dates covered by this table."
    mjd::Vector{Float64}
    "Contains the last available data point for every table."
    lastdate::Dict{Symbol, Float64}
    "Difference between UT1 and TAI."
    UT1_TAI::Akima
    "Set of leap second dates"
    leaps::Set{Int}
    "North pole x-coordinate."
    xp::Akima
    "Error in North pole x-coordinate."
    xp_err::Akima
    "North pole y-coordinate."
    yp::Akima
    "Error in North pole y-coordinate."
    yp_err::Akima
    "Difference between UT1 and UTC."
    ΔUT1::Akima
    "Error in difference between UT1 and UTC."
    ΔUT1_err::Akima
    "Excess length of day."
    lod::Akima
    "Error in excess length of day."
    lod_err::Akima
    "Ecliptic nutation correction."
    dψ::Akima
    "Error in ecliptic nutation correction."
    dψ_err::Akima
    "Ecliptic obliquity correction"
    dϵ::Akima
    "Error in ecliptic obliquity correction"
    dϵ_err::Akima
    "Celestial pole x-coordinate correction."
    dx::Akima
    "Error in celestial pole x-coordinate correction."
    dx_err::Akima
    "Celestial pole y-coordinate correction."
    dy::Akima
    "Error in celestial pole y-coordinate correction."
    dy_err::Akima

    EOParams(date, mjd) = new(date, mjd, Dict{Symbol,Float64}())
end

@OptionalData EOP_DATA EOParams "Call 'EarthOrientation.update()' to load it."

const columns = Dict(
    :xp => 6,
    :xp_err => 7,
    :yp => 8,
    :yp_err => 9,
    :ΔUT1 => 11,
    :ΔUT1_err => 12,
    :lod => 13,
    :lod_err => 14,
    :dψ => 16,
    :dψ_err => 17,
    :dϵ => 18,
    :dϵ_err => 19,
    :dx => 20,
    :dx_err => 21,
    :dy => 22,
    :dy_err => 23,
)

"""
    EOParams(iau1980file::String, iau2000file::String)

Parse IERS data files into a `EOParams` object.
`iau1980file` and `iau2000file` are the paths to a 'finals.all' and a
'finals2000A.all' CSV file, respectively.
"""
function EOParams(iau1980file::String, iau2000file::String)
    data80, header80 = readdlm(iau1980file, ';', header=true)
    data00, header00 = readdlm(iau2000file, ';', header=true)
    date = getdate(data80)
    mjd = Vector{Float64}(data80[:,1])
    eop = EOParams(date, mjd)
    for field in fieldnames(EOParams)[6:end]
        col = columns[field]
        data = col < 20 ? data80 : data00
        row = findfirst(isempty.(data[:,col])) - 1
        merge!(eop.lastdate, Dict(field => mjd[row]))
        setfield!(eop, field,
                  Akima(mjd[1:row], Vector{Float64}(data[1:row,col])))
    end
    Δut1_tai = similar(eop.ΔUT1.y)
    for (i, (mjd, Δut1_utc)) in enumerate(zip(eop.ΔUT1.x, eop.ΔUT1.y))
        jd = mjd + MJD_EPOCH
        Δat = offset_tai_utc(jd)
        Δut1_tai[i] = Δut1_utc - Δat
    end
    eop.UT1_TAI = Akima(copy(eop.ΔUT1.x), Δut1_tai)
    eop.leaps = Set(LeapSeconds.LS_EPOCHS)
    return eop
end

Base.show(io::IO, eop::EOParams) = print(io, "EOParams($(eop.date))")

@inline function interpolate(eop::EOParams, field::Symbol, jd::Float64; outside_range=:warn)
    mjd = jd - MJD_EPOCH
    before = mjd < eop.mjd[1]
    after = mjd > eop.lastdate[field]
    if after || before
        when = after ? "after" : "before"
        lim = after ? eop.lastdate[field] : eop.mjd[1]
        if outside_range == :nothing
            # continue
        elseif outside_range == :error
            throw(OutOfRangeError(lim, when))
        elseif outside_range == :warn
            warn_range(lim, when)
        else
            throw(ArgumentError("Unknown value for keyword argument 'outside_range'."))
        end
    end

    if field == :ΔUT1
        Δat = offset_tai_utc(jd)
        imjd = trunc(Int, mjd)
        if imjd == mjd && imjd in eop.leaps
            Δat -= 1.0
        end
        return interpolate(eop.UT1_TAI, mjd) + Δat
    end

    interpolate(getfield(eop, field), mjd)
end

function interpolate(eop::EOParams, field::Symbol, dt::DateTime; args...)
    interpolate(eop, field, datetime2julian(dt); args...)
end

################
# Polar motion #
################

"""
    getxp(date; outside_range=:warn)

Get the x-coordinate of Earth's north pole w.r.t. the CIO for a certain `date` in arcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getxp(eop, date; args...) = interpolate(eop, :xp, date; args...)
getxp(date; args...) = getxp(get(EOP_DATA), date; args...)

"""
    getxp_err(date; outside_range=:warn)

Get the error for the x-coordinate of Earth's north pole w.r.t. the CIO
for a certain `date` in arcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getxp_err(eop, date; args...) = interpolate(eop, :xp_err, date; args...)
getxp_err(date; args...) = getxp_err(get(EOP_DATA), date; args...)

"""
    getyp(date; outside_range=:warn)

Get the y-coordinate of Earth's north pole w.r.t. the CIO for a certain `date` in arcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getyp(eop, date; args...) = interpolate(eop, :yp, date; args...)
getyp(date; args...) = getyp(get(EOP_DATA), date; args...)

"""
    getyp_err(date; outside_range=:warn)

Get the error for the y-coordinate of Earth's north pole w.r.t. the CIO
for a certain `date` in arcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getyp_err(eop, date; args...) = interpolate(eop, :yp_err, date; args...)
getyp_err(date; args...) = getyp_err(get(EOP_DATA), date; args...)

"""
    polarmotion(date; outside_range=:warn)

Get the coordinates of Earth's north pole w.r.t. the CIO for a certain `date` in arcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
polarmotion(eop, date; args...) = (getxp(eop, date; args...),
    getyp(eop, date; outside_range=:warn, args...))
polarmotion(date; args...) = (getxp(date; args...),
    getyp(date; outside_range=:warn, args...))

################
# ΔUT1 and LOD #
################

"""
    getΔUT1(date; outside_range=:warn)

Get the difference between UTC and UT1 for a certain `date` in seconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getΔUT1(eop, date; args...) = interpolate(eop, :ΔUT1, date; args...)
getΔUT1(date; args...) = getΔUT1(get(EOP_DATA), date; args...)

"""
    getΔUT1_err(date; outside_range=:warn)

Get the error in the difference between UTC and UT1 for a certain `date` in seconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getΔUT1_err(eop, date; args...) = interpolate(eop, :ΔUT1_err, date; args...)
getΔUT1_err(date; args...) = getΔUT1_err(get(EOP_DATA), date; args...)

"""
    getlod(date; outside_range=:warn)

Get the excess length of day for a certain `date` in milliseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getlod(eop, date; args...) = interpolate(eop, :lod, date; args...)
getlod(date; args...) = getlod(get(EOP_DATA), date; args...)

"""
    getlod_err(date; outside_range=:warn)

Get the error in the excess length of day for a certain `date` in milliseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getlod_err(eop, date; args...) = interpolate(eop, :lod_err, date; args...)
getlod_err(date; args...) = getlod_err(get(EOP_DATA), date; args...)

#######################################
# IAU 1980 precession/nutation theory #
#######################################

"""
    getdψ(date; outside_range=:warn)

Get the ecliptic nutation correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdψ(eop, date; args...) = interpolate(eop, :dψ, date; args...)
getdψ(date; args...) = getdψ(get(EOP_DATA), date; args...)

"""
    getdψ_err(date; outside_range=:warn)

Get the error in the ecliptic nutation correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdψ_err(eop, date; args...) = interpolate(eop, :dψ_err, date; args...)
getdψ_err(date; args...) = getdψ_err(get(EOP_DATA), date; args...)

"""
    getdϵ(date; outside_range=:warn)

Get the ecliptic obliquity correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdϵ(eop, date; args...) = interpolate(eop, :dϵ, date; args...)
getdϵ(date; args...) = getdϵ(get(EOP_DATA), date; args...)

"""
    getdϵ_err(date; outside_range=:warn)

Get the error in the ecliptic obliquity correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdϵ_err(eop, date; args...) = interpolate(eop, :dϵ_err, date; args...)
getdϵ_err(date; args...) = getdϵ_err(get(EOP_DATA), date; args...)

"""
    precession_nutation80(date; outside_range=:warn)

Get the ecliptic corrections for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
precession_nutation80(eop, date; args...) = (getdψ(eop, date; args...),
    getdϵ(eop, date; outside_range=:warn, args...))
precession_nutation80(date; args...) = (getdψ(date; args...),
    getdϵ(date; outside_range=:warn, args...))

############################################
# IAU 2000/2006 precession/nutation theory #
############################################

"""
    getdx(date; outside_range=:warn)

Get the celestial pole x-coordinate correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdx(eop, date; args...) = interpolate(eop, :dx, date; args...)
getdx(date; args...) = getdx(get(EOP_DATA), date; args...)

"""
    getdx_err(date; outside_range=:warn)

Get the error in celestial pole x-coordinate correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdx_err(eop, date; args...) = interpolate(eop, :dx_err, date; args...)
getdx_err(date; args...) = getdx_err(get(EOP_DATA), date; args...)

"""
    getdy(date; outside_range=:warn)

Get the celestial pole y-coordinate correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdy(eop, date; args...) = interpolate(eop, :dy, date; args...)
getdy(date; args...) = getdy(get(EOP_DATA), date; args...)

"""
    getdy_err(date; outside_range=:warn)

Get the error in celestial pole y-coordinate correction for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
getdy_err(eop, date; args...) = interpolate(eop, :dy_err, date; args...)
getdy_err(date; args...) = getdy_err(get(EOP_DATA), date; args...)

"""
    precession_nutation00(date; outside_range=:warn)

Get the celestial pole coordinate corrections for a certain `date` in milliarcseconds.

`date` can either be a `DateTime` object or a Julian date represented by a `Float64`.
The `outside_range` argument determines what to do if no data is available for `date`:
- `:warn`: The last valid value is returned and a warning will be displayed.
- `:nothing`: The last valid value is returned.
- `:error`: An `OutOfRangeError` is thrown.
"""
precession_nutation00(eop, date; args...) = (getdx(eop, date; args...), getdy(eop, date; outside_range=:warn, args...))
precession_nutation00(date; args...) = (getdx(date; args...), getdy(date; outside_range=:warn, args...))

end # module


text:
function GraggBulirschStoerIntegrator(arg0::jdouble, arg1::jdouble, arg2::Vector{jdouble}, arg3::Vector{jdouble})
    return GraggBulirschStoerIntegrator((jdouble, jdouble, Vector{jdouble}, Vector{jdouble}), arg0, arg1, arg2, arg3)
end

function GraggBulirschStoerIntegrator(arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return GraggBulirschStoerIntegrator((jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function integrate(obj::GraggBulirschStoerIntegrator, arg0::ExpandableODE, arg1::ODEState, arg2::jdouble)
    return jcall(obj, "integrate", ODEStateAndDerivative, (ExpandableODE, ODEState, jdouble), arg0, arg1, arg2)
end

function set_control_factors(obj::GraggBulirschStoerIntegrator, arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return jcall(obj, "setControlFactors", void, (jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function set_interpolation_control(obj::GraggBulirschStoerIntegrator, arg0::jboolean, arg1::jint)
    return jcall(obj, "setInterpolationControl", void, (jboolean, jint), arg0, arg1)
end

function set_order_control(obj::GraggBulirschStoerIntegrator, arg0::jint, arg1::jdouble, arg2::jdouble)
    return jcall(obj, "setOrderControl", void, (jint, jdouble, jdouble), arg0, arg1, arg2)
end

function set_stability_check(obj::GraggBulirschStoerIntegrator, arg0::jboolean, arg1::jint, arg2::jint, arg3::jdouble)
    return jcall(obj, "setStabilityCheck", void, (jboolean, jint, jint, jdouble), arg0, arg1, arg2, arg3)
end