classdef ProbKB < handle
    % Author:
    %     David Sacharny
    %     Amelia Lessen
    %     Tom Henderson
    %     University of Utah
    %     Spring 2022 
    
    properties
        % kb struct 
        m_kb = struct("clauses", [], "prob", 0);
        %map with keys as string name and value as number for kb
        m_string_map
        
        m_bw_map = strings(1,32);
        % The number of atoms in this KB
        m_num_atoms
        
        m_num_clauses
      
    end

methods (Access = public)
        
        
function obj = ProbKB()
%initialize ProbKB
obj.m_string_map = containers.Map;
obj.m_num_atoms = 0;
obj.m_num_clauses = 0;

end
  

function ns = formatString(kb,s)
% formatString- format ProbKB input
% On input:
%     kb: ProbKB
%     s (string vector): input to tellClause
% On output:
%     ns (string): properly formatted string
% Call:
%    kb.formatString(["a"]);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
%   
NEG_OP = '-';

l = length(s);
nums = zeros(1,l);
for i = 1:l
    s2 = s(i);
    y = contains(s2, NEG_OP);
    s2 = strrep(s2, NEG_OP, "");
    
   
    if isKey(kb.m_string_map, s2)
        nums(i) = kb.m_string_map(s2);
    else
        kb.m_num_atoms = kb.m_num_atoms + 1;
        nums(i) = kb.m_num_atoms;
        kb.m_string_map(s2) = nums(i);
        kb.m_bw_map(nums(i)) = s2;
    end
    if(y)
        nums(i) = -nums(i);
    
    end
end
ns = nums;

end


function ns = formatStringDoNotAdd(kb,s)
% formatString- format ProbKB input but do not change ProbKB
% On input:
%     kb: ProbKB
%     s (string vector): input to tellClause
% On output:
%     ns (string): properly formatted string
% Call:
%    kb.formatStringDoNotAdd(["a"]);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
%  
NEG_OP = '-';

l = length(s);
nums = zeros(1,l);
for i = 1:l
    s2 = s(i);
    y = contains(s2, NEG_OP);
    s2 = strrep(s2, NEG_OP, "");
    
    
    if isKey(kb.m_string_map, s2)
        nums(i) = kb.m_string_map(s2);
    else
        ns = false;
        return;
      
    end
    if(y)
        nums(i) = -nums(i);
    
    end
end
ns = nums;

end
%--------------------------------------------------------------------------
function tellClause(kb, clause, prob)
% tellClause- add clause to kb
% On input:
%     kb: ProbKB
%     clause (string vector): input to tellClause
%     prob (float): probability of clause
% On output:
%     c (bool) : does current solution set add to one
% Call:
%    kb.tellClause(["a"], 0.7);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
% 
ns = formatString(kb, clause);
b = checkForClause(kb, ns, prob);

if(b == 0)
    kb.m_num_clauses = kb.m_num_clauses + 1;
    if kb.m_num_clauses == 1
        kb.m_kb(kb.m_num_clauses).clauses = [kb.m_kb(kb.m_num_clauses).clauses, ns];
    else
    kb.m_kb(kb.m_num_clauses).clauses = ns;
    end
    kb.m_kb(kb.m_num_clauses).prob = prob;
end

%p_cc = BR_Nilsson(kb.m_kb);

%if(sum(p_cc) - 1 <= 0.00001)
 %   c = true;
%else
 %   c = false;
%end

end
        

function tellClause2(kb, clause)
% tellClause- add clause to kb
% On input:
%     kb: ProbKB
%     clause (string vector): input to tellClause
%     prob (float): probability of clause
% On output:
%     c (bool) : does current solution set add to one
% Call:
%    kb.tellClause(["a"], 0.7);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
% 
ns = formatString(kb, clause);
updateClause(kb, ns);

%p_cc = BR_Nilsson(kb.m_kb);

%if(sum(p_cc) - 1 <= 0.00001)
 %   c = true;
%else
 %   c = false;
%end

end

function atom_probs = askClause(kb, clause)
% askClause- return probabilities
% On input:
%     kb: ProbKB
%     clause (string vector): a clause in the ProbKB
% On output: 
%     P_cc (kx1 vector): model probabilities
%     p_query (float): probability of query
% Call:
%    kb.askClause(["a"]);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
% 
ns = formatStringDoNotAdd(kb, clause);
if(ns == false)
   disp("Clause does not exist in current KB."); %how to exit?
    %do we want disp message???
    atom_probs = 0;  %return 0???
    return;
end


 BR_KB2F_ind(kb.m_kb,'KBWorkingi');
 num_vars = max(BR_vars(kb.m_kb,[]));
 x0 = 0.5*ones(num_vars,1); % should initialize 0/1 probs
 [atom_probs,~,~,~] = BR_KB_Jacobian(kb.m_kb,'KBWorkingi',...
        1,0.001,50,x0);
%sum and give query prob back
end  

%------------------------------------------------------------------------    

function b = checkForClause(kb, newClause, newProb)
% checkForClause- checks if clause is in kb and updates prob if it is
% On input:
%     kb: ProbKB
%     newClause (string vector): the clause to check if in the ProbKB
%     newProb (float): update to this prob if clause in kb
% On output: 
%     b: 1 if updated, 0 if not in kb
% Call:
%    kb.checkForClause(["a"], 0.5);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
% 

    b = 0;
    i = length(kb.m_kb);
    for j = 1:i
        if isequal(unique(kb.m_kb(j).clauses),unique(newClause))
            if(kb.m_kb(j).prob ~= 0 && kb.m_kb(j).prob ~= 1)
                kb.m_kb(j).prob = newProb;
            end
            b = 1;
            return;
        end
    end
end

function updateClause(kb, newClause)
% checkForClause- checks if clause is in kb and updates prob if it is
% On input:
%     kb: ProbKB
%     newClause (string vector): the clause to check if in the ProbKB
%     newProb (float): update to this prob if clause in kb
% On output: 
%     b: 1 if updated, 0 if not in kb
% Call:
%    kb.checkForClause(["a"], 0.5);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
% 

    b = 0;
    i = length(kb.m_kb);
    for j = 1:i
        if isequal(unique(kb.m_kb(j).clauses),unique(newClause))
            if(kb.m_kb(j).prob ~= 0 && kb.m_kb(j).prob ~= 1)
                kb.m_kb(j).prob = kb.m_kb(j).prob + 0.2;
            end
            return;
        end
    end
end

function b = checkExists(kb, newClause)
% checkForClause- checks if clause is in kb
% On input:
%     kb: ProbKB
%     newClause (string vector): the clause to check if in the ProbKB
% On output: 
%     b: 1 if in kb, 0 if not in kb
% Call:
%    kb.checkExists(["a"]);
% Author:
%     Amelia Lessen
%     UU
%     Spring 2022
% 
    b = 0;
    i = length(kb.m_kb);
    for j = 1:i
        if isequal(unique(kb.m_kb(j).clauses), unique(newClause))
            b = j;
            return;
        end
        
    end
end
end
end

