classdef VDJ_base_model
    % Base model class that defines some generic methods used for the EM algorithm.
    % All VDJ models inherit from this class.
    
    properties (SetAccess = private)
    end
    
    methods
        function model = VDJ_base_model()
            % These are default values for these settings. They can be
            % passed in at run time by calling the function settings()
            
        end
        
        % Function that sets values of properties if passed in as a struct.
        % Used for initialization.
        function model = settings(model,model_settings)
            snames = fieldnames(model_settings);
            for i=1:length(snames)
                model.(snames{i}) = model_settings.(snames{i});
            end
            
        end
        
        % Make checks on completeness/sanity of model parameter values.
        function validate_model(model)
            ps = fieldnames(model);
            for p=1:length(ps)
                if isempty(model.(ps{p}))
                    disp(['Model has empty value for ' ps{p} '.']);
                    disp('Quitting.')
                    assert(1==2);
                end
            end
        end
        
        % Normalizes all probability distributions in model.
        function norm_model = normalize_model(model)
            norm_model = model;
            ps = fieldnames(norm_model);
            for p=1:length(ps)
                if ps{p}(1)=='P'
                    if isempty(strfind(ps{p},'_given_')) 
                        norm_model.(ps{p}) = norm_model.(ps{p})/sum( reshape(norm_model.(ps{p}),[],1) );
                    else % If it's a conditional probability, then normalize appropriately
                        nd = ndims(norm_model.(ps{p}));
                        norm_model.(ps{p}) = conditional_distribution(norm_model.(ps{p}), nd);
                    end
                end
            end
        end
        
        % Initializes the counter variable.
        function counter = initialize_counter(model)
            
            counter.loglikelihood = 0; % This adds up the log likelihood of generating a sequence given the model.
            counter.N_processed = 0;
            % For all parameters of model that begin with a 'P' or an 'M' AND don't have the word 'specific' in them,
            % make a corresponding variable in counter, with an 'n' prefixed to it.
            ps = fieldnames(model);
            for p=1:length(ps)
                if (ps{p}(1)=='P' || ps{p}(1)=='M')
                    counter.(['n' ps{p}]) = zeros(size(model.(ps{p})));
                    
                elseif (ps{p}(1)=='R')
                    % Rate parameter. So make two 'nM' variables in counter, one for numerator and one for denominator.
                    % Names of these two should be separated by '_per_'.
                    % The ratio of these will be estimate of the model parameter.
                    k = strfind(ps{p},'_per_');
                    assert(~isempty(k));
                    numerator_name = ps{p}(2:k-1);
                    denominator_name = ps{p}(k+5:end);
                    counter.(['nM' numerator_name]) = zeros(size(model.(ps{p})));
                    counter.(['nM' denominator_name]) = zeros(size(model.(ps{p})));
                end
            end
            % Add any other fields you might want to keep track of, but are
            % not in model, below. ALWAYS Begin the name with an 'nP'
            % or 'nM'. And don't use the word '_specific'
            
            % These additional fields are added in the derived model class.
            
        end
        
        function model = get_model_from_counter(oldmodel,counter)
            
            model = oldmodel;
            N = counter.N_processed;
            % = round( sum( reshape( counter.([nps{p}]) ,[],1) ) ); % N is simply the sum of all its elements.
            
            % normalize counter variables to get model
            ps = fieldnames(model);
            for p=1:length(ps)
                if (ps{p}(1)=='P') && isfield(counter,['n' ps{p}])
                    % Divide by N to either normalize distribution or get mean value from sum.
                    cvar = counter.(['n' ps{p}]);
                    model.(ps{p}) = cvar./sum(cvar(:));
                    % if model property contains the word '_given_', then get conditional distribution from the joint.
                    if ~isempty(strfind(ps{p},'_given_'))
                        Pc = conditional_distribution(model.(ps{p}), ndims(model.(ps{p})) );
                        Pc(isnan(Pc)) = oldmodel.(ps{p})(isnan(Pc)); % if divided by zero, set those elements to old values.
                        model.(ps{p}) = Pc;
                    end
                    
                elseif ps{p}(1)=='R'
                    % Divide the two appropriate counter variables to get rate estimate.
                    k = strfind(ps{p},'_per_');
                    assert(~isempty(k));
                    numerator_name = ps{p}(2:k-1);
                    denominator_name = ps{p}(k+5:end);
                    model.(ps{p}) = counter.(['nM' numerator_name])./counter.(['nM' denominator_name]);
                elseif (ps{p}(1)=='M') && isfield(counter,['n' ps{p}])
                    % Divide by N to get mean value from sum.
                    model.(ps{p}) = counter.(['n' ps{p}])/N;
                end
            end
        end
        
        % Initialize assigns struct
        function assigns = initialize_assigns(model, counter)
            % For every variable in counter, make a corresponding one in assigns.
            nps = fieldnames(counter);
            for p=1:length(nps)
                if strcmp(nps{p}(1:2),'nP')
                    % If counter variable begins with 'nP' i.e. the model variable is a probability distribution P(x),
                    % then the assigns variable should store the index of the value of x, for each assignment for a sequence.
                    % eg. counter.nPpVdelV -> assigns.pVdelV which stores the indices of palindrome half-length and deletions for V as [1 + npV, 1 + ndV]
                    k_dims = ndim(counter.(nps{p}));
                    assigns.(nps{p}(3:end)) = zeros(model.max_assignments,k_dims);
                elseif  strcmp(nps{p}(1:2),'nM')
                    % If counter variable begins with 'nM' i.e. the model variable is a mean value of x,
                    % then the assigns variable should store the value of x, for each assignment for a sequence.
                    % eg. counter.nMerror_vs_position -> assigns.error_vs_position which stores number of errors at each position for each assignment for a sequence.
                    k_size = size(counter.(nps{p}));
                    assigns.(nps{p}(3:end)) = zeros([model.max_assignments,k_size]);
                end
            end
            assigns.proba = zeros(model.max_assignments,1); % probability of each assignment
            assigns.event_probability = zeros(model.max_assignments,1); % probability of each assignment, assuming zero error rate.
            
            assigns.n_assignments=0; % number of valid assignments
            assigns.skips = 0;
            assigns.likelihood = 0; % total probability of generating the sequence (sum over assignments).
            assigns.generation_probability = 0; % total probability of generating the sequence (assuming zero error rate).
            assigns.max_proba_index = 0;
        end
        
        
        
        % Function that updates the counter according to the assignments and their probabilities for each sequence.
        function [counter, delta_counter] = update_counter(model, assigns, counter)
            if assigns.n_assignments > 0 && assigns.likelihood > 0
                
                n_as = assigns.n_assignments;
                % Add up log likelihood of generating each sequence, given the generative model.
                
                delta_counter = model.initialize_counter();
                delta_counter.loglikelihood = log(assigns.likelihood);
                delta_counter.N_processed = 1;
                
%                 counter.loglikelihood = counter.loglikelihood  + delta_counter.loglikelihood;
%                 counter.N_processed = counter.N_processed + 1;
                
                % Normalize probabilities of assignments.
                assigns.proba=assigns.proba/assigns.likelihood;
                
                ps = fieldnames(assigns);
                for p=1:length(ps)
                    nps = ['nP' ps{p}];
                    nms = ['nM' ps{p}];
                    if isfield(counter, nps)
                        % For probability distributions P(x), figure out count contributions of assignments to each value of x
                        % returns total probability values for each distinct value of x
                        % Only consider assignments where the values of the variable are nonzero. This allows XOR conditions.
                        k_nonzeros = all(assigns.(ps{p})(1:n_as,:)>0,2);
                        if ~isempty(find(k_nonzeros,1)) > 0
                            tmp = accumarray(assigns.(ps{p})(k_nonzeros,:),assigns.proba(k_nonzeros)./sum(assigns.proba(k_nonzeros)));
                            subsize = [size(tmp), ones(1, ndims(counter.(nps)) - ndims(tmp)) ];
                            numDims = numel(subsize);
                            idx = cell(1,numDims);
                            for k=1:numDims
                                idx{k} = 1:subsize(k); % figures out indices of counter variable elements that need to be updated.
                            end
                            delta_counter.(nps)(idx{:}) = tmp;
%                             counter.(nps)(idx{:}) = counter.(nps)(idx{:}) + tmp; % add the count contributions
                        
                        end
                    elseif isfield(counter,nms)
                        % For mean value variables, just add expectation value from this set of assignments for this sequence to counter variable.
                        delta_counter.(nms) = shiftdim(dot2( assigns.proba, assigns.(ps{p}),1,1),1);
%                         counter.(nms) = counter.(nms) + shiftdim(dot2( assigns.proba, assigns.(ps{p}),1,1),1); %This used to be multiprod
                    end
                end
            end            
            counter = model.sum_counters([counter, delta_counter]);            
        end
        
    end % methods
    
    methods (Static = true)
        
        
        
        
        % Function that adds all variables of a set of counters.
        function counter = sum_counters(counters)
            fs = fieldnames(counters(1));
            for k=1:length(fs)
                counter.(fs{k}) = counters(1).(fs{k});
            end
            
            for i=2:length(counters)
                for k=1:length(fs)
                    counter.(fs{k}) = counter.(fs{k}) + counters(i).(fs{k});
                end
            end
        end

        % Add only variables that have the same dimensions
        function counter = sum_compatible_counters(counters)
            fs = fieldnames(counters(1));
            for k=1:length(fs)
                counter.(fs{k}) = counters(1).(fs{k});
            end
            
            for i=2:length(counters)
                for k=1:length(fs)
                    if all(size(counter.(fs{k}))==size(counters(i).(fs{k})))
                        counter.(fs{k}) = counter.(fs{k}) + counters(i).(fs{k});
                    end
                end
            end
        end

        
        % Function that reduces sizes of arrays in assigns to number of assignments. I don't really use this. It wastes time.
        function assigns = remove_unused_assignments(assigns)
            ps = fieldnames(assigns);
            for p=1:length(ps)
                if ~strcmp(ps{p},'n_assignments')
                    assigns.(ps{p})(assigns.n_assignments+1:end,:) = [];
                end
            end
            %assigns.proba(assn_assigns+1:end) = [];
        end
    end
    
end % classdef
