#!/usr/bin/env Rscript

# =============================================================================
# Command-line arguments
# =============================================================================
# Keep it simple

N_JOBS <- 22

possible_cmd_args <- c(as.character(1:N_JOBS), "all")
cmdargs <- commandArgs(trailingOnly=TRUE)
cat("Command-line arguments:", paste(cmdargs, collapse=" "), "\n")
if (!all(cmdargs %in% possible_cmd_args)) {
    stop("Invalid arguments; must be one or more of: ",
         paste(possible_cmd_args, collapse=", "))
}


# =============================================================================
# Libraries
# =============================================================================

library(Rcpp)  # or Stan will fail with: could not find function "cpp_object_initializer"  -- also, load it early on osprey (version conflicts)
library(data.table)
library(ggplot2)
library(lme4)
library(lmerTest)
library(parallel)
library(plyr)
library(readr)
library(rstan)  # to install: install.packages("rstan")
library(semver)
library(shinystan)
library(bridgesampling)

RLIB_PREFIX = "http://egret.psychol.cam.ac.uk/rlib"

source(paste(RLIB_PREFIX, "debugfunc.R", sep="/"))
source(paste(RLIB_PREFIX, "listassign.R", sep="/"))
source(paste(RLIB_PREFIX, "listfunc.R", sep="/"))
source(paste(RLIB_PREFIX, "miscfile.R", sep="/"))
source(paste(RLIB_PREFIX, "miscmath.R", sep="/"))
source(paste(RLIB_PREFIX, "miscstat.R", sep="/"))
source(paste(RLIB_PREFIX, "rpm.R", sep="/"))
source(paste(RLIB_PREFIX, "stanfunc.R", sep="/"))

CODE_COMMON_FUNCTIONS <- readr::read_file(paste(
    RLIB_PREFIX, "commonfunc.stan", sep="/"))

# As advised by Stan:

rstan::rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())


# =============================================================================
# Directories
# =============================================================================

THIS_SCRIPT_DIR = miscfile$current_script_directory()
MAIN_ANALYSIS_DIR <- file.path(THIS_SCRIPT_DIR, "..")
FITS_DIR <- file.path(THIS_SCRIPT_DIR, "pr_fits")


# =============================================================================
# Import synthetic data
# =============================================================================

setwd(THIS_SCRIPT_DIR)
simdata <- data.table(read.csv("synth_parameter_recovery.csv"))

param_combinations <- simdata[
    , .(alpha_reward, alpha_punishment, tau_reinforcement_sensitivity,
        tau_location_stickiness, tau_stimulus_stickiness)]
param_combinations <- unique(param_combinations)


# =============================================================================
# Stan code
# =============================================================================

DATABLOCK <- '
    int<lower=0> N_SUBJECTS;  // number of subjects
    int<lower=0> N_TRIALS;  // TOTAL number of trials
    int<lower=0> N_STIMULI;  // number of pictures (stimuli!)

    int<lower=1, upper=N_SUBJECTS> subject[N_TRIALS];
    int<lower=1, upper=N_STIMULI> left_stimulus[N_TRIALS];
    int<lower=1, upper=N_STIMULI> right_stimulus[N_TRIALS];
    int<lower=0, upper=1> responded_right[N_TRIALS];
    int<lower=0, upper=1> outcome[N_TRIALS];

    // DEBUGGING // real<lower=0, upper=1> source_p_choose_right[N_TRIALS];
    // DEBUGGING // real<lower=0, upper=1> source_v_right[N_TRIALS];
'

STANCODE_M2_CUTDOWN <- '

// Note special *limited* analysis for measuring parameter recovery from the
// core RL algorithm, including no intersubject variability of parameters,
// since there is none; see
//      sim_from_model.py --parameters parameter_recovery_demo

// See comments in reversals_2.stan

functions {
// #include commonfunc.stan
}

data {
// #include datablock.stan
}

transformed data {
// #include constants.stan
}

parameters {
    real<lower=0, upper=1> reward_rate;
    real<lower=0, upper=1> punish_rate;
    real<lower=0> reinf_sensitivity;
    real side_stickiness;
    real stimulus_stickiness;
}

model {
    vector[N_TRIALS] p_choose_rhs;

    sampleBeta_RRR_lp(reward_rate, PRIOR_BETA_SHAPE1, PRIOR_BETA_SHAPE2);
    sampleBeta_RRR_lp(punish_rate, PRIOR_BETA_SHAPE1, PRIOR_BETA_SHAPE2);
    sampleGamma_RRR_lp(reinf_sensitivity, PRIOR_GAMMA_ALPHA_FOR_REINF_SENSITIVITY_MEAN, PRIOR_GAMMA_BETA_FOR_REINF_SENSITIVITY_MEAN);  // positive only
    sampleNormal_RRR_lp(side_stickiness, 0, 1);
    sampleNormal_RRR_lp(stimulus_stickiness, 0, 1);

    {
        int s = 0; // current subject, starting with an invalid value

        vector[N_STIMULI] stim_value;  // value of each stimulus; range -1 to +1

        // Temporary variables: RL calculations
        vector[N_SIDES] values_left_right;
        vector[N_SIDES] side_stickiness_left_right;  // integer contents, but needs to be real for vector multiplication
        vector[N_SIDES] stimulus_stickiness_left_right;  // integer contents, but needs to be real for vector multiplication
        vector[N_SIDES] softmax_inputs;
        int first_trial;  // stickiness only makes sense for trialnum > 1
        int chosen_stimulus_index;  // takes values 1 ... N_STIMULI
        real predicted_outcome;
        real actual_outcome;
        real prediction_error;
        real value_change;

        real temp_new_value;

        for (t in 1:N_TRIALS) {
            if (s != subject[t]) {
                s = subject[t];
                first_trial = 1;

                for (stim in 1:N_STIMULI) {  // better way to assign to a vector??
                    stim_value[stim] = 0;  // all stimuli start with neutral value
                }

            } else {
                first_trial = 0;
            }

            values_left_right[LEFT] = stim_value[left_stimulus[t]];
            values_left_right[RIGHT] = stim_value[right_stimulus[t]];

            // Side stickiness:
            if (!first_trial) {
                side_stickiness_left_right[RIGHT] = responded_right[t - 1];
                side_stickiness_left_right[LEFT] = 1 - responded_right[t - 1];
            } else {
                side_stickiness_left_right[RIGHT] = 0;
                side_stickiness_left_right[LEFT] = 0;
            }

            // Stimulus stickiness:
            if (!first_trial) {
                int stimulus_chosen_last_trial =
                    responded_right[t - 1]
                    ? right_stimulus[t - 1]
                    : left_stimulus[t - 1];
                stimulus_stickiness_left_right[RIGHT] = right_stimulus[t] == stimulus_chosen_last_trial ? 1 : 0;
                stimulus_stickiness_left_right[LEFT] = left_stimulus[t] == stimulus_chosen_last_trial ? 1 : 0;
            } else {
                stimulus_stickiness_left_right[RIGHT] = 0;
                stimulus_stickiness_left_right[LEFT] = 0;
            }

            // Calculate p
            softmax_inputs = (
                reinf_sensitivity * values_left_right +
                side_stickiness * side_stickiness_left_right +
                stimulus_stickiness * stimulus_stickiness_left_right
            );

            // Choose
            p_choose_rhs[t] = softmaxNth(softmax_inputs, RIGHT);

            // Update

            if (responded_right[t]) {
                chosen_stimulus_index = right_stimulus[t];
            } else {
                chosen_stimulus_index = left_stimulus[t];
            }

            predicted_outcome = stim_value[chosen_stimulus_index];
            actual_outcome = outcome[t];
            prediction_error = actual_outcome - predicted_outcome;

            if (prediction_error > 0) {
                value_change = prediction_error * reward_rate;
            } else {
                value_change = prediction_error * punish_rate;
            }
            temp_new_value = predicted_outcome + value_change;
            if (temp_new_value < 0.0 || temp_new_value > 1.0) {
                reject("Error: stimulus value out of range [0, 1]!");
            }
            stim_value[chosen_stimulus_index] = temp_new_value;
        }
    }

    // Final fit to behaviour

    sampleBernoulli_AV_lp(responded_right, p_choose_rhs);
}

'


STANCODE_M7_CUTDOWN <- '

functions {
// #include commonfunc.stan
}

data {
// #include datablock.stan
}

transformed data {
// #include constants.stan
}

parameters {
    real<lower=0, upper=1> alpha;
    real<lower=0> beta;
}

model {
    vector[N_TRIALS] p_choose_rhs;

    sampleBeta_RRR_lp(alpha, PRIOR_BETA_SHAPE1, PRIOR_BETA_SHAPE2);
    sampleGamma_RRR_lp(beta, PRIOR_GAMMA_ALPHA_FOR_REINF_SENSITIVITY_MEAN, PRIOR_GAMMA_BETA_FOR_REINF_SENSITIVITY_MEAN);  // positive only

    {
        int s = 0; // current subject, starting with an invalid value

        vector[N_STIMULI] stim_value;  // value of each stimulus; range -1 to +1

        // Temporary variables: RL calculations
        vector[N_SIDES] values_left_right;
        vector[N_SIDES] softmax_inputs;
        int chosen_stimulus_index;  // takes values 1 ... N_STIMULI
        real predicted_outcome;
        real actual_outcome;
        real prediction_error;
        real value_change;

        real temp_new_value;

        for (t in 1:N_TRIALS) {
            if (s != subject[t]) {
                s = subject[t];

                for (stim in 1:N_STIMULI) {  // better way to assign to a vector??
                    stim_value[stim] = 0;  // all stimuli start with neutral value
                }
            }

            values_left_right[LEFT] = stim_value[left_stimulus[t]];
            values_left_right[RIGHT] = stim_value[right_stimulus[t]];

            // Calculate p
            softmax_inputs = (
                beta * values_left_right
            );

            // Choose
            p_choose_rhs[t] = softmaxNth(softmax_inputs, RIGHT);

            // Update

            if (responded_right[t]) {
                chosen_stimulus_index = right_stimulus[t];
            } else {
                chosen_stimulus_index = left_stimulus[t];
            }

            predicted_outcome = stim_value[chosen_stimulus_index];
            actual_outcome = outcome[t];
            prediction_error = actual_outcome - predicted_outcome;

            value_change = prediction_error * alpha;
            temp_new_value = predicted_outcome + value_change;
            if (temp_new_value < 0.0 || temp_new_value > 1.0) {
                reject("Error: stimulus value out of range [0, 1]!");
            }
            stim_value[chosen_stimulus_index] = temp_new_value;
        }
    }

    // Final fit to behaviour

    sampleBernoulli_AV_lp(responded_right, p_choose_rhs);
}

'


STANCODE_M8_CUTDOWN <- '

functions {
// #include commonfunc.stan
}

data {
// #include datablock.stan
}

transformed data {
// #include constants.stan
}

parameters {
    real<lower=0, upper=1> alpha;
}

model {
    vector[N_TRIALS] p_choose_rhs;
    // DEBUGGING // vector[N_TRIALS] v_right;

    sampleBeta_RRR_lp(alpha, PRIOR_BETA_SHAPE1, PRIOR_BETA_SHAPE2);

    {
        int s = 0; // current subject, starting with an invalid value

        vector[N_STIMULI] stim_value;  // value of each stimulus; range -1 to +1

        // Temporary variables: RL calculations
        vector[N_SIDES] values_left_right;
        vector[N_SIDES] softmax_inputs;
        int first_trial;  // stickiness only makes sense for trialnum > 1
        int chosen_stimulus_index;  // takes values 1 ... N_STIMULI
        real predicted_outcome;
        real actual_outcome;
        real prediction_error;
        real value_change;

        real temp_new_value;

        for (t in 1:N_TRIALS) {
            if (s != subject[t]) {
                s = subject[t];
                first_trial = 1;

                for (stim in 1:N_STIMULI) {  // better way to assign to a vector??
                    stim_value[stim] = 0;  // all stimuli start with neutral value
                }

            } else {
                first_trial = 0;
            }

            // DEBUGGING // v_right[t] = stim_value[right_stimulus[t]];

            values_left_right[LEFT] = stim_value[left_stimulus[t]];
            values_left_right[RIGHT] = stim_value[right_stimulus[t]];

            // Calculate p
            softmax_inputs = (
                values_left_right
            );

            // Choose
            p_choose_rhs[t] = softmaxNth(softmax_inputs, RIGHT);

            // Update

            if (responded_right[t]) {
                chosen_stimulus_index = right_stimulus[t];
            } else {
                chosen_stimulus_index = left_stimulus[t];
            }

            predicted_outcome = stim_value[chosen_stimulus_index];
            actual_outcome = outcome[t];
            prediction_error = actual_outcome - predicted_outcome;

            value_change = prediction_error * alpha;
            temp_new_value = predicted_outcome + value_change;
            if (temp_new_value < 0.0 || temp_new_value > 1.0) {
                reject("Error: stimulus value out of range [0, 1]!");
            }

            stim_value[chosen_stimulus_index] = temp_new_value;
        }
    }

    // Final fit to behaviour

    sampleBernoulli_AV_lp(responded_right, p_choose_rhs);

    // DEBUGGING // source_v_right ~ normal(v_right, 0.01);
    // DEBUGGING // print("alpha=", alpha, ", target=", target());
}

'

DEBUGGING_JUNK <- '
            print("Trial ", t,
                  ": stim_value[1]=", stim_value[1],
                  ", stim_value[2]=", stim_value[2],
                  ", stim_value[3]=", stim_value[3],
                  ", stim_value[4]=", stim_value[4],
                  ", left_stimulus=", left_stimulus[t],
                  ", right_stimulus=", right_stimulus[t],
                  ", p_choose_rhs=", p_choose_rhs[t],
                  ", responded_right=", responded_right[t],
                  ", outcome=", outcome[t],
                  ", prediction_error=", prediction_error,
                  ", alpha=", alpha,
                  ", value_change=", value_change);
'


CODE_CONSTANTS <- readr::read_file(paste(
    MAIN_ANALYSIS_DIR, "constants.stan", sep="/"))

STANCODE <- STANCODE_M2_CUTDOWN
FILEPREFIX = "m2"
# DEBUGGING # STANCODE <- STANCODE_M7_CUTDOWN
# DEBUGGING # FILEPREFIX = "m7"
# DEBUGGING # STANCODE <- STANCODE_M8_CUTDOWN
# DEBUGGING # FILEPREFIX = "m8"

STANCODE <- gsub("// #include datablock.stan", DATABLOCK, STANCODE)
STANCODE <- gsub("// #include commonfunc.stan", CODE_COMMON_FUNCTIONS, STANCODE)
STANCODE <- gsub("// #include constants.stan", CODE_CONSTANTS, STANCODE)


# =============================================================================
# Stan data
# =============================================================================

make_subset_data <- function(simdata,
                             alpha_reward, alpha_punishment,
                             tau_reinforcement_sensitivity,
                             tau_location_stickiness, tau_stimulus_stickiness)
{
    cat(paste0("Full data has ", nrow(simdata), " rows\n"))
    whichrows = (
        simdata$alpha_reward == alpha_reward &
        simdata$alpha_punishment == alpha_punishment &
        simdata$tau_reinforcement_sensitivity == tau_reinforcement_sensitivity &
        simdata$tau_location_stickiness == tau_location_stickiness &
        simdata$tau_stimulus_stickiness == tau_stimulus_stickiness
    )
    subset_data <- simdata[whichrows]
    cat(paste0("Subset has ", nrow(subset_data), " rows\n"))
    setkey(subset_data, subject_num, trial_number)
    return(subset_data)
}


make_standata <- function(subset_data)
{
    standata <- list(
        N_SUBJECTS = length(unique(subset_data$subject_num)),
        N_TRIALS = nrow(subset_data),
        N_STIMULI = 4,
        subject = subset_data$subject_num,
        left_stimulus = subset_data$left_stimulus,
        right_stimulus = subset_data$right_stimulus,
        responded_right = subset_data$responded_right,
        # DEBUGGING # source_p_choose_right = subset_data$p_choose_right,
        # DEBUGGING # source_v_right = subset_data$v_right,
        outcome = subset_data$outcome
    )
    return(standata)
}


ensure_sensible <- function(subset_data)
{
    s <- subset_data
    b_tib <- s[
        ,
        list(
            pct_correct = sum(chose_correctly) / nrow(.SD)
        ),
        by = .(block, trial_in_block)
    ]
    setkey(b_tib, block, trial_in_block)
    p <- (
        ggplot(s, aes(x=trial_number, y=chose_correctly,
                      by=subject_num, colour=block)) +
        geom_line()
    )
    pcheck <- list(
        p_right_low_chose_right = (
            sum(s$p_choose_right < 0.5 & s$responded_right) /
            sum(s$p_choose_right < 0.5)
        ),
        p_right_half_chose_right = (
            sum(s$p_choose_right == 0.5 & s$responded_right) /
            sum(s$p_choose_right == 0.5)
        ),
        p_right_high_chose_right = (
            sum(s$p_choose_right > 0.5 & s$responded_right) /
            sum(s$p_choose_right > 0.5)
        )
    )
    return(list(
        b_tib=b_tib,
        p=p,
        pcheck=pcheck
    ))
}


summarize_fit <- function(fit) {
    s <- stanfunc$summary_data_table(fit)
    f <- function(x) {
        # NOT THIS: format(x, scientific = FALSE, digits = 3, trim = TRUE)
        # ... doesn't handle "digit" well
        y <- sprintf("%.3f", x)
        y <- gsub("-", "–", y)
        return(y)
    }
    maketextcol <- function(m, a, b, r) {
        paste0(f(m), " [", f(a), ", ", f(b), "] (R=", f(r), ")")
    }
    s[, summary := maketextcol(mean, s[["2.5%"]], s[["97.5%"]], Rhat)]
    s <- s[]  # for data table display bug
    return(s)
}


run_parameter_recovery <- function(cmdargs, write_text_output = FALSE)
{
    startsink <- function(filename, overwrite = FALSE) {
        if (write_text_output) {
            sink(file=filename, split=TRUE, append=!overwrite)
        }
    }
    endsink <- function() {
        if (write_text_output) {
            sink()
        }
    }

    for (i in 1:nrow(param_combinations)) {
        if (!(i %in% cmdargs || "all" %in% cmdargs)) {
            next
        }

        alpha_reward <- param_combinations[i, alpha_reward]
        alpha_punishment <- param_combinations[i, alpha_punishment]
        tau_reinforcement_sensitivity <- param_combinations[i, tau_reinforcement_sensitivity]
        tau_location_stickiness <- param_combinations[i, tau_location_stickiness]
        tau_stimulus_stickiness <- param_combinations[i, tau_stimulus_stickiness]
        subset_data <- make_subset_data(
                simdata, alpha_reward, alpha_punishment,
                tau_reinforcement_sensitivity,
                tau_location_stickiness, tau_stimulus_stickiness)
        standata <- make_standata(subset_data)
        fit_filename <- paste0(FITS_DIR,
                               "/", FILEPREFIX, "_pr",
                               "_ar", alpha_reward,
                               "_ap", alpha_punishment,
                               "_tr", tau_reinforcement_sensitivity,
                               "_tl", tau_location_stickiness,
                               "_ts", tau_stimulus_stickiness,
                               ".rds")

        OUTPUT_FILENAME <- paste0(FITS_DIR, "/", FILEPREFIX, "_pr_output_", i, ".txt")
        startsink(OUTPUT_FILENAME, overwrite = TRUE)
        cat(paste0("Starting parameter recovery analysis for job: ", i, "\n"))
        print(param_combinations[i])
        cat(paste0(
            "SIMULATION PARAMETERS:\n",
            "alpha_reward=", alpha_reward,
            ", alpha_punishment=", alpha_punishment,
            ", tau_reinforcement_sensitivity=", tau_reinforcement_sensitivity,
            ", tau_location_stickiness=", tau_location_stickiness,
            ", tau_stimulus_stickiness=", tau_stimulus_stickiness,
            "\n"))
        endsink()

        CHAINS <- 8  # rstan default: 4
        ITER <- 2000  # rstan default: 2000
        INIT <- "random"  # rstan default: "random"
        # ALTERNATIVE # INIT <- "0"  # rstan default: "random"
        SEED <- 1234  # rstan default: a random number

        ADAPT_DELTA <- 0.95  # range 0-1; rstan default: 0.8 (https://cran.r-project.org/web/packages/rstan/rstan.pdf)
        # ADAPT_DELTA <- 0.99  # range 0-1; rstan default: 0.8 (https://cran.r-project.org/web/packages/rstan/rstan.pdf)
        STEPSIZE <- 1  # rstan default: 1 [https://github.com/stan-dev/rstan/blob/develop/rstan3/R/AllClass.R]
        MAX_TREEDEPTH <- 10  # rstan default: 10 [https://github.com/stan-dev/rstan/blob/develop/rstan3/R/AllClass.R]

        model_name <- "Kanen_M2_param_recovery"

        fit <- stanfunc$load_or_run_stan(
            fit_filename = fit_filename,
            model_name = model_name,
            model_code = STANCODE,
            data = standata,
            chains = CHAINS,
            iter = ITER,
            init = INIT,
            seed = SEED,
            control = list(
                adapt_delta = ADAPT_DELTA,
                stepsize = STEPSIZE,
                max_treedepth = MAX_TREEDEPTH
            )
        )

        stansummary <- summarize_fit(fit)
        startsink(OUTPUT_FILENAME)
        cat("STAN RESULTS:\n")
        print(stansummary)
        endsink()
    }
}


DEBUG_CODE <- FALSE
if (DEBUG_CODE) {
    sink("_tmp_stancode.txt")
    cat(STANCODE)
    sink()
    count_char_occurrences <- function(char, s) {
        s2 <- gsub(char, "", s, fixed = TRUE)
        return (nchar(s) - nchar(s2))
    }
    print(count_char_occurrences("{", STANCODE))
    print(count_char_occurrences("}", STANCODE))  # same
}


run_parameter_recovery(cmdargs, write_text_output = TRUE)

# For human use:

debugfunc$wideScreen()

# Examples

TEST <- '

subset1 <- make_subset_data(simdata, 0.1, 0.5, 4.5, 0, 0)
subset3 <- make_subset_data(simdata, 0.5, 0.5, 4.5, 0, 0)
subset5 <- make_subset_data(simdata, 0.9, 0.5, 4.5, 0, 0)
subset10 <- make_subset_data(simdata, 0.5, 0.5, 1, 0, 0)

standata10 <- make_standata(subset10)

ensure_sensible(subset1)
ensure_sensible(subset3)
ensure_sensible(subset10)

m2fit6 <- readRDS(file.path(FITS_DIR, "m2_pr_ar0.5_ap0.1_tr4.5_tl0_ts0.rds"))
m2fit13 <- readRDS(file.path(FITS_DIR, "m2_pr_ar0.5_ap0.5_tr6_tl0_ts0.rds"))

stanfunc$test_specific_parameter_from_stanfit(m2fit6, "stimulus_stickiness", hdi_proportion=0.99)
stanfunc$test_specific_parameter_from_stanfit(m2fit13, "side_stickiness", hdi_proportion=0.99)

m7fit1 <- readRDS(file.path(FITS_DIR, "m7_pr_ar0.1_ap0.5_tr4.5_tl0_ts0.rds"))
m7fit3 <- readRDS(file.path(FITS_DIR, "m7_pr_ar0.5_ap0.5_tr4.5_tl0_ts0.rds"))
m7fit5 <- readRDS(file.path(FITS_DIR, "m7_pr_ar0.9_ap0.5_tr4.5_tl0_ts0.rds"))

m8fit1 <- readRDS(file.path(FITS_DIR, "m8_pr_ar0.1_ap0.5_tr4.5_tl0_ts0.rds"))
m8fit3 <- readRDS(file.path(FITS_DIR, "m8_pr_ar0.5_ap0.5_tr4.5_tl0_ts0.rds"))
m8fit5 <- readRDS(file.path(FITS_DIR, "m8_pr_ar0.9_ap0.5_tr4.5_tl0_ts0.rds"))
m8fit10 <- readRDS(file.path(FITS_DIR, "m8_pr_ar0.5_ap0.5_tr1_tl0_ts0.rds"))

summarize_fit(m7fit1)

'