%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%% BREATHING FILTER DETECTION TASK ANALYSIS %%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% -------------------------------------------------------------------------
% Author: Olivia Harrison
% Created: 14/08/2018
%
% This software is free software: you can redistribute it and/or modify it
% under the terms of the GNU General Public License as published by the
% Free Software Foundation, either version 3 of the License, or (at your
% option) any later version. This software is distributed in the hope that
% it will be useful, but WITHOUT ANY WARRANTY; without even the implied
% warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
% GNU General Public License for more details: http://www.gnu.org/licenses/
% -------------------------------------------------------------------------

% This analysis script uses either Brian Maniscalco's single subect
% analysis or Steve Fleming's Hierarchical Bayesian toolbox (for group
% analysis) to calculate perceptual decision and metacognitive metrics,
% specific to data produced by running the tapas_filter_detection_task.

% IMPORTANT ANALYSIS NOTES:
%   1) The code requires JAGS software to run, which can be found here:
%      http://mcmc-jags.sourceforge.net/
%      The code has been tested on JAGS 3.4.0; there are compatibility
%      issues between matjags and JAGS 4.X.
%   2) The code requires the HMeta-d toolbox (provided as a submodule
%      from https://github.com/metacoglab/HMeta-d) to be on your path.
%   3) This script is looking for specific results files that were output
%      from the tapas_filter_detection_task, which are located in the
%      FDT/results/ folder. The results are expected to contain a
%      results.filterThreshold.xx structure, where values for xx:
%         - results.filterThreshold.filterNum: a value to determine the
%           number of filters where trials were performed
%         - results.filterThreshold.filters: a vector with 1 meaning
%           filters were presented, 0 when filters were absent (dummy)
%         - results.filterThreshold.response: a vector with 1 meaning
%           response was 'yes', 0 when response was 'no'
%         - results.filterThreshold.confidence: a vector containing
%           confidence score (1-10) on each trial
%      If your results are formatted differently, refer directly to the
%      original scripts from the HMeta-d toolbox
%      (https://github.com/metacoglab/HMeta-d).

% TO RUN THIS SCRIPT:
% type tapas_filter_detection_analysis into the MATLAB terminal from the
% main FDT folder, and follow the prompts.

% ANALYSIS OPTIONS:
% This script allows you to run either:
%    - A single subject analysis: A non-Bayesian estimation of meta-d',
%      fitting one subject at a time using a single point maximum
%      likelihood estimation (from the original Maniscalco & Lau 2012
%      paper, more info can be found here:
%      http://www.columbia.edu/~bsm2105/type2sdt/archive/index.html).
%      HOWEVER: Simulations have shown this analysis to be VERY unstable
%      when estimating meta-d' using 60 trials. It is STRONGLY encouraged
%      to utilise the hierarchical models, or collect many more trials
%      (200+ trials) for each subject to have a more reliable measure of
%      meta-d'.
%    - A group mean analysis: A hierarchical Bayesian analysis, whereby all
%      subjects are fit together and information from the group mean is
%      used as prior information for each subject. This hierarchical model
%      helps with much more accurate estimations of meta-d' with small
%      trial numbers, such as using 60 trials per subject.
%    - A group regression analysis: A hierarchical Bayesian regression
%      analysis that is an extension on the group mean analysis, whereby
%      subjects are fit together in one model, and the relationship between
%      single subject log(meta-d'/d') values and the covariate(s) is
%      estimated within the hierarchical model structure.
%    - A group difference analysis: A hierarchical Bayesian analysis,
%      whereby each group of subjects is fitted separately, and the groups
%      are then compared. Frequentist statistics (such as parametric
%      unpaired T-tests, or non-parametric Wilcoxon signed-rank tests) can
%      be used for all values that are not fitted using hierarchical
%      information, such as d', c, filter number, accuracy and average
%      confidence. As the group values for log(meta-d'/d') are calculated
%      using two separate hierarchical models, group differences are then
%      inferred via comparison of the resulting log(meta-d'/d')
%      distributions, and whether the 95% highest-density interval (HDI) of
%      the difference between the distributions spans zero. The HDI is the
%      shortest possible interval containing 95% of the MCMC samples, and
%      may not be symmetric. NOTE: If single subject values for
%      log(meta-d'/d') (or any related meta-d' metric) are required for
%      further analyses that span both groups (such as entry into general
%      linear models), it is recommended to fit all subjects together in
%      one regression model with a regressor denoting group identity.
%    - A session difference analysis (paired group difference): A
%      hierarchical Bayesian analysis, whereby two sessions / measures from
%      the same participants are fitted in a single model using a
%      multivariate normal distribution. This distribution allows for the
%      non-independence between sessions for each participant. NOTE:
%      Participants must be listed in the same order for analysis, and must
%      have data for both sessions / each measure.

% NOTES ON NON-BAYESIAN SINGLE SUBJECT ANALYSIS FROM AUTHORS' WEBSITE
% (columbia.edu/~bsm2105/type2sdt/archive/index.html):
% This analysis is intended to quantify metacognitive sensitivity (i.e. the
% efficacy with which confidence ratings discriminate between correct and
% incorrect judgments) in a signal detection theory framework. A central
% idea is that primary task performance can influence metacognitive
% sensitivity, and it is informative to take this influence into account.
% Description of the methodology can be found here: Maniscalco, B., & Lau,
% H. (2012). A signal detection theoretic approach for estimating
% metacognitive sensitivity from confidence ratings. Consciousness and
% Cognition, 21(1), 422–430. doi:10.1016/j.concog.2011.09.021
% If you use these analysis files, please reference the Consciousness &
% Cognition paper and website. Brian Maniscalco: brian@psych.columbia.edu

% NOTES ON HIERARCHICAL TOOLBOX FROM ORIGINAL SCRIPTS (FOR GROUP ANALYSES):
% This MATLAB toolbox implements the meta-d’ model (Maniscalco & Lau, 2012)
% in a hierarchical Bayesian framework using Matlab and JAGS, a program for
% conducting MCMC inference on arbitrary Bayesian models. A paper with more
% details on the method and the advantages of estimating meta-d’ in a
% hierarchal Bayesian framework is available here https://academic.oup.com/
% nc/article/doi/10.1093/nc/nix007/3748261/HMeta-d-hierarchical-Bayesian-
% estimation-of. For a more general introduction to Bayesian models of
% cognition see Lee & Wagenmakers, Bayesian Cognitive Modeling: A Practical
% Course http://bayesmodels.com/. This code is being released with a
% permissive open-source license. You should feel free to use or adapt the
% utility code as long as you follow the terms of the license provided. If
% you use the toolbox in a publication we ask that you cite the following
% paper: Fleming, S.M. (2017) HMeta-d: hierarchical Bayesian estimation of
% metacognitive efficiency from confidence ratings, Neuroscience of
% Consciousness, 3(1) nix007, https://doi.org/10.1093/nc/nix007. Copyright
% (c) 2017, Stephen Fleming. For more information and/or licences, please
% see the original code: https://github.com/metacoglab/HMeta-d.

% KEY ANALYSIS OUTPUT VALUES:
% The primary measures to come out of this analysis can be found in
% analysis.{single/groupMean/groupDiff}.xx, and they are:
%   1) xx = filterNum: Number of filters. Less filters means more sensitive
%           to changes in breathing.
%   2) xx = d1: d prime, discriminability between filter and dummy trials.
%           Larger d' means more discriminability at specific filter number
%   3) xx = c1: Decision criterion, or where the decision boundary exists.
%           Negative criterion values mean bias towards 'yes' (lower, more
%           liberal criterion), while positive values mean bias towards
%           'no' response (higher, more strngent criterion).
%   4) xx = meta_d: A measure of metacognition, or 'type 2' sensitivity.
%           This reflects how much information, in signal-to-noise units,
%           is available for metacognition (see Maniscalco & Lau, 2012).
%           NB: ISSUES WITH THIS ESTIMATION (AND ALL RELATED ESTIMATIONS)
%           IF USING SINGLE SUBJECT ANALYSES, AS DESCRIBED ABOVE.
%   5) xx = Mratio: Meta-d' to d' ratio, as a measure of metacognitive
%           'efficiency' (see Rouault et al., 2018).
%   6) xx = log_Mratio: log(meta_d/d1), reported in papers to help with
%           normalisation of data (see Rouault et al., 2018).
%   7) xx = avgConfidence: A second measure of metacognition, thought to be
%           independent from meta-d'. NB: Average confidence should only be
%           compared between two groups if there is no difference in d'
%           between the groups (i.e. task difficulty was comparable).

% ADDITIONAL GROUP DIFFERENCE OUTPUT VALUES:
% If the analysis is a two-group or two-session (paired) difference, the
% following will also be calculated for the non-hierarchical measures
% (test = groupDiff or sessionDiff; xx = filterNum, d1, c1 and
% avgConfidence):
%   1) analysis.test.xx.h = Results of null hypothesis test, where h =
%      1 for rejection of the null, and h = 0 for no rejection.
%   2) analysis.test.xx.p = The p-value for the statistical test.
%   3) analysis.test.xx.stats = Further statistics (such as tstats,
%      number of samples, degrees of freedom) associated with the test.
%   4) analysis.test.xx.ci = Confidence interval of the difference,
%      calculated only when T-test is specified.
% The highest density interval (HDI) for the difference in log(meta-d'/d')
% between the groups / sessions will be calculated and recorded, as
% frequentist statistics cannot be used here. A summary Figure for each of
% these metrics will be created and saved in the FDT/analysis/ folder.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% SET UP THE ANALYSIS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function tapas_filter_detection_analysis()

% Check folder location is main FDT folder
[~,dir_name] = fileparts(pwd);
if ~strcmp(dir_name,'FDT')
   error('Not currently in main FDT folder. Please move to FDT folder and try again.');
end

% Add relevant paths
addpath('analysis');

% Display setup on screen
fprintf('\n________________________________________\n\n    SET UP ANALYSIS FOR FILTER TASK\n________________________________________\n');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% CHOOSE THE TYPE OF ANALYSIS TO RUN AND SPECIFY FILES
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

try
    analysis.type = input('Type of analysis (single or group) = ', 's'); % Ask for type of analysis to be run
catch
    fprintf('\nOption does not exist, please try again.\n');
    return
end

try
    if strcmp(analysis.type,'single') == 1 % If single subject analysis specified
        analysis.PPIDs = input('PPID = ', 's'); % Ask for PPID
        try
            analysis.data = load(fullfile('results', ['filter_task_results_', analysis.PPIDs, '.mat'])); % Load data
            analysis.ratings = analysis.data.results.setup.confidenceUpper - analysis.data.results.setup.confidenceLower + 1; % Calculate number of confidence rating bins
        catch
            fprintf('\nInvalid PPID.\n');
            return
        end
    elseif strcmp(analysis.type,'group') == 1 % If group analysis specified
        analysis.type = input('Type of analysis (mean, diff, paired or regress) = ', 's'); % Ask for type of group analysis to be run
        if strcmp(analysis.type,'mean') == 1 % If group mean analysis specified
            analysis.PPIDs = input('Input group PPIDs (e.g. {''001'', ''002'', ''003'',...}) = '); % Ask for PPIDs
            analysis.groupsize(1) = length(analysis.PPIDs);
            for n = 1:length(analysis.PPIDs)
                PPID = char(analysis.PPIDs(n));
                try
                    analysis.data(n) = load(fullfile('results', ['filter_task_results_', PPID, '.mat'])); % Load data
                catch
                    fprintf('\nInvalid PPIDs.\n');
                    return
                end
            end
            analysis.ratings = analysis.data(1).results.setup.confidenceUpper - analysis.data(1).results.setup.confidenceLower + 1; % Calculate number of confidence rating bins
        elseif strcmp(analysis.type,'diff') == 1 % If two group difference analysis specified
            analysis.PPIDs.group1 = input('Input group 1 PPIDs (e.g. {''001'', ''002'', ''003'',...}) = '); % Ask for PPIDs
            analysis.PPIDs.group2 = input('Input group 2 PPIDs (e.g. {''004'', ''005'', ''006'',...}) = '); % Ask for PPIDs
            analysis.groupDiff.test = input('Type 1 variable tests (ttest or wilcoxon) = ', 's'); % Ask for parametric or non-parmetric test options for type 1 variables (d', c etc.)
            analysis.groupsize(1) = length(analysis.PPIDs.group1);
            analysis.groupsize(2) = length(analysis.PPIDs.group2);
            for n = 1:analysis.groupsize(1)
                PPID = char(analysis.PPIDs.group1(n));
                try
                    analysis.group1.data(n) = load(fullfile('results', ['filter_task_results_', PPID, '.mat'])); % Load data
                catch
                    fprintf('\nInvalid PPIDs for group 1.\n');
                    return
                end
            end
            for n = 1:analysis.groupsize(2)
                PPID = char(analysis.PPIDs.group2(n));
                try
                    analysis.group2.data(n) = load(fullfile('results', ['filter_task_results_', PPID, '.mat'])); % Load data
                catch
                    fprintf('\nInvalid PPIDs for group 2.\n');
                    return
                end
            end
            analysis.ratings = analysis.group1.data(1).results.setup.confidenceUpper - analysis.group1.data(1).results.setup.confidenceLower + 1; % Calculate number of confidence rating bins
        elseif strcmp(analysis.type,'paired') == 1 % If paired difference analysis specified
            fprintf('\nNOTE: PPIDs from each session need to be paired and in the same order\n');
            analysis.PPIDs.session1 = input('Input session 1 PPIDs (e.g. {''001a'', ''002a'', ''003a'',...}) = '); % Ask for PPIDs
            analysis.PPIDs.session2 = input('Input session 2 PPIDs (e.g. {''001b'', ''002b'', ''003b'',...}) = '); % Ask for PPIDs
            analysis.sessionDiff.test = input('Type 1 variable tests (ttest or wilcoxon) = ', 's'); % Ask for parametric or non-parmetric test options for type 1 variables (d', c etc.)
            analysis.sessionSize(1) = length(analysis.PPIDs.session1);
            analysis.sessionSize(2) = length(analysis.PPIDs.session2);
            if analysis.sessionSize(1) ~= analysis.sessionSize(2)
                error('Number of PPIDs in each session does not match!');
            end
            for n = 1:analysis.sessionSize(1)
                PPID = char(analysis.PPIDs.session1(n));
                try
                    analysis.session1.data(n) = load(fullfile('results', ['filter_task_results_', PPID, '.mat'])); % Load data
                catch
                    fprintf('\nInvalid PPIDs for session 1.\n');
                    return
                end
            end
            for n = 1:analysis.sessionSize(2)
                PPID = char(analysis.PPIDs.session2(n));
                try
                    analysis.session2.data(n) = load(fullfile('results', ['filter_task_results_', PPID, '.mat'])); % Load data
                catch
                    fprintf('\nInvalid PPIDs for session 2.\n');
                    return
                end
            end
            analysis.ratings = analysis.session1.data(1).results.setup.confidenceUpper - analysis.session1.data(1).results.setup.confidenceLower + 1; % Calculate number of confidence rating bins
        elseif strcmp(analysis.type,'regress') == 1 % If regression analysis specified
            analysis.PPIDs = input('Input group PPIDs (e.g. {''001'', ''002'', ''003'',...}) = '); % Ask for PPIDs
            analysis.groupsize(1) = length(analysis.PPIDs);
            for n = 1:length(analysis.PPIDs)
                PPID = char(analysis.PPIDs(n));
                try
                    analysis.data(n) = load(fullfile('results', ['filter_task_results_', PPID, '.mat'])); % Load data
                catch
                    fprintf('\nInvalid PPIDs.\n');
                    return
                end
            end
            analysis.ratings = analysis.data(1).results.setup.confidenceUpper - analysis.data(1).results.setup.confidenceLower + 1; % Calculate number of confidence rating bins
            fprintf('\nNote: Covariate text file required to be placed in results folder\n--> Scores need to be in the SAME ORDER as PPID input order.\n');
            analysis.covariate.fileName = input('Input covariate file name (e.g. covariate_example.txt) = ', 's'); % Ask for covariate file name
            try
                analysis.covariate.data = load(fullfile('results', analysis.covariate.fileName)); % Load data
            catch
                fprintf('\nCovariate file cannot be found in results folder.\n');
                return
            end
            [a,b] = size(analysis.covariate.data);
            if a > b % Re-shape vector if needed
                analysis.covariate.data = analysis.covariate.data';
            end
        end
    end
catch
    fprintf('\nInvalid.\n');
    return
end

% Save results
if strcmp(analysis.type,'single') == 1
    resultsFile = fullfile('analysis', ['filter_task_analysis_', analysis.type, analysis.PPIDs]); % Create figure file name
else
    resultsFile = fullfile('analysis', ['filter_task_analysis_', analysis.type]); % Create figure file name
end
save(resultsFile, 'analysis'); % Save results

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% RUN TRIALS2COUNTS SCRIPT
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if strcmp(analysis.type,'single') == 1 || strcmp(analysis.type,'mean') == 1 || strcmp(analysis.type,'regress') == 1

    % Run trials2counts script for single subject or whole group
    for n = 1:length(analysis.data)
        stimID = analysis.data(n).results.thresholdTrials.filters;
        response = analysis.data(n).results.thresholdTrials.response;
        rating = analysis.data(n).results.thresholdTrials.confidence;
        [analysis.trials2counts.nR_S1{n}, analysis.trials2counts.nR_S2{n}] = trials2counts(stimID,response,rating,analysis.ratings,0); % Run for number of ratings and no cell padding
    end

elseif strcmp(analysis.type,'diff') == 1

    % Run trials2counts script for group 1
    for n = 1:length(analysis.group1.data)
        stimID = analysis.group1.data(n).results.thresholdTrials.filters;
        response = analysis.group1.data(n).results.thresholdTrials.response;
        rating = analysis.group1.data(n).results.thresholdTrials.confidence;
        [analysis.group1.trials2counts.nR_S1{n}, analysis.group1.trials2counts.nR_S2{n}] = trials2counts(stimID,response,rating,analysis.ratings,0); % Run for number of ratings and no cell padding
    end
    % Run trials2counts script for group 2
    for n = 1:length(analysis.group2.data)
        stimID = analysis.group2.data(n).results.thresholdTrials.filters;
        response = analysis.group2.data(n).results.thresholdTrials.response;
        rating = analysis.group2.data(n).results.thresholdTrials.confidence;
        [analysis.group2.trials2counts.nR_S1{n}, analysis.group2.trials2counts.nR_S2{n}] = trials2counts(stimID,response,rating,analysis.ratings,0); % Run for number of ratings and no cell padding
    end

elseif strcmp(analysis.type,'paired') == 1

    % Run trials2counts script for session 1
    for n = 1:length(analysis.session1.data)
        stimID = analysis.session1.data(n).results.thresholdTrials.filters;
        response = analysis.session1.data(n).results.thresholdTrials.response;
        rating = analysis.session1.data(n).results.thresholdTrials.confidence;
        [analysis.trials2counts.nR_S1(1).counts{n}, analysis.trials2counts.nR_S2(1).counts{n}] = trials2counts(stimID,response,rating,analysis.ratings,0); % Run for number of ratings and no cell padding
    end
    % Run trials2counts script for session 2
    for n = 1:length(analysis.session2.data)
        stimID = analysis.session2.data(n).results.thresholdTrials.filters;
        response = analysis.session2.data(n).results.thresholdTrials.response;
        rating = analysis.session2.data(n).results.thresholdTrials.confidence;
        [analysis.trials2counts.nR_S1(2).counts{n}, analysis.trials2counts.nR_S2(2).counts{n}] = trials2counts(stimID,response,rating,analysis.ratings,0); % Run for number of ratings and no cell padding
    end

end

% Save results
save(resultsFile, 'analysis');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% RUN SINGLE SUBJECT ANALYSIS IF SPECIFIED
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if strcmp(analysis.type,'single') == 1 % If single subject analysis specified

    % Pad ratings so that zero counts do not interfere with fit
    analysis.trials2counts.nR_S1_padded = analysis.trials2counts.nR_S1{1} + 1/(2*10);
    analysis.trials2counts.nR_S2_padded = analysis.trials2counts.nR_S2{1} + 1/(2*10);

    % Fit data
    analysis.single.fit = fit_meta_d_MLE(analysis.trials2counts.nR_S1_padded, analysis.trials2counts.nR_S2_padded);

    % Pull out key variables
    analysis.single.filterNum = mean(analysis.data.results.thresholdTrials.filterNum);
    analysis.single.accuracy = analysis.data.results.thresholdTrials.accuracyTotal;
    analysis.single.d1 = analysis.single.fit.d1;
    analysis.single.c1 = analysis.single.fit.c1;
    analysis.single.meta_d = analysis.single.fit.meta_d;
    analysis.single.Mratio = analysis.single.fit.M_ratio;
    analysis.single.fit.log_Mratio = log(analysis.single.fit.M_ratio);
    analysis.single.avgConfidence = mean(analysis.data(1).results.thresholdTrials.confidence);

    % Save results
    save(resultsFile, 'analysis'); % Save results

    % Display completion message on screen
    fprintf('\n________________________________________\n\n   COMPLETED SINGLE SUBJECT ANALYSIS\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    fprintf('\n             MRATIO = %.2f\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n', analysis.single.Mratio);
    fprintf('\n        RESULTS CAN BE FOUND IN: \n             FDT/analysis/\n________________________________________\n');

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% RUN GROUP MEAN ANALYSIS IF SPECIFIED
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if strcmp(analysis.type,'mean') == 1 % If group mean analysis specified

    % Specify parameters
    analysis.groupMean.mcmc_params = fit_meta_d_params;

    % Fit group data all at once
    analysis.groupMean.fit = fit_meta_d_mcmc_group(analysis.trials2counts.nR_S1, analysis.trials2counts.nR_S2, analysis.groupMean.mcmc_params);

    % Calculate log of Mratio for single subjects values
    for n = 1:length(analysis.groupMean.fit.Mratio)
        analysis.groupMean.log_Mratio.singleSubject(n) = log(analysis.groupMean.fit.Mratio(n));
    end

    % Pull out group mean values
    analysis.groupMean.d1.groupMean = mean(analysis.groupMean.fit.d1);
    analysis.groupMean.d1.singleSubject = analysis.groupMean.fit.d1;
    analysis.groupMean.c1.groupMean = mean(analysis.groupMean.fit.c1);
    analysis.groupMean.c1.singleSubject = analysis.groupMean.fit.c1;
    analysis.groupMean.meta_d.groupMean = mean(analysis.groupMean.fit.meta_d);
    analysis.groupMean.meta_d.singleSubject = analysis.groupMean.fit.meta_d;
    analysis.groupMean.Mratio.groupMean = exp(analysis.groupMean.fit.mu_logMratio);
    analysis.groupMean.Mratio.singleSubject = analysis.groupMean.fit.Mratio;
    analysis.groupMean.Mratio.hdi = calc_HDI(exp(analysis.groupMean.fit.mcmc.samples.mu_logMratio(:)));
    analysis.groupMean.log_Mratio.groupMean = analysis.groupMean.fit.mu_logMratio;
    analysis.groupMean.log_Mratio.hdi = calc_HDI(analysis.groupMean.fit.mcmc.samples.mu_logMratio(:));

    % Calculate average confidence, add filter number and accuracy to analysis results
    for n = 1:length(analysis.data)
        analysis.groupMean.avgConfidence.singleSubject(n) = mean(analysis.data(n).results.thresholdTrials.confidence);
        analysis.groupMean.filterNum.singleSubject(n) = mean(analysis.data(n).results.thresholdTrials.filterNum);
        analysis.groupMean.accuracy.singleSubject(n) = analysis.data(n).results.thresholdTrials.accuracyTotal;
    end
    analysis.groupMean.avgConfidence.groupMean = mean(analysis.groupMean.avgConfidence.singleSubject);
    analysis.groupMean.filterNum.groupMean = mean(analysis.groupMean.filterNum.singleSubject);
    analysis.groupMean.accuracy.groupMean = mean(analysis.groupMean.accuracy.singleSubject);

    % Save results
    save(resultsFile, 'analysis'); % Save results

    % Display completion message on screen
    fprintf('\n________________________________________\n\n     COMPLETED GROUP MEAN ANALYSIS\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    fprintf('\n GROUP MRATIO = %.2f (HDI: %.2f to %.2f)\n', analysis.groupMean.Mratio.groupMean, analysis.groupMean.Mratio.hdi(1), analysis.groupMean.Mratio.hdi(2));
    if (analysis.groupMean.Mratio.hdi(1) > 0 && analysis.groupMean.Mratio.hdi(2) > 0) || (analysis.groupMean.Mratio.hdi(1) < 0 && analysis.groupMean.Mratio.hdi(2) < 0)
        fprintf('   HDI DOES NOT SPAN ZERO: MRATIO SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    else
        fprintf('   HDI SPANS ZERO: MRATIO NOT SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    end
    fprintf('\n        RESULTS CAN BE FOUND IN: \n             FDT/analysis/\n________________________________________\n');

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% RUN GROUP DIFFERENCE ANALYSIS (UNPAIRED) IF SPECIFIED
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if strcmp(analysis.type,'diff') == 1 % If group difference analysis specified

    % For two-group difference analysis: Print message about groups being fit separately --> no t-tests possible on meta-d' metrics
    fprintf('\nNOTE: Groups are fit using separate hierarchical models.\nFrequentist statistics (such as t-tests) are not possible on any metrics involving meta-d''.\nFrequentist statistics will still be applied to all type 1 metrics and average confidence values.\n');

    % Specify parameters
    analysis.groupDiff.mcmc_params = fit_meta_d_params;

    % Fit each group using a separate hierarchical Bayesian model
    analysis.group1.fit = fit_meta_d_mcmc_group(analysis.group1.trials2counts.nR_S1, analysis.group1.trials2counts.nR_S2, analysis.groupDiff.mcmc_params);
    analysis.group2.fit = fit_meta_d_mcmc_group(analysis.group2.trials2counts.nR_S1, analysis.group2.trials2counts.nR_S2, analysis.groupDiff.mcmc_params);

    % Compute HDI of difference for log(meta-d'/d')
    analysis.group1.logMratio.groupMean = analysis.group1.fit.mu_logMratio;
    analysis.group1.logMratio.singleSubject = log(analysis.group1.fit.Mratio);
    analysis.group1.Mratio.groupMean = exp(analysis.group1.fit.mu_logMratio);
    analysis.group1.Mratio.singleSubject = analysis.group1.fit.Mratio;
    analysis.group2.logMratio.groupMean = analysis.group2.fit.mu_logMratio;
    analysis.group2.logMratio.singleSubject = log(analysis.group2.fit.Mratio);
    analysis.group2.Mratio.groupMean = exp(analysis.group2.fit.mu_logMratio);
    analysis.group2.Mratio.singleSubject = analysis.group2.fit.Mratio;
    analysis.groupDiff.log_Mratio.sampleDiff = analysis.group1.fit.mcmc.samples.mu_logMratio - analysis.group2.fit.mcmc.samples.mu_logMratio;
    analysis.groupDiff.log_Mratio.mean = analysis.group1.fit.mu_logMratio - analysis.group2.fit.mu_logMratio;
    analysis.groupDiff.log_Mratio.hdi = calc_HDI(analysis.groupDiff.log_Mratio.sampleDiff(:));
    analysis.groupDiff.Mratio.mean = exp(analysis.group1.fit.mu_logMratio) - exp(analysis.group2.fit.mu_logMratio);
    analysis.groupDiff.Mratio.hdi = calc_HDI((exp(analysis.group1.fit.mcmc.samples.mu_logMratio(:)) - exp(analysis.group2.fit.mcmc.samples.mu_logMratio(:))));

    % Pull out group mean values for all non-hierarchical metrics
    analysis.group1.d1.groupMean = mean(analysis.group1.fit.d1);
    analysis.group1.d1.singleSubject = analysis.group1.fit.d1;
    analysis.group1.c1.groupMean = mean(analysis.group1.fit.c1);
    analysis.group1.c1.singleSubject = analysis.group1.fit.c1;
    analysis.group2.d1.groupMean = mean(analysis.group2.fit.d1);
    analysis.group2.d1.singleSubject = analysis.group2.fit.d1;
    analysis.group2.c1.groupMean = mean(analysis.group2.fit.c1);
    analysis.group2.c1.singleSubject = analysis.group2.fit.c1;

    % Calculate average confidence, add filter number and accuracy to results for group 1
    for n = 1:length(analysis.group1.data)
        analysis.group1.avgConfidence.singleSubject(n) = mean(analysis.group1.data(n).results.thresholdTrials.confidence);
        analysis.group1.filterNum.singleSubject(n) = mean(analysis.group1.data(n).results.thresholdTrials.filterNum);
        analysis.group1.accuracy.singleSubject(n) = analysis.group1.data(n).results.thresholdTrials.accuracyTotal;
    end
    analysis.group1.avgConfidence.groupMean = mean(analysis.group1.avgConfidence.singleSubject);
    analysis.group1.filterNum.groupMean = mean(analysis.group1.filterNum.singleSubject);
    analysis.group1.accuracy.groupMean = mean(analysis.group1.accuracy.singleSubject);

    % Calculate average confidence, add filter number and accuracy to results for group 2
    for n = 1:length(analysis.group2.data)
        analysis.group2.avgConfidence.singleSubject(n) = mean(analysis.group2.data(n).results.thresholdTrials.confidence);
        analysis.group2.filterNum.singleSubject(n) = mean(analysis.group2.data(n).results.thresholdTrials.filterNum);
        analysis.group2.accuracy.singleSubject(n) = analysis.group2.data(n).results.thresholdTrials.accuracyTotal;
    end
    analysis.group2.avgConfidence.groupMean = mean(analysis.group2.avgConfidence.singleSubject);
    analysis.group2.filterNum.groupMean = mean(analysis.group2.filterNum.singleSubject);
    analysis.group2.accuracy.groupMean = mean(analysis.group2.accuracy.singleSubject);

    % Calculate group difference values for all non-hierarchical metrics
    analysis.groupDiff.filterNum.meanDiff = analysis.group1.filterNum.groupMean - analysis.group2.filterNum.groupMean;
    analysis.groupDiff.accuracy.meanDiff = analysis.group1.accuracy.groupMean - analysis.group2.accuracy.groupMean;
    analysis.groupDiff.d1.meanDiff = analysis.group1.d1.groupMean - analysis.group2.d1.groupMean;
    analysis.groupDiff.c1.meanDiff = analysis.group1.c1.groupMean - analysis.group2.c1.groupMean;
    analysis.groupDiff.avgConfidence.meanDiff = analysis.group1.avgConfidence.groupMean - analysis.group2.avgConfidence.groupMean;

    % Run specified statistical test on all non-hierarchical metrics
    if strcmp(analysis.groupDiff.test,'wilcoxon') == 1
        analysis.groupDiff.testType = 'Wilcoxon rank sum test';
        [analysis.groupDiff.filterNum.p,analysis.groupDiff.filterNum.h,analysis.groupDiff.filterNum.stats] = ranksum(analysis.group1.filterNum.singleSubject,analysis.group2.filterNum.singleSubject);
        [analysis.groupDiff.accuracy.p,analysis.groupDiff.accuracy.h,analysis.groupDiff.accuracy.stats] = ranksum(analysis.group1.accuracy.singleSubject,analysis.group2.accuracy.singleSubject);
        [analysis.groupDiff.d1.p,analysis.groupDiff.d1.h,analysis.groupDiff.d1.stats] = ranksum(analysis.group1.d1.singleSubject,analysis.group2.d1.singleSubject);
        [analysis.groupDiff.c1.p,analysis.groupDiff.c1.h,analysis.groupDiff.c1.stats] = ranksum(analysis.group1.c1.singleSubject,analysis.group2.c1.singleSubject);
        [analysis.groupDiff.avgConfidence.p,analysis.groupDiff.avgConfidence.h,analysis.groupDiff.avgConfidence.stats] = ranksum(analysis.group1.avgConfidence.singleSubject,analysis.group2.avgConfidence.singleSubject);
    elseif strcmp(analysis.groupDiff.test,'ttest') == 1
        analysis.groupDiff.testType = 'Unpaired T-test';
        [analysis.groupDiff.filterNum.h,analysis.groupDiff.filterNum.p,analysis.groupDiff.filterNum.ci,analysis.groupDiff.filterNum.stats] = ttest2(analysis.group1.filterNum.singleSubject,analysis.group2.filterNum.singleSubject);
        [analysis.groupDiff.accuracy.h,analysis.groupDiff.accuracy.p,analysis.groupDiff.accuracy.ci,analysis.groupDiff.accuracy.stats] = ttest2(analysis.group1.accuracy.singleSubject,analysis.group2.accuracy.singleSubject);
        [analysis.groupDiff.d1.h,analysis.groupDiff.d1.p,analysis.groupDiff.d1.ci,analysis.groupDiff.d1.stats] = ttest2(analysis.group1.d1.singleSubject,analysis.group2.d1.singleSubject);
        [analysis.groupDiff.c1.h,analysis.groupDiff.c1.p,analysis.groupDiff.c1.ci,analysis.groupDiff.c1.stats] = ttest2(analysis.group1.c1.singleSubject,analysis.group2.c1.singleSubject);
        [analysis.groupDiff.avgConfidence.h,analysis.groupDiff.avgConfidence.p,analysis.groupDiff.avgConfidence.ci,analysis.groupDiff.avgConfidence.stats] = ttest2(analysis.group1.avgConfidence.singleSubject,analysis.group2.avgConfidence.singleSubject);
    end

    % Save results
    save(resultsFile, 'analysis'); % Save results

    % Create Figure to display results
    figure
    set(gcf, 'Units', 'Normalized', 'OuterPosition', [0 0 0.6 0.6]);
    labels = ['Group 1'; 'Group 2'];
    % Plot Type 1 metrics on top line
    ax1 = subplot(2,4,1);
    bar(ax1, [analysis.group1.filterNum.groupMean; analysis.group2.filterNum.groupMean]);
    hold on
    errorbar(ax1, [analysis.group1.filterNum.groupMean; analysis.group2.filterNum.groupMean], [std(analysis.group1.filterNum.singleSubject), std(analysis.group2.filterNum.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.groupDiff.filterNum.h == 0
        title('FILTER NUMBER: NON-SIG DIFF')
    elseif analysis.groupDiff.filterNum.h > 0
        title('FILTER NUMBER: SIG DIFF')
    end
    hold off
    ax2 = subplot(2,4,2);
    bar(ax2, [analysis.group1.accuracy.groupMean; analysis.group2.accuracy.groupMean]);
    hold on
    errorbar(ax2, [analysis.group1.accuracy.groupMean; analysis.group2.accuracy.groupMean], [std(analysis.group1.accuracy.singleSubject), std(analysis.group2.accuracy.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.groupDiff.accuracy.h == 0
        title('ACCURACY: NON-SIG DIFF')
    elseif analysis.groupDiff.accuracy.h > 0
        title('ACCURACY: SIG DIFF')
    end
    ax3 = subplot(2,4,3);
    bar(ax3, [analysis.group1.d1.groupMean; analysis.group2.d1.groupMean]);
    hold on
    errorbar(ax3, [analysis.group1.d1.groupMean; analysis.group2.d1.groupMean], [std(analysis.group1.d1.singleSubject), std(analysis.group2.d1.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.groupDiff.d1.h == 0
        title('d'': NON-SIG DIFF')
    elseif analysis.groupDiff.d1.h > 0
        title('d'': SIG DIFF')
    end
    hold off
    ax4 = subplot(2,4,4);
    bar(ax4, [analysis.group1.c1.groupMean; analysis.group2.c1.groupMean]);
    hold on
    errorbar(ax4, [analysis.group1.c1.groupMean; analysis.group2.c1.groupMean], [std(analysis.group1.c1.singleSubject), std(analysis.group2.c1.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.groupDiff.c1.h == 0
        title('CRITERION: NON-SIG DIFF')
    elseif analysis.groupDiff.c1.h > 0
        title('CRITERION: SIG DIFF')
    end
    hold off
    % Plot group Mratio and the difference
    ax5 = subplot(2,4,5);
    histogram(ax5, exp(analysis.group1.fit.mcmc.samples.mu_logMratio))
    xlim([0 2])
    xlabel('MRatio');
    ylabel('Sample count');
    hold on
    histogram(ax5, exp(analysis.group2.fit.mcmc.samples.mu_logMratio))
    hold off
    lgd = legend('Group 1', 'Group 2', 'Location', 'northeast');
    lgd.FontSize = 10;
    if (analysis.groupDiff.Mratio.hdi(1) > 0 && analysis.groupDiff.Mratio.hdi(2) > 0) || (analysis.groupDiff.Mratio.hdi(1) < 0 && analysis.groupDiff.Mratio.hdi(2) < 0)
        title('MRATIO: SIG DIFF')
    else
        title('MRATIO: NON-SIG DIFF')
    end
    ax6 = subplot(2,4,6);
    histogram(ax6, analysis.group1.fit.mcmc.samples.mu_logMratio)
    xlabel('log(MRatio)');
    ylabel('Sample count');
    hold on
    histogram(ax6, analysis.group2.fit.mcmc.samples.mu_logMratio)
    hold off
    lgd = legend('Group 1', 'Group 2', 'Location', 'northwest');
    lgd.FontSize = 10;
    if (analysis.groupDiff.log_Mratio.hdi(1) > 0 && analysis.groupDiff.log_Mratio.hdi(2) > 0) || (analysis.groupDiff.log_Mratio.hdi(1) < 0 && analysis.groupDiff.log_Mratio.hdi(2) < 0)
        title('LOG(MRATIO): SIG DIFF')
    else
        title('LOG(MRATIO): NON-SIG DIFF')
    end
    ax7 = subplot(2,4,7);
    histogram(ax7, analysis.groupDiff.log_Mratio.sampleDiff)
    xlabel('log(MRatio)');
    ylabel('Sample count');
    if (analysis.groupDiff.log_Mratio.hdi(1) > 0 && analysis.groupDiff.log_Mratio.hdi(2) > 0) || (analysis.groupDiff.log_Mratio.hdi(1) < 0 && analysis.groupDiff.log_Mratio.hdi(2) < 0)
        title('LOG(MRATIO) DIFF: SIG DIFF')
    else
        title('LOG(MRATIO) DIFF: NON-SIG DIFF')
    end
    hold on
    ln2 = line([analysis.groupDiff.log_Mratio.hdi(1) analysis.groupDiff.log_Mratio.hdi(1)], [0 1800]);
    ln2.Color = 'r';
    ln2.LineWidth = 1.5;
    ln2.LineStyle = '--';
    ln2 = line([analysis.groupDiff.log_Mratio.hdi(2) analysis.groupDiff.log_Mratio.hdi(2)], [0 1800]);
    ln2.Color = 'r';
    ln2.LineWidth = 1.5;
    ln2.LineStyle = '--';
    hold off
    % Plot average confidence
    ax8 = subplot(2,4,8);
    bar(ax8, [analysis.group1.avgConfidence.groupMean; analysis.group2.avgConfidence.groupMean]);
    hold on
    errorbar(ax8, [analysis.group1.avgConfidence.groupMean; analysis.group2.avgConfidence.groupMean], [std(analysis.group1.avgConfidence.singleSubject), std(analysis.group2.avgConfidence.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.groupDiff.avgConfidence.h == 0
        title('CONFIDENCE: NON-SIG DIFF')
    elseif analysis.groupDiff.avgConfidence.h > 0
        title('CONFIDENCE: SIG DIFF')
    end
    hold off

    % Print figure
    figureFile = fullfile('analysis', ['filter_task_analysis_', analysis.type]); % Create figure file name
    print(figureFile, '-dtiff');

    % Display completion message on screen
    fprintf('\n________________________________________\n\n  COMPLETED GROUP DIFFERENCE ANALYSIS\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    fprintf('\n DIFF MRATIO = %.2f (HDI: %.2f to %.2f)\n', analysis.groupDiff.Mratio.mean, analysis.groupDiff.Mratio.hdi(1), analysis.groupDiff.Mratio.hdi(2));
    if (analysis.groupDiff.Mratio.hdi(1) > 0 && analysis.groupDiff.Mratio.hdi(2) > 0) || (analysis.groupDiff.Mratio.hdi(1) < 0 && analysis.groupDiff.Mratio.hdi(2) < 0)
        fprintf('HDI DOES NOT SPAN ZERO: MRATIO DIFF SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    else
        fprintf('  HDI SPANS ZERO: MRATIO DIFF NOT SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    end
    fprintf('\n        RESULTS CAN BE FOUND IN: \n             FDT/analysis/\n________________________________________\n');

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% RUN SESSION DIFFERENCE ANALYSIS (PAIRED) IF SPECIFIED
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if strcmp(analysis.type,'paired') == 1 % If session difference analysis specified

    % For two-session difference analysis: Print message about groups being fit together --> t-tests possible on meta-d' metrics
    fprintf('\nNOTE: Groups are fit using one hierarchical model.\nFrequentist statistics (such as t-tests) are possible on any metrics involving meta-d''.\nFrequentist statistics are applied to all type 1 metrics and average confidence values.\n');

    % Specify parameters
    analysis.sessionDiff.mcmc_params = fit_meta_d_params;

    % Fit each group using a separate hierarchical Bayesian model
    analysis.sessionDiff.fit = fit_meta_d_mcmc_groupCorr(analysis.trials2counts.nR_S1, analysis.trials2counts.nR_S2, analysis.sessionDiff.mcmc_params);

    % Compute HDI of difference for log(meta-d'/d')
    analysis.session1.logMratio.sessionMean = analysis.sessionDiff.fit.mu_logMratio(1);
    analysis.session1.logMratio.singleSubject = log(analysis.sessionDiff.fit.Mratio(:,1));
    analysis.session1.Mratio.sessionMean = exp(analysis.sessionDiff.fit.mu_logMratio(1));
    analysis.session1.Mratio.singleSubject = analysis.sessionDiff.fit.Mratio(:,1);
    analysis.session2.logMratio.sessionMean = analysis.sessionDiff.fit.mu_logMratio(2);
    analysis.session2.logMratio.singleSubject = log(analysis.sessionDiff.fit.Mratio(:,2));
    analysis.session2.Mratio.sessionMean = exp(analysis.sessionDiff.fit.mu_logMratio(2));
    analysis.session2.Mratio.singleSubject = analysis.sessionDiff.fit.Mratio(:,2);
    analysis.sessionDiff.log_Mratio.sampleDiff = analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,1) - analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,2);
    analysis.sessionDiff.log_Mratio.mean = analysis.session1.logMratio.sessionMean - analysis.session2.logMratio.sessionMean;
    analysis.sessionDiff.log_Mratio.hdi = calc_HDI(analysis.sessionDiff.log_Mratio.sampleDiff(:));
    analysis.sessionDiff.Mratio.mean = exp(analysis.session1.logMratio.sessionMean) - exp(analysis.session2.logMratio.sessionMean);
    analysis.sessionDiff.Mratio.hdi = calc_HDI((exp(analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,1)) - exp(analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,2))));

    % Pull out session mean values for all non-hierarchical metrics
    analysis.session1.d1.sessionMean = mean(analysis.sessionDiff.fit.d1(:,1));
    analysis.session1.d1.singleSubject = analysis.sessionDiff.fit.d1(:,1);
    analysis.session1.c1.sessionMean = mean(analysis.sessionDiff.fit.c1(:,1));
    analysis.session1.c1.singleSubject = analysis.sessionDiff.fit.c1(:,1);
    analysis.session2.d1.sessionMean = mean(analysis.sessionDiff.fit.d1(:,2));
    analysis.session2.d1.singleSubject = analysis.sessionDiff.fit.d1(:,2);
    analysis.session2.c1.sessionMean = mean(analysis.sessionDiff.fit.c1(:,2));
    analysis.session2.c1.singleSubject = analysis.sessionDiff.fit.c1(:,2);

    % Calculate average confidence, add filter number and accuracy to results for session 1
    for n = 1:length(analysis.session1.data)
        analysis.session1.avgConfidence.singleSubject(n) = mean(analysis.session1.data(n).results.thresholdTrials.confidence);
        analysis.session1.filterNum.singleSubject(n) = mean(analysis.session1.data(n).results.thresholdTrials.filterNum);
        analysis.session1.accuracy.singleSubject(n) = analysis.session1.data(n).results.thresholdTrials.accuracyTotal;
    end
    analysis.session1.avgConfidence.sessionMean = mean(analysis.session1.avgConfidence.singleSubject);
    analysis.session1.filterNum.sessionMean = mean(analysis.session1.filterNum.singleSubject);
    analysis.session1.accuracy.sessionMean = mean(analysis.session1.accuracy.singleSubject);

    % Calculate average confidence, add filter number and accuracy to results for session 2
    for n = 1:length(analysis.session2.data)
        analysis.session2.avgConfidence.singleSubject(n) = mean(analysis.session2.data(n).results.thresholdTrials.confidence);
        analysis.session2.filterNum.singleSubject(n) = mean(analysis.session2.data(n).results.thresholdTrials.filterNum);
        analysis.session2.accuracy.singleSubject(n) = analysis.session2.data(n).results.thresholdTrials.accuracyTotal;
    end
    analysis.session2.avgConfidence.sessionMean = mean(analysis.session2.avgConfidence.singleSubject);
    analysis.session2.filterNum.sessionMean = mean(analysis.session2.filterNum.singleSubject);
    analysis.session2.accuracy.sessionMean = mean(analysis.session2.accuracy.singleSubject);

    % Calculate session difference values for all non-hierarchical metrics
    analysis.sessionDiff.filterNum.meanDiff = analysis.session1.filterNum.sessionMean - analysis.session2.filterNum.sessionMean;
    analysis.sessionDiff.accuracy.meanDiff = analysis.session1.accuracy.sessionMean - analysis.session2.accuracy.sessionMean;
    analysis.sessionDiff.d1.meanDiff = analysis.session1.d1.sessionMean - analysis.session2.d1.sessionMean;
    analysis.sessionDiff.c1.meanDiff = analysis.session1.c1.sessionMean - analysis.session2.c1.sessionMean;
    analysis.sessionDiff.avgConfidence.meanDiff = analysis.session1.avgConfidence.sessionMean - analysis.session2.avgConfidence.sessionMean;

    % Run specified statistical test on all non-hierarchical metrics
    if strcmp(analysis.sessionDiff.test,'wilcoxon') == 1
        analysis.sessionDiff.testType = 'Wilcoxon signed rank test';
        [analysis.sessionDiff.filterNum.p,analysis.sessionDiff.filterNum.h,analysis.sessionDiff.filterNum.stats] = signrank(analysis.session1.filterNum.singleSubject,analysis.session2.filterNum.singleSubject);
        [analysis.sessionDiff.accuracy.p,analysis.sessionDiff.accuracy.h,analysis.sessionDiff.accuracy.stats] = signrank(analysis.session1.accuracy.singleSubject,analysis.session2.accuracy.singleSubject);
        [analysis.sessionDiff.d1.p,analysis.sessionDiff.d1.h,analysis.sessionDiff.d1.stats] = signrank(analysis.session1.d1.singleSubject,analysis.session2.d1.singleSubject);
        [analysis.sessionDiff.c1.p,analysis.sessionDiff.c1.h,analysis.sessionDiff.c1.stats] = signrank(analysis.session1.c1.singleSubject,analysis.session2.c1.singleSubject);
        [analysis.sessionDiff.avgConfidence.p,analysis.sessionDiff.avgConfidence.h,analysis.sessionDiff.avgConfidence.stats] = signrank(analysis.session1.avgConfidence.singleSubject,analysis.session2.avgConfidence.singleSubject);
    elseif strcmp(analysis.sessionDiff.test,'ttest') == 1
        analysis.sessionDiff.testType = 'Paired T-test';
        [analysis.sessionDiff.filterNum.h,analysis.sessionDiff.filterNum.p,analysis.sessionDiff.filterNum.ci,analysis.sessionDiff.filterNum.stats] = ttest(analysis.session1.filterNum.singleSubject,analysis.session2.filterNum.singleSubject);
        [analysis.sessionDiff.accuracy.h,analysis.sessionDiff.accuracy.p,analysis.sessionDiff.accuracy.ci,analysis.sessionDiff.accuracy.stats] = ttest(analysis.session1.accuracy.singleSubject,analysis.session2.accuracy.singleSubject);
        [analysis.sessionDiff.d1.h,analysis.sessionDiff.d1.p,analysis.sessionDiff.d1.ci,analysis.sessionDiff.d1.stats] = ttest(analysis.session1.d1.singleSubject,analysis.session2.d1.singleSubject);
        [analysis.sessionDiff.c1.h,analysis.sessionDiff.c1.p,analysis.sessionDiff.c1.ci,analysis.sessionDiff.c1.stats] = ttest(analysis.session1.c1.singleSubject,analysis.session2.c1.singleSubject);
        [analysis.sessionDiff.avgConfidence.h,analysis.sessionDiff.avgConfidence.p,analysis.sessionDiff.avgConfidence.ci,analysis.sessionDiff.avgConfidence.stats] = ttest(analysis.session1.avgConfidence.singleSubject,analysis.session2.avgConfidence.singleSubject);
    end

    % Save results
    save(resultsFile, 'analysis'); % Save results

    % Create Figure to display results
    figure
    set(gcf, 'Units', 'Normalized', 'OuterPosition', [0 0 0.6 0.6]);
    labels = ['Session 1'; 'Session 2'];
    % Plot Type 1 metrics on top line
    ax1 = subplot(2,4,1);
    bar(ax1, [analysis.session1.filterNum.sessionMean; analysis.session2.filterNum.sessionMean]);
    hold on
    errorbar(ax1, [analysis.session1.filterNum.sessionMean; analysis.session2.filterNum.sessionMean], [std(analysis.session1.filterNum.singleSubject), std(analysis.session2.filterNum.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.sessionDiff.filterNum.h == 0
        title('FILTER NUMBER: NON-SIG DIFF')
    elseif analysis.sessionDiff.filterNum.h > 0
        title('FILTER NUMBER: SIG DIFF')
    end
    hold off
    ax2 = subplot(2,4,2);
    bar(ax2, [analysis.session1.accuracy.sessionMean; analysis.session2.accuracy.sessionMean]);
    hold on
    errorbar(ax2, [analysis.session1.accuracy.sessionMean; analysis.session2.accuracy.sessionMean], [std(analysis.session1.accuracy.singleSubject), std(analysis.session2.accuracy.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.sessionDiff.accuracy.h == 0
        title('ACCURACY: NON-SIG DIFF')
    elseif analysis.sessionDiff.accuracy.h > 0
        title('ACCURACY: SIG DIFF')
    end
    ax3 = subplot(2,4,3);
    bar(ax3, [analysis.session1.d1.sessionMean; analysis.session2.d1.sessionMean]);
    hold on
    errorbar(ax3, [analysis.session1.d1.sessionMean; analysis.session2.d1.sessionMean], [std(analysis.session1.d1.singleSubject), std(analysis.session2.d1.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.sessionDiff.d1.h == 0
        title('d'': NON-SIG DIFF')
    elseif analysis.sessionDiff.d1.h > 0
        title('d'': SIG DIFF')
    end
    hold off
    ax4 = subplot(2,4,4);
    bar(ax4, [analysis.session1.c1.sessionMean; analysis.session2.c1.sessionMean]);
    hold on
    errorbar(ax4, [analysis.session1.c1.sessionMean; analysis.session2.c1.sessionMean], [std(analysis.session1.c1.singleSubject), std(analysis.session2.c1.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.sessionDiff.c1.h == 0
        title('CRITERION: NON-SIG DIFF')
    elseif analysis.sessionDiff.c1.h > 0
        title('CRITERION: SIG DIFF')
    end
    hold off
    % Plot session Mratio and the difference
    ax5 = subplot(2,4,5);
    histogram(ax5, exp(analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,1)))
    xlim([0 2])
    xlabel('MRatio');
    ylabel('Sample count');
    hold on
    histogram(ax5, exp(analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,2)))
    hold off
    lgd = legend('Session 1', 'Session 2', 'Location', 'northeast');
    lgd.FontSize = 10;
    if (analysis.sessionDiff.Mratio.hdi(1) > 0 && analysis.sessionDiff.Mratio.hdi(2) > 0) || (analysis.sessionDiff.Mratio.hdi(1) < 0 && analysis.sessionDiff.Mratio.hdi(2) < 0)
        title('MRATIO: SIG DIFF')
    else
        title('MRATIO: NON-SIG DIFF')
    end
    ax6 = subplot(2,4,6);
    histogram(ax6, analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,1))
    xlabel('log(MRatio)');
    ylabel('Sample count');
    hold on
    histogram(ax6, analysis.sessionDiff.fit.mcmc.samples.mu_logMratio(:,:,2))
    hold off
    lgd = legend('Session 1', 'Session 2', 'Location', 'northwest');
    lgd.FontSize = 10;
    if (analysis.sessionDiff.log_Mratio.hdi(1) > 0 && analysis.sessionDiff.log_Mratio.hdi(2) > 0) || (analysis.sessionDiff.log_Mratio.hdi(1) < 0 && analysis.sessionDiff.log_Mratio.hdi(2) < 0)
        title('LOG(MRATIO): SIG DIFF')
    else
        title('LOG(MRATIO): NON-SIG DIFF')
    end
    ax7 = subplot(2,4,7);
    histogram(ax7, analysis.sessionDiff.log_Mratio.sampleDiff)
    xlabel('log(MRatio)');
    ylabel('Sample count');
    if (analysis.sessionDiff.log_Mratio.hdi(1) > 0 && analysis.sessionDiff.log_Mratio.hdi(2) > 0) || (analysis.sessionDiff.log_Mratio.hdi(1) < 0 && analysis.sessionDiff.log_Mratio.hdi(2) < 0)
        title('LOG(MRATIO) DIFF: SIG DIFF')
    else
        title('LOG(MRATIO) DIFF: NON-SIG DIFF')
    end
    hold on
    ln2 = line([analysis.sessionDiff.log_Mratio.hdi(1) analysis.sessionDiff.log_Mratio.hdi(1)], [0 1800]);
    ln2.Color = 'r';
    ln2.LineWidth = 1.5;
    ln2.LineStyle = '--';
    ln2 = line([analysis.sessionDiff.log_Mratio.hdi(2) analysis.sessionDiff.log_Mratio.hdi(2)], [0 1800]);
    ln2.Color = 'r';
    ln2.LineWidth = 1.5;
    ln2.LineStyle = '--';
    hold off
    % Plot average confidence
    ax8 = subplot(2,4,8);
    bar(ax8, [analysis.session1.avgConfidence.sessionMean; analysis.session2.avgConfidence.sessionMean]);
    hold on
    errorbar(ax8, [analysis.session1.avgConfidence.sessionMean; analysis.session2.avgConfidence.sessionMean], [std(analysis.session1.avgConfidence.singleSubject), std(analysis.session2.avgConfidence.singleSubject)], 'o', 'marker', 'none', 'linewidth', 1, 'Color','k');
    set(gca,'XTickLabel',labels);
    set(gca,'XTickLabelRotation',90)
    if analysis.sessionDiff.avgConfidence.h == 0
        title('CONFIDENCE: NON-SIG DIFF')
    elseif analysis.sessionDiff.avgConfidence.h > 0
        title('CONFIDENCE: SIG DIFF')
    end
    hold off

    % Print figure
    figureFile = fullfile('analysis', ['filter_task_analysis_', analysis.type]); % Create figure file name
    print(figureFile, '-dtiff');

    % Display completion message on screen
    fprintf('\n________________________________________\n\n  COMPLETED SESSION DIFFERENCE ANALYSIS\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    fprintf('\n DIFF MRATIO = %.2f (HDI: %.2f to %.2f)\n', analysis.sessionDiff.Mratio.mean, analysis.sessionDiff.Mratio.hdi(1), analysis.sessionDiff.Mratio.hdi(2));
    if (analysis.sessionDiff.Mratio.hdi(1) > 0 && analysis.sessionDiff.Mratio.hdi(2) > 0) || (analysis.sessionDiff.Mratio.hdi(1) < 0 && analysis.sessionDiff.Mratio.hdi(2) < 0)
        fprintf('HDI DOES NOT SPAN ZERO: MRATIO DIFF SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    else
        fprintf('  HDI SPANS ZERO: MRATIO DIFF NOT SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    end
    fprintf('\n        RESULTS CAN BE FOUND IN: \n             FDT/analysis/\n________________________________________\n');

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% RUN GROUP REGRESSION ANALYSIS IF SPECIFIED
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if strcmp(analysis.type,'regress') == 1 % If group regression analysis specified

    % Specify parameters
    analysis.groupRegression.mcmc_params = fit_meta_d_params;

    % Standardise covariate
    analysis.groupRegression.cov = zscore(analysis.covariate.data);

    % Fit group data all at once
    analysis.groupRegression.fit = fit_meta_d_mcmc_regression(analysis.trials2counts.nR_S1, analysis.trials2counts.nR_S2, analysis.groupRegression.cov, analysis.groupRegression.mcmc_params);

    % Calculate log of Mratio for single subjects values
    for n = 1:length(analysis.groupRegression.fit.Mratio)
        analysis.groupRegression.log_Mratio.singleSubject(n) = log(analysis.groupRegression.fit.Mratio(n));
    end

    % Pull out group and single subject values
    analysis.groupRegression.d1.groupMean = mean(analysis.groupRegression.fit.d1);
    analysis.groupRegression.d1.singleSubject = analysis.groupRegression.fit.d1;
    analysis.groupRegression.c1.groupMean = mean(analysis.groupRegression.fit.c1);
    analysis.groupRegression.c1.singleSubject = analysis.groupRegression.fit.c1;
    analysis.groupRegression.meta_d.groupMean = mean(analysis.groupRegression.fit.meta_d);
    analysis.groupRegression.meta_d.singleSubject = analysis.groupRegression.fit.meta_d;
    analysis.groupRegression.Mratio.groupMean = exp(analysis.groupRegression.fit.mu_logMratio);
    analysis.groupRegression.Mratio.singleSubject = analysis.groupRegression.fit.Mratio;
    analysis.groupRegression.Mratio.hdi = calc_HDI(exp(analysis.groupRegression.fit.mcmc.samples.mu_logMratio(:)));
    analysis.groupRegression.log_Mratio.groupMean = analysis.groupRegression.fit.mu_logMratio;
    analysis.groupRegression.log_Mratio.hdi = calc_HDI(analysis.groupRegression.fit.mcmc.samples.mu_logMratio(:));
    analysis.groupRegression.beta1.groupMean = analysis.groupRegression.fit.mu_beta1;
    analysis.groupRegression.beta1.hdi = calc_HDI(analysis.groupRegression.fit.mcmc.samples.mu_beta1(:));

    % Calculate average confidence, add filter number and accuracy to analysis results
    for n = 1:length(analysis.data)
        analysis.groupRegression.avgConfidence.singleSubject(n) = mean(analysis.data(n).results.thresholdTrials.confidence);
        analysis.groupRegression.filterNum.singleSubject(n) = mean(analysis.data(n).results.thresholdTrials.filterNum);
        analysis.groupRegression.accuracy.singleSubject(n) = analysis.data(n).results.thresholdTrials.accuracyTotal;
    end
    analysis.groupRegression.avgConfidence.groupMean = mean(analysis.groupRegression.avgConfidence.singleSubject);
    analysis.groupRegression.filterNum.groupMean = mean(analysis.groupRegression.filterNum.singleSubject);
    analysis.groupRegression.accuracy.groupMean = mean(analysis.groupRegression.accuracy.singleSubject);

    % Save results
    save(resultsFile, 'analysis'); % Save results

    % Display completion message on screen plus results
    fprintf('\n________________________________________\n\n  COMPLETED GROUP REGRESSION ANALYSIS\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    fprintf('\n GROUP MRATIO = %.2f (HDI: %.2f to %.2f)\n', analysis.groupRegression.Mratio.groupMean, analysis.groupRegression.Mratio.hdi(1), analysis.groupRegression.Mratio.hdi(2));
    if analysis.groupRegression.Mratio.hdi(1) > 0
        fprintf('   HDI DOES NOT SPAN ZERO: MRATIO SIG.\n');
    else
        fprintf('   HDI SPANS ZERO: MRATIO NOT SIG.\n');
    end
    fprintf('\n GROUP BETA = %.2f (HDI: %.2f to %.2f)\n', analysis.groupRegression.beta1.groupMean, analysis.groupRegression.beta1.hdi(1), analysis.groupRegression.beta1.hdi(2));
    if analysis.groupRegression.beta1.hdi(1) > 0
        fprintf('    HDI DOES NOT SPAN ZERO: BETA SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    else
        fprintf('      HDI SPANS ZERO: BETA NOT SIG.\n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _\n');
    end
    fprintf('\n        RESULTS CAN BE FOUND IN: \n             FDT/analysis/\n________________________________________\n');

end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


function [model,output] = normalizeExponentialCone(model)

% N.B: YALMIP Definition % x2*exp(x1/x2)  <= x3

% Temporarily extract the data in a more standard name min c'*y, b + Ay >= 0
output.problem = 0;
if ~isempty(model.evalMap)

    % Temporarily extract the data in a more standard name min c'*y, b + Ay >= 0
    data.A = model.F_struc(:,2:end);
    data.b = full(model.F_struc(:,1));
    data.c = full(model.c);
    cones = model.K;
    % Clear away the SDP cone to easier add new rows. Will be put back in
    % place later
    data.A(end-sum(model.K.s.^2)+1:end,:)=[];
    data.b(end-sum(model.K.s.^2)+1:end) = [];
    sdpData = model.F_struc(end-sum(model.K.s.^2)+1:end,:);

    % First check that we have only exponentials
    convexFunctions = [];
    concaveFunctions = [];
    vectorentropy = [];
    vectorkullback = [];
    vectorlogsumexp = [];
    for i = 1:length(model.evalMap)
        switch model.evalMap{i}.fcn
            case {'exp','pexp'}
                convexFunctions = [convexFunctions model.evalMap{i}.computes];
            case {'log','plog','slog'}
                concaveFunctions = [concaveFunctions model.evalMap{i}.computes];
            case 'entropy'
                concaveFunctions = [concaveFunctions model.evalMap{i}.computes];
                if length(model.evalMap{i}.variableIndex) > 1
                    vectorentropy = [vectorentropy i];
                end
            case 'kullbackleibler'
                convexFunctions = [convexFunctions model.evalMap{i}.computes];
                if length(model.evalMap{i}.variableIndex) > 2
                    vectorkullback = [vectorkullback i];
                end
            case 'logsumexp'
                convexFunctions = [convexFunctions model.evalMap{i}.computes];
                vectorlogsumexp = [vectorlogsumexp i];
            otherwise
                % Standard interface, return solver not applicable
                % This should not be able to happen as we check this
                % earlier
                output = createoutput([],[],[],-4,model.solver.tag,[],[],0);
                return
        end
    end
    % Check that all exp/log enter in a convex fashion
    if model.K.f > 0
        if nnz(data.A(1:model.K.f,convexFunctions))>0 || nnz(data.A(1:model.K.f,concaveFunctions))>0
            output = createoutput([],[],[],19,model.solver.tag,[],[],0);
            return
        end
    end
    % Check sign in objective on exp/log terms
    if any(data.c(convexFunctions) < 0) || any(data.c(concaveFunctions) > 0)
        output = createoutput([],[],[],19,model.solver.tag,[],[],0);
        return
    end
    % Check sign in elementwise inequalities on exp/log terms
    if model.K.l > 0
        if nnz(data.A(1+model.K.f:model.K.f+model.K.l,convexFunctions)>0) || nnz(data.A(1+model.K.f:model.K.f+model.K.l,concaveFunctions)<0)
            output = createoutput([],[],[],19,model.solver.tag,[],[],0);
            return
        end
    end
    % Check so there are no exp/log terms in quadratic or SDP cone
    if sum(model.K.q) + sum(model.K.s) > 0
        if nnz(data.A(1+model.K.f+model.K.l:end,convexFunctions))>0 || nnz(data.A(1+model.K.f+model.K.l:end,concaveFunctions))>0
            output = createoutput([],[],[],19,model.solver.tag,[],[],0);
            return
        end
    end
    % Scalarize entropy operator
    if ~isempty(vectorentropy)
        for i = vectorentropy(:)'
            involved = model.evalMap{i}.variableIndex;
            computes = model.evalMap{i}.computes;
            n = length(data.c);
            m = length(involved);
            data.A = [sparse(ones(1,m+1),[computes n+(1:m)],[-1 ones(1,m)],1,n+m);
                      data.A spalloc(size(data.A,1),m,0)];
            data.b = [0;data.b];
            data.c = [data.c;zeros(m,1)];
            cones.f = cones.f + 1;
            % The original strucuture now just relates to the first new term
            model.evalMap{i}.computes = n+1;
            model.evalMap{i}.variableIndex = involved(1);
            % and then we add some
            for j = 2:length(involved)
                model.evalMap{end+1} = model.evalMap{i};
                model.evalMap{end}.variableIndex = involved(j);
                model.evalMap{end}.computes = n+j;
            end
        end
    end

    % Scalarize kullbackleibler operator
    if ~isempty(vectorkullback)
        for i = vectorkullback(:)'
            involved = model.evalMap{i}.variableIndex;
            computes = model.evalMap{i}.computes;
            n = length(data.c);
            m = length(involved)/2;
            data.A = [sparse(ones(1,m+1),[computes n+(1:m)],[-1 ones(1,m)],1,n+m);
                      data.A spalloc(size(data.A,1),m,0)];
            data.b = [0;data.b];
            data.c = [data.c;zeros(m,1)];
            cones.f = cones.f + 1;
            % The original strucuture now just relates to the first new term
            model.evalMap{i}.computes = n+1;
            model.evalMap{i}.variableIndex = involved([1 m+1]);
            % and then we add some
            for j = 2:length(involved)/2
                model.evalMap{end+1} = model.evalMap{i};
                model.evalMap{end}.variableIndex = involved([j j+m]);
                model.evalMap{end}.computes = n+j;
            end
        end
    end

    % Scalarize logsumexp operator
    % write log(expx1+expx2..)<= z as exp(x1-z)+exp(x2-z)+...<=1
    if ~isempty(vectorlogsumexp)
        for i = vectorlogsumexp(:)'
            involved = model.evalMap{i}.variableIndex;
            computes = model.evalMap{i}.computes;
            % Orignal #variables
            n = length(data.c);
            % Number of terms in logsumexp
            m = length(involved);

            % exp(x1-z)+exp(x2-z)+...<=1
            data.A = [data.A(1:cones.f,:) spalloc(cones.f,m,0);
                      spalloc(1,n,0) -ones(1,m);
                      data.A(cones.f+1:end,:) spalloc(cones.l+sum(cones.q)+sum(cones.s),m,0)];
            data.b = [data.b(1:cones.f);
                      1;
                      data.b(cones.f+1:end)];
            data.c = [data.c;zeros(m,1)];
            cones.l = cones.l + 1;

            % The original strucuture now just relates to the first new term
            model.evalMap{i}.computes = n+1;
            model.evalMap{i}.variableIndex = [involved(1) computes];
            model.evalMap{i}.fcn = 'expdiff';
            % and then we add some new
            for j = 2:length(involved)
                model.evalMap{end+1} = model.evalMap{i};
                model.evalMap{end}.variableIndex = [involved(j) computes];
                model.evalMap{end}.computes = n+j;
            end
        end
    end

    % Describe all exponential cones
    m = length(model.evalMap);
    cones.e = model.K.e + m;
    for i = 1:m
        switch model.evalMap{i}.fcn
            case 'exp'
                % 1*exp(xv/1) <= xc
                % y*exp(x/y)  <= z.  y new variable, xv the original variable, and xc the "computed"
                x = model.evalMap{i}.variableIndex;
                z = model.evalMap{i}.computes;
                data.A = [data.A;
                          -sparse([1;3],[x z],-1,3,size(data.A,2))];
                data.b = [data.b;[0;1;0]];
            case 'pexp'
                % xv(1)*exp(xv(2)/xv(1)) <= xc
                x = model.evalMap{i}.variableIndex(2);
                y = model.evalMap{i}.variableIndex(1);
                z = model.evalMap{i}.computes;
                data.A = [data.A;
                          -sparse([1;2;3],[x y z],[-1 -1 -1],3,size(data.A,2))];
                data.b = [data.b;[0;0;0]];
            case 'log'
                % log(xv) >= xc i.e. xv >= exp(xc/1)*1
                z = model.evalMap{i}.variableIndex;
                x = model.evalMap{i}.computes;
                data.A = [data.A;
                          -sparse([1;3],[x z],-1,3,size(data.A,2))];
                data.b = [data.b;[0;1;0]];
            case 'plog'
                % xv(1)log(xv(2)/xv(1))>=xc i.e.
                % -xc >= -xv(1)log(xv(2)/xv(1))
                z = model.evalMap{i}.computes;
                y = model.evalMap{i}.variableIndex(1);
                x = model.evalMap{i}.variableIndex(2);
                data.A = [data.A;
                          -sparse([1;2;3],[x y z],[1 -1 1],3,size(data.A,2))];
                data.b = [data.b;[0;0;0]];
            case 'slog'
                % log(1+xv) >= xc i.e. (1+xv) >= exp(xc/1)*1
                z = model.evalMap{i}.variableIndex;
                x = model.evalMap{i}.computes;
                data.A = [data.A;
                          -sparse([1;3],[x z],-1,3,size(data.A,2))];
                data.b = [data.b;0;1;1];
            case 'entropy'
                % -xv*log(xv)>=xc i.e. 1 >= exp(xc/xv)*xv
                x = model.evalMap{i}.computes;
                y = model.evalMap{i}.variableIndex;
                data.A = [data.A;
                          -sparse([1;2],[x y],[-1 -1],3,size(data.A,2))];
                data.b = [data.b;[0;0;1]];
            case 'kullbackleibler'
                % -xv1*log(xv1/xv2)>=-xc i.e. xv2 >= exp(-xc/xv1)*xv1
                x = model.evalMap{i}.computes;
                y = model.evalMap{i}.variableIndex(1);
                z = model.evalMap{i}.variableIndex(2);
                data.A = [data.A;
                          -sparse([1;2;3],[x y z],[1 -1 -1],3,size(data.A,2))];
                data.b = [data.b;[0;0;0]];
            case 'expdiff'
                % exp(xv(1)-xv(2)) <= xc
                x = model.evalMap{i}.variableIndex;
                z = model.evalMap{i}.computes;
                data.A = [data.A;
                          -sparse([1;1;3],[x(1) x(2) z],[-1 1 -1],3,size(data.A,2))];
                data.b = [data.b;[0;1;0]];
            otherwise
        end
    end
    model.F_struc = [data.b data.A];
    model.c = data.c;
    model.K = cones;
    % Put back the SDP cone in place
    if model.K.s(1)>0
        if size(sdpData,2) < size(model.F_struc,2)
            sdpDAta(end,size(model.F_struc,2)) = 0;
        end
        model.F_struc = [model.F_struc;sdpData];
    end
end
n = size(model.F_struc,2)-1;
if n > length(model.lb)
    missing = n - length(model.lb);
    model.lb = [model.lb;-inf(missing,1)];
end
if size(model.F_struc,2)-1 > length(model.ub)
    missing = n - length(model.ub);
    model.ub = [model.ub;inf(missing,1)];
end


function F = coth(F, varargin)
%COTH   Hyperbolic cotangent of a CHEBFUN.
%   COTH(F) computes the hyperbolic cotangent of the CHEBFUN F.
%
%   COTH(F, PREF) does the same but uses the CHEBFUNPREF object PREF when
%   computing the composition.
%
% See also ACOTH.

% Copyright 2017 by The University of Oxford and The Chebfun Developers.
% See http://www.chebfun.org/ for Chebfun information.

% Call the compose method:
F = compose(F, @coth, varargin{:});

end


function f = restrict(f, dom)
% RESTRICT    Restrict the domain of a CHEBFUN3.
%
%   F = RESTRICT(F, DOM) returns a
%   - CHEBFUN3 on the domain DOM that approximates F on that domain if DOM
%     is a nondegenarate cuboid,
%   - CHEBFUN2 that approximates F on the domain if DOM is a plane, and
%   - CHEBFUN that approximates F on the domain if DOM is a line.
%   DOM should be a vector of length 6 specifying the underlying cuboid.

% Copyright 2017 by The University of Oxford and The Chebfun Developers.
% See http://www.chebfun.org/ for Chebfun information.

if ( isa(dom, 'double') )
    if ( numel(dom) == 6 )                   % Restrict to DOM.
        xlen = diff(dom(1:2));
        ylen = diff(dom(3:4));
        zlen = diff(dom(5:6));

        if ( ( xlen == 0 ) && ( ylen == 0) && ( zlen == 0 ) )
            % DOM is a point.
            f = feval(f, dom(1), dom(3), dom(5));

        elseif ( ( xlen ~= 0 ) && ( ylen == 0 ) && ( zlen == 0 ) )
            % DOM is a vertical line (corresponding to a column).
            cols = restrict(f.cols, dom(1:2));
            rows = feval(f.rows, dom(3));
            tubes = feval(f.tubes, dom(5));
            core = chebfun3.txm(chebfun3.txm(f.core, rows, 2), tubes, 3);
            f = chebfun(cols*core);

        elseif ( ( xlen == 0 ) && ( ylen ~= 0 ) && ( zlen == 0 ) )
            % DOM is a horizontal line (corresponding to a row).
            cols = feval(f.cols, dom(1));
            rows = restrict(f.rows, dom(3:4));
            tubes = feval(f.tubes, dom(5));
            core = chebfun3.txm(chebfun3.txm(f.core, cols, 1), tubes, 3);
            f = chebfun(rows*core.');

        elseif ( ( xlen == 0 ) && ( ylen == 0 ) && ( zlen ~= 0 ) )
            % DOM is an oblique line (corresponding to a tube).
            cols = feval(f.cols, dom(1));
            rows = feval(f.rows, dom(3));
            tubes = restrict(f.tubes, dom(5:6));
            core = squeeze(chebfun3.txm(chebfun3.txm(f.core, cols, 1), ...
                rows, 2));
            f = chebfun(tubes*core);

        elseif ( ( xlen == 0 ) && ( ylen ~= 0 ) && ( zlen ~= 0 ) )
            % DOM is a horizontal plane (corresponding to a horizontal slice).
            cols = feval(f.cols, dom(1));
            rows = restrict(f.rows, dom(3:4));
            tubes = restrict(f.tubes, dom(5:6));
            core = squeeze(chebfun3.txm(f.core, cols, 1));
            f = chebfun2(tubes*core.'*rows.');

        elseif ( ( xlen ~= 0 ) && ( ylen == 0 ) && ( zlen ~= 0 ) )
            % DOM is a lateral plane (corresponding to a lateral slice).
            cols = restrict(f.cols, dom(1:2));
            rows = feval(f.rows, dom(3));
            tubes = restrict(f.tubes, dom(5:6));
            core = squeeze(chebfun3.txm(f.core, rows, 2));
            f = chebfun2(tubes*core.'*cols.');

        elseif ( ( xlen ~= 0 ) && ( ylen ~= 0 ) && ( zlen == 0 ) )
            % DOM is a frontal plane (corresponding to a frontal slice).
            cols = restrict(f.cols, dom(1:2));
            rows = restrict(f.rows, dom(3:4));
            tubes = feval(f.tubes,dom(5));
            core = chebfun3.txm(f.core, tubes, 3);
            f = chebfun2(rows*core.'*cols.');

        else
            % DOM is not degenerate
            f.cols = restrict(f.cols, dom(1:2));
            f.rows = restrict(f.rows, dom(3:4));
            f.tubes = restrict(f.tubes, dom(5:6));
            f.domain = dom;
        end
    else
        error('CHEBFUN:CHEBFUN3:restrict:domain', 'Domain not determined.');
    end

else
    error('CHEBFUN:CHEBFUN3:restrict:domain', 'Unrecognizable domain.');
end

end


function f = constructor(f, op, varargin)
%CONSTRUCTOR   Main CHEBFUN3T constructor.
%   Classical full tensor approach for 3D functions. No low-rank technique
%   is involved here. This is experimental code, not for ordinary use.
%
% See also CHEBFUN3.

%   The 'trig' flag is NOT implemented.

% Copyright 2017 by The University of Oxford and The Chebfun Developers.
% See http://www.chebfun.org/ for Chebfun information.

% Parse the inputs:
[op, dom, pref, vectorize] = parseInputs(op, varargin{:});

% Set preferences:
tech        = pref.tech();
tpref       = tech.techPref;
grid = tpref.minSamples;
maxSample = 363;
% Since 363*363*363=(6916)^2, this is equivalent to working with a matrix
% of size 6916 * 6916 which needs 380MB of memory.
pseudoLevel = pref.cheb3Prefs.chebfun3eps;
prefStruct = pref.cheb2Prefs;
passSampleTest = prefStruct.sampleTest;

[xx, yy, zz] = ndgrid(dom(1:2), dom(3:4), dom(5:6));
A = op(xx, yy, zz);
if ( isscalar(A) )
    op = @(x,y,z) op(x,y,z) + 0*x + 0*y + 0*z;
end

%%
isHappy = 0;
m = grid; n = grid; p = grid;
for i=1:10
    %% Compute values at 3D chebpts
    out = tech.tensorGrid([m, n, p], dom);
    xx = out{1};
    yy = out{2};
    zz = out{3};
    F = op(xx,yy,zz);

    % Does the function blow up or evaluate to NaN?:
    vscale = max(abs(F(:)));
    if ( isinf(vscale) )
        error('CHEBFUN:CHEBFUN3T:constructor:inf', ...
            'Function returned INF when evaluated');
    elseif ( any(isnan(F(:)) ) )
        error('CHEBFUN:CHEBFUN3T:constructor:nan', ...
            'Function returned NaN when evaluated');
    end

    %% Call 3D vals2coeffs
    coeffs3D = chebfun3t.vals2coeffs(F);
    % happinessCheck
    [isHappyX, isHappyY, isHappyZ, cutoffX2, cutoffY2, cutoffZ2] = ...
        happinessCheck3D(coeffs3D, pref);
    isHappy = isHappyX & isHappyY & isHappyZ;

    %% PHASE II
    % This phase refines the grid ONLY in "unhappy" directions.
    % We don't refine here if ALL the directions are nonhappy (as mentioned
    % in the following). If, however, less than 3 directions are nonhappy,
    % we refine in those directions only.

    m2 = m; n2 = n; p2 = p;
    while (~isHappy)
        % We need to refine the grid. We refine in 2 different ways: If all
        % 3 directions are unhappy, the new grid is bigger in each
        % direction by a factor of sqrt(2). If 1 or 2 of the directions are
        % unhappy, we refine (only those unhappy directions) by a factor of
        % 2.

        if (~isHappyX && ~isHappyY && ~isHappyZ)
            % All directions are unhappy. To avoid getting an unncecessarily
            % big tensor, refine differently.
            break;
        end
        if (~isHappyX)
            m2 = gridRefinePhase2(m2, pref);
        end
        if (~isHappyY)
            n2 = gridRefinePhase2(n2, pref);
        end
        if (~isHappyZ)
            p2 = gridRefinePhase2(p2, pref);
        end
        out = tech.tensorGrid([m2, n2, p2], dom);
        xx = out{1};
        yy = out{2};
        zz = out{3};
        F = op(xx,yy,zz);
        coeffs3D = chebfun3t.vals2coeffs(F);
        [isHappyX, isHappyY, isHappyZ, cutoffX2, cutoffY2, cutoffZ2] = ...
            happinessCheck3D(coeffs3D,pref);
        isHappy = isHappyX & isHappyY & isHappyZ;
    end

     if ( isHappy )
         % Chop the tail:
         coeffs3D = coeffs3D(1:cutoffX2, 1:cutoffY2, 1:cutoffZ2);

         % Construct a CHEBFUN3T object:
         f.coeffs = coeffs3D;
         f.vscale = max(abs(F(:)));
         f.domain = dom;

         % Step 2: Apply a SAMPLETEST:
         [m2, n2, p2] = size(coeffs3D);
         if ( passSampleTest )
             tol2 = getTol3D(xx, yy, zz, F, length(xx), dom, pseudoLevel);
             pass = sampleTest(f, op, tol2, vectorize);
             if ( pass ) % :)
                 break
             end
         end
     end
    m = gridRefinePhase1(m);
    n = gridRefinePhase1(n);
    p = gridRefinePhase1(p);
end

% Reached maximum allowed tensor size and still unhappy?
if ( i == 10 && ~isHappy )
    % Throw a warning and construct from the latest unhappy approximation
    warning('CHEBFUN:CHEBFUN3T:constructor:notResolved', ...
                'Unresolved with maximum CHEBFUN3T length: %u.', maxSample);
    f.coeffs = coeffs3D;
    f.vscale = max(abs(F(:)));
    f.domain = dom;
end

end

%%
function [isHappyX, isHappyY, isHappyZ, cutoffX2, cutoffY2, cutoffZ2] = ...
    happinessCheck3D(simple_3D_coeffs, pref)
% An alternative technique: check whether all the rows, columns and tubes have been resolved.
% Somewhat similar to Chebfun2 constructor (but with all the rows, cols and tubes instead of the pivot ones)

% TODO: THIS DOES NOT PROPERLY WORK WITH TRIGs.

colChebtech = sum(chebfun3t.unfold(abs(simple_3D_coeffs), 1), 2);
fCol = chebtech2({[], colChebtech});
[isHappyX, cutoffX2] = happinessCheck(fCol, [], [], [], pref);

rowChebtech = sum(chebfun3t.unfold(abs(simple_3D_coeffs), 2), 2);
fRow = chebtech2({[], rowChebtech});
[isHappyY, cutoffY2] = happinessCheck(fRow, [], [], [], pref);

tubeChebtech = sum(chebfun3t.unfold(abs(simple_3D_coeffs), 3), 2);
fTube = chebtech2({[], tubeChebtech});
[isHappyZ, cutoffZ2] = happinessCheck(fTube, [], [], [], pref);
end
%%
function grid = gridRefinePhase1(grid)
% Grid refinement strategy in Phase 1. It does the same for all TECHS.
grid = floor(sqrt(2)^(floor(2*log2(grid)) + 1)) + 1;
end

%%
function grid = gridRefinePhase2(grid, pref)
% Grid refinement strategy for tech in Phase 2

% What tech am I based on?:
tech = pref.tech();

% What is the next grid size?
if ( isa(tech, 'chebtech2') )
    % Double sampling on tensor grid:
    grid = 2*grid-1;
elseif ( isa(tech, 'trigtech') )
    % Double sampling on tensor grid:
    grid = 2^(floor(log2(grid)+1));
elseif ( isa(tech, 'chebtech1') )
    grid = 3 * grid;
else
    error('CHEBFUN:CHEBFUN3:constructor:gridRefinePhase2:techType', ...
        'Technology is unrecognized.');
end

end
%%

function tol = getTol3D(xx, yy, zz, vals, grid, dom, pseudoLevel)
%See https://github.com/chebfun/chebfun/issues/1491
[m,n,p] = size(vals);
% Remove some edge values so that df_dx, df_dy and df_dz have the same size.
df_dx = diff(vals(:,1:n-1,1:p-1),1,1) ./ diff(xx(:,1:n-1,1:p-1),1,1); % xx changes in the first mode.
df_dy = diff(vals(1:m-1,:,1:p-1),1,2) ./ diff(yy(1:m-1,:,1:p-1),1,2); % yy changes row-wise (2nd mode).
df_dz = diff(vals(1:m-1,1:n-1,:),1,3) ./ diff(zz(1:m-1,1:n-1,:),1,3); % zz changes tube-wise (3rd mode).

J = max(max(abs(df_dx),abs(df_dy)), abs(df_dz));
Jac_norm = max(J(:)); % An approximation for the norm of the Jacobian over the whole domain.
vscale = max(abs(vals(:)));
tol = grid^(4/5) * max( abs(dom(:) ) ) * max(Jac_norm, vscale) * pseudoLevel;
end
%%
function [op, dom, pref, vectorize] = parseInputs(op, varargin)
vectorize = 0;
pref = chebfunpref();

if ( isa(op, 'char') )     % CHEBFUN3T( CHAR )
    op = str2op(op);
end

if nargin == 1
    dom = [-1 1 -1 1 -1 1];
elseif ( ( (nargin == 2) && ~any(strcmpi(varargin{1}, {'trig', 'periodic'})) ))
    % DOMAIN is specified
    dom = varargin{1};
elseif ( ( (nargin == 2) && any(strcmpi(varargin{1}, {'trig', 'periodic'})) ))
    % TECH is specified
    pref.tech = @trigtech;
    dom = [-1 1 -1 1 -1 1];
elseif ( ( (nargin == 3) && any(strcmpi(varargin{2}, {'trig', 'periodic'})) ))
    % DOMAIN and TECH are specified
    dom = varargin{1};
    pref.tech = @trigtech;
elseif ( ( (nargin == 3) && any(strcmpi(varargin{1}, {'trig', 'periodic'})) ))
    % TECH and DOMAIN are specified
    pref.tech = @trigtech;
    dom = varargin{3};
elseif ( (nargin == 3) && strcmpi(varargin{1}, 'eps') )
    % EPS is specified
    pref.chebfuneps = varargin{2};
    dom = [-1 1 -1 1 -1 1];
elseif ( (nargin == 4) && strcmpi(varargin{1}, 'eps') ...
        && any(strcmpi(varargin{3}, {'trig', 'periodic'})) )
    % EPS and TECH are specified
    pref.tech = @trigtech;
    pref.chebfuneps = varargin{3};
    dom = [-1 1 -1 1 -1 1];
elseif ( (nargin == 4) && strcmpi(varargin{2}, 'eps') && ...
        any(strcmpi(varargin{1}, {'trig', 'periodic'})) )
    % TECH and EPS are specified
    pref.tech = @trigtech;
    pref.chebfuneps = varargin{3};
    dom = [-1 1 -1 1 -1 1];
elseif ( (nargin == 4) && strcmpi(varargin{2}, 'eps') && ...
        ~any(strcmpi(varargin{1}, {'trig', 'periodic'})) )
    % DOMAIN and EPS are specified
    pref.chebfuneps = varargin{3};
    dom = varargin{1};
elseif ( (nargin == 4) && strcmpi(varargin{1}, 'eps') && ...
        ~any(strcmpi(varargin{3}, {'trig', 'periodic'})) )
    % EPS and DOMAIN are specified
    pref.chebfuneps = varargin{2};
    dom = varargin{3};
elseif ( (nargin == 5) && any(strcmpi(varargin{2}, {'trig', 'periodic'}))...
        &&  strcmpi(varargin{3}, 'eps') )
    % DOMAIN, TECH and EPS are specified
    dom = varargin{1};
    pref.tech = @trigtech;
    pref.chebfuneps = varargin{4};
elseif ( (nargin == 5) && strcmpi(varargin{2}, 'eps') ...
        && any(strcmpi(varargin{4}, {'trig', 'periodic'})) )
    % DOMAIN, EPS and TECH are specified
    dom = varargin{1};
    pref.chebfuneps = varargin{3};
    pref.tech = @trigtech;
elseif ( (nargin == 5) && any(strcmpi(varargin{1}, {'trig', 'periodic'}))...
        &&  strcmpi(varargin{2}, 'eps') )
    % TECH, EPS and DOMAIN are specified
    pref.tech = @trigtech;
    pref.chebfuneps = varargin{3};
    dom = varargin{4};
elseif ( (nargin == 5) && any(strcmpi(varargin{1}, {'trig', 'periodic'}))...
        &&  strcmpi(varargin{3}, 'eps') )
    % TECH, DOMAIN and EPS are specified
    pref.tech = @trigtech;
    dom = varargin{2};
    pref.chebfuneps = varargin{4};
elseif ( (nargin == 5) && any(strcmpi(varargin{4}, {'trig', 'periodic'}))...
        &&  strcmpi(varargin{1}, 'eps') )
    % EPS, DOMAIN and TECH are specified
    pref.chebfuneps = varargin{2};
    dom = varargin{3};
    pref.tech = @trigtech;
elseif ( (nargin == 5) && any(strcmpi(varargin{3}, {'trig', 'periodic'}))...
        &&  strcmpi(varargin{1}, 'eps') )
    % EPS, TECH and DOMAIN are specified
    pref.chebfuneps = varargin{2};
    pref.tech = @trigtech;
    dom = varargin{4};
else
    error('CHEBFUN:CHEBFUN3T:constructor', ...
                    'Domain not valid or fully determined.');
end
end
%%
function op = str2op( op )
% OP = STR2OP(OP), finds independent variables in a string and returns an op
% handle than can be evaluated.

vars = symvar(op); % Independent variables
if ( numel(vars) > 3)
    error('CHEBFUN:CHEBFUN3T:constructor:str2op:depvars', ...
        'Too many independent variables in string input.');
else
    op = eval(['@(' vars{1} ',' vars{2} ',' vars{3} ')' op]);
end

end


% Coupled nonlinear PDE's
% Buckley Leverett equation
% dependent variables: pressure and water saturation
% Prepared for educational purposes by ** AAE **
% For whatever reason, the code is very slow and does not converge
% DOES NOT WORK WELL
% Written by Ali A. Eftekhari
% Last checked: June 2021
clc; clear;
%% define the geometry
Nx = 50; % number of cells in x direction
Ny = 30; % number of cells in y direction
W = 50; % [m] length of the domain in x direction
H = 30; % [m] length of the domain in y direction
x1=linspace(0,W, Nx);
x2=x1+0.001;
x=zeros(2*Nx,1);
j=0;
for i=1:Nx
    j=j+1;
    x(j)=x1(i);
    j=j+1;
    x(j)=x2(i);
end
y1=linspace(0,H, Ny);
y2=y1+0.01;
y=zeros(2*Ny,1);
j=0;
for i=1:Ny
    j=j+1;
    y(j)=y1(i);
    j=j+1;
    y(j)=y2(i);
end
m = createMesh2D(x, y); % creates a 2D mesh
[X,Y]=ndgrid(m.cellsize.x, m.cellsize.y);
%% define the physical parametrs
p0 = 1e5; % [bar] pressure
pin = 50e5; % [bar] injection pressure at the left boundary
q_in=5e-5; % [m^3/s] water injection
sw0 = 0; % initial water saturation
swin = 1;
mu_oil = 1e-3; % [Pa.s] oil viscosity
mu_water = 1e-3; % [Pa.s] water viscosity
% reservoir
k0 = 2e-12; % [m^2] average reservoir permeability
k_frac=100e-12; % [m^2] fracture permeability
phi0 = 0.2; % average porosity
clx=0.05;
cly=0.05;
V_dp=0.1; % Dykstra-Parsons coef.
perm_val= field2d(2*Nx,2*Ny,k0,V_dp,clx,cly);
k=createCellVariable(m, k0);
for i=4:2:2*Nx
    k.value(i,3:end-2)=k_frac;
end
for j=4:2:2*Ny
    k.value(3:end-2,j)=k_frac;
end
k.value(1:2,:)=k_frac;
k.value(end-1:end,:)=k_frac;
k.value(:,1:2)=k_frac;
k.value(:,end-1:end)=k_frac;
%k.value(1:2,:)=k_frac;
phi=createCellVariable(m, phi0);
krw0 = 0.4;
kro0 = 0.8;
nw = 1;
no = 1;
krw = @(sw)(krw0*sw.^nw);
dkrwdsw = @(sw)(krw0*nw*sw.^(nw-1));
kro = @(sw)(kro0*(1-sw).^no);
dkrodsw = @(sw)(-kro0*no*(1-sw).^(no-1));
lw = geometricMean(k)/mu_water;
lo = geometricMean(k)/mu_oil;
Lw = @(sw)(krw(sw));
Lo = @(sw)(k/mu_oil*kro(sw));
dLwdsw = @(sw)(k/mu_water*dkrwdsw(sw));
dLodsw = @(sw)(k/mu_oil*dkrodsw(sw));
%% Define the boundaries
BCp = createBC(m); % Neumann BC for pressure
BCs = createBC(m); % Neumann BC for saturation
% change the right boandary to constant pressure (Dirichlet)
BCp.right.a(end-5:end)=0; BCp.right.b(end-5:end)=1; BCp.right.c(end-5:end)=p0;
% change the left boundary to constant flow
BCp.left.a(1:5)=-k_frac/mu_water; BCp.left.b(1:5)=0; BCp.left.c(1:5)=q_in;
% change the left boundary to constant saturation (Dirichlet)
BCs.left.a(1:5)=0; BCs.left.b(1:5)=1; BCs.left.c(1:5)=1;
%% define the time step and solver properties
dt = 10; % [s] time step
t_end = 1000*dt; % [s] final time
eps_p = 1e-5; % pressure accuracy
eps_sw = 1e-5; % saturation accuracy
%% define the variables
sw_old = createCellVariable(m, sw0, BCs);
p_old = createCellVariable(m, p0, BCp);
sw = sw_old;
p = p_old;
uw = -gradientTerm(p_old); % an estimation of the water velocity
%% start the main loop
% generate intial pressure profile (necessary to initialize the fully
% implicit solver)

t = 0;
while (t<t_end)
    error_p = 1e5;
    error_sw = 1e5;
    % estimation loop: sequential
%     for i = 1:2
%         pgrad = gradientTerm(m, p.value);
%         sw_face = upwindMean(m, -pgrad, sw.value); % average value of water saturation
%         labdao = lo.*funceval(kro, sw_face);
%         labdaw = lw.*funceval(krw, sw_face);
%         labda = labdao+labdaw;
%         Mdiffp1 = diffusionTerm(m, -labda);
%         [Mbcp, RHSbcp] = boundaryCondition(m, BCp);
%         p.value = reshapeCell(m, full((Mdiffp1+Mbcp)\RHSbcp));
%         for j = 1:2
%             pgrad = gradientTerm(m, p.value);
%             sw_face = upwindMean(m, -pgrad, sw.value);
%             labdaw = lw.*funceval(krw, sw_face);
%             dlabdawdsw = lw.*funceval(dkrwdsw, sw_face);
%             [Mbcsw, RHSbcsw] = boundaryCondition(m, BCs);
%             [Mtranssw2, RHStrans2] = transientTerm(m, phi, dt, sw);
%             Mconvsw2 = convectionUpwindTerm(m, -dlabdawdsw.*pgrad);
%             RHSs = RHSbcsw+RHStrans2+divergenceTerm(m, (labdaw-dlabdawdsw.*sw_face).*pgrad);
%             Ms = Mbcsw+Mconvsw2+Mtranssw2;
%             sw.value = reshapeCell(m, full(Ms\RHSs));
%         end
%     end
    % Implicit loop
    hold all
    while ((error_p>eps_p) || (error_sw>eps_sw))
        % calculate parameters
        pgrad = gradientTerm(p);
        sw_face = upwindMean(sw, -pgrad); % average value of water saturation
        labdao = lo.*funceval(kro, sw_face);
        labdaw = lw.*funceval(krw, sw_face);
        dlabdaodsw = lo.*funceval(dkrodsw, sw_face);
        dlabdawdsw = lw.*funceval(dkrwdsw, sw_face);
        labda = labdao+labdaw;
        dlabdadsw = dlabdaodsw+dlabdawdsw;
        % compute [Jacobian] matrices
        Mdiffp1 = diffusionTerm(-labda);
        Mdiffp2 = diffusionTerm(-labdaw);
        Mconvsw1 = convectionUpwindTerm(-dlabdadsw.*pgrad);
        Mconvsw2 = convectionUpwindTerm(-dlabdawdsw.*pgrad);
        [Mtranssw2, RHStrans2] = transientTerm(sw_old, dt, phi);
        % Compute RHS values
        RHS1 = divergenceTerm(-dlabdadsw.*sw_face.*pgrad);
        RHS2 = divergenceTerm(-dlabdawdsw.*sw_face.*pgrad);
        % include boundary conditions
        [Mbcp, RHSbcp] = boundaryCondition(BCp);
        [Mbcsw, RHSbcsw] = boundaryCondition(BCs);
        % Couple the equations; BC goes into the block on the main diagonal
        M = [Mdiffp1+Mbcp Mconvsw1; Mdiffp2 Mconvsw2+Mtranssw2+Mbcsw];
        RHS = [RHS1+RHSbcp; RHS2+RHStrans2+RHSbcsw];
        % solve the linear system of equations
        x = M\RHS;
        % x = agmg(M, RHS, [], 1e-10, 500, [], [p.value(:); sw.value(:)]);
        % separate the variables from the solution
        p_new = reshapeCell(m,full(x(1:(2*Nx+1)*(2*Ny+1))));
        sw_new = reshapeCell(m,full(x((2*Nx+1)*(2*Ny+1)+1:end)));
        % calculate error values
        error_p = max(abs((p_new(:)-p.value(:))./p_new(:)));
        error_sw = max(abs(sw_new(:)-sw.value(:)));
        % assign new values of p and sw
        p.value = p_new;
        sw.value = sw_new;
    end
    t=t+dt;
    p_old = p;
    sw_old = sw;
    figure(1);visualizeCells(sw);shading interp; drawnow;
    sw_tot=sw.value.*X.*Y/(W*H);
    oil_prod=(sum(sum(sw_tot(2:end-1,2:end-1))));
    figure(2);semilogx(t, p_new(2,2), 'o'); drawnow;
end


% Like a durational HMM, except we use soft evidence on the observed nodes.
% Should give the same results as HSMM/test_mgram2.

past = 1;
% If past=1, P(Yt|Qt=j,Dt=d) = P(y_{t-d+1:t}|j)
% If past=0, P(Yt|Qt=j,Dt=d) = P(y_{t:t+d-1}|j) - future evidence

words = {'the', 't', 'h', 'e'};
data = 'the';
nwords = length(words);
word_len = zeros(1, nwords);
word_prob = normalise(ones(1,nwords));
word_logprob = log(word_prob);
for wi=1:nwords
  word_len(wi)=length(words{wi});
end
D = max(word_len);

alphasize = 26*2;
data = letter2num(data);
T = length(data);

% node numbers
W = 1; % top level state = word id
L = 2; % bottom level state = letter position within word
F = 3;
O = 4;

ss = 4;
intra = zeros(ss,ss);
intra(W,[F L O])=1;
intra(L,[O F])=1;

inter = zeros(ss,ss);
inter(W,W)=1;
inter(L,L)=1;
inter(F,[W L O])=1;

% node sizes
ns = zeros(1,ss);
ns(W) = nwords;
ns(L) = D;
ns(F) = 2;
ns(O) = alphasize;
ns2 = [ns ns];

% Make the DBN
bnet = mk_dbn(intra, inter, ns, 'observed', O);
eclass = bnet.equiv_class;

% uniform start distrib over words, uniform trans mat
Wstart = normalise(ones(1,nwords));
Wtrans = mk_stochastic(ones(nwords,nwords));
%Wtrans = ones(nwords,nwords);

% always start in state d = length(word) for each bottom level HMM
Lstart = zeros(nwords, D);
for i=1:nwords
  l = length(words{i});
  Lstart(i,l)=1;
end

% make downcounters
RLtrans = mk_rightleft_transmat(D, 0); % 0 self loop prob
Ltrans = repmat(RLtrans, [1 1 nwords]);

% Finish when downcoutner = 1
Fprob = zeros(nwords, D, 2);
Fprob(:,1,2)=1;
Fprob(:,2:end,1)=1;

% Define CPDs for slice 1
bnet.CPD{eclass(W,1)} = tabular_CPD(bnet, W, 'CPT', Wstart);
bnet.CPD{eclass(L,1)} = tabular_CPD(bnet, L, 'CPT', Lstart);
bnet.CPD{eclass(F,1)} = tabular_CPD(bnet, F, 'CPT', Fprob);

% Define CPDs for slice 2
bnet.CPD{eclass(W,2)} = hhmmQ_CPD(bnet, W+ss, 'Fbelow', F, 'startprob', Wstart,  'transprob', Wtrans);
bnet.CPD{eclass(L,2)} = hhmmQ_CPD(bnet, L+ss, 'Fself', F, 'Qps', W+ss, 'startprob', Lstart, 'transprob', Ltrans);

if 0
% To test it is generating correctly, we create an artificial
% observation process that capitalizes at the start of a new segment
% Oprob(Ft-1,Qt,Dt,Yt)
Oprob = zeros(2,nwords,D,alphasize);
Oprob(1,1,3,letter2num('t'),1)=1;
Oprob(1,1,2,letter2num('h'),1)=1;
Oprob(1,1,1,letter2num('e'),1)=1;
Oprob(2,1,3,letter2num('T'),1)=1;
Oprob(2,1,2,letter2num('H'),1)=1;
Oprob(2,1,1,letter2num('E'),1)=1;
Oprob(1,2,1,letter2num('a'),1)=1;
Oprob(2,2,1,letter2num('A'),1)=1;
Oprob(1,3,1,letter2num('b'),1)=1;
Oprob(2,3,1,letter2num('B'),1)=1;
Oprob(1,4,1,letter2num('c'),1)=1;
Oprob(2,4,1,letter2num('C'),1)=1;

% Oprob1(Qt,Dt,Yt)
Oprob1 = zeros(nwords,D,alphasize);
Oprob1(1,3,letter2num('t'),1)=1;
Oprob1(1,2,letter2num('h'),1)=1;
Oprob1(1,1,letter2num('e'),1)=1;
Oprob1(2,1,letter2num('a'),1)=1;
Oprob1(3,1,letter2num('b'),1)=1;
Oprob1(4,1,letter2num('c'),1)=1;

bnet.CPD{eclass(O,2)} = tabular_CPD(bnet, O+ss, 'CPT', Oprob);
bnet.CPD{eclass(O,1)} = tabular_CPD(bnet, O, 'CPT', Oprob1);

evidence = cell(ss,T);
%evidence{W,1}=1;
sample = cell2num(sample_dbn(bnet, 'length', T, 'evidence', evidence));
str = num2letter(sample(4,:))
end

if 1

[log_obslik, obslik, match] = mk_mgram_obslik(lower(data), words, word_len, word_prob);
% obslik(j,t,d)
softCPDpot = cell(ss,T);
ens = ns;
ens(O)=1;
ens2 = [ens ens];
for t=2:T
  dom = [F W+ss L+ss O+ss];
  % tab(Ft-1, Q2, Dt)
  tab = ones(2, nwords, D);
  if past
    tab(1,:,:)=1; % if haven't finished previous word, likelihood is 1
    %tab(2,:,:) = squeeze(obslik(:,t,:)); % otherwise likelihood of this segment
    for d=1:min(t,D)
      tab(2,:,d) = squeeze(obslik(:,t,d));
    end
  else
    for d=1:max(1,min(D,T+1-t))
      tab(2,:,d) = squeeze(obslik(:,t+d-1,d));
    end
  end
  softCPDpot{O,t} = dpot(dom, ens2(dom), tab);
end
t = 1;
dom = [W L O];
% tab(Q2, Dt)
tab = ones(nwords, D);
if past
  %tab = squeeze(obslik(:,t,:));
  tab(:,1) = squeeze(obslik(:,t,1));
else
  for d=1:min(D,T-t)
    tab(:,d) = squeeze(obslik(:,t+d-1,d));
  end
end
softCPDpot{O,t} = dpot(dom, ens(dom), tab);

%bnet.observed = [];
% uniformative observations
%bnet.CPD{eclass(O,2)} = tabular_CPD(bnet, O+ss, 'CPT', mk_stochastic(ones(2,nwords,D,alphasize)));
%bnet.CPD{eclass(O,1)} = tabular_CPD(bnet, O, 'CPT', mk_stochastic(ones(nwords,D,alphasize)));

engine = jtree_dbn_inf_engine(bnet);
evidence = cell(ss,T);
% we add dummy data to O to force its effective size to be 1.
% The actual values have already been incorporated into softCPDpot
evidence(O,:) = num2cell(ones(1,T));
[engine, ll_dbn] = enter_evidence(engine, evidence, 'softCPDpot', softCPDpot);

%evidence(F,:) = num2cell(2*ones(1,T));
%[engine, ll_dbn] = enter_evidence(engine, evidence);

gamma = zeros(nwords, T);
for t=1:T
  m = marginal_nodes(engine, [W F], t);
  gamma(:,t) = m.T(:,2);
end

gamma

xidbn = zeros(nwords, nwords);
for t=1:T-1
  m = marginal_nodes(engine, [W F W+ss], t);
  xidbn = xidbn + squeeze(m.T(:,2,:));
end

% thee
% xidbn(1,4)  = 0.9412  the->e
% (2,3)=0.0588 t->h
% (3,4)=0.0588 h-e
% (4,4)=0.0588 e-e

end