classdef (Sealed) hmmEntityModel
%

% Copyright 2022 The MathWorks, Inc.
    
    properties(Access=private)
        internalMdl;
        Tags;
    end

    properties(SetAccess=private)
        Entities;
    end
    
    methods

        function tbl = predict(mdl,documents)
            arguments
                mdl (1,1) hmmEntityModel
                documents (1,:) tokenizedDocument
            end
            if all(doclength(documents)==0)
                tokens = strings(0,1);
                tags = categorical(strings(0,1),string(mdl.Tags))';
            else
                tags = mdl.predictTags(documents);
                tokens = doc2cell(documents);
                tokens = [tokens{:}];
            end
            tbl = table(tokens', tags, 'VariableNames',["Token" "Entity"]);
        end
    end

    methods (Hidden)
        function mdl = hmmEntityModel()
        end

        function tags = predictTags(mdl, documents)
            w = warning ('off','textanalytics:corpusLanguage:UnsupportedScript');
            cleanupObj = onCleanup(@() warning(w));
            documents = addSentenceDetails(documents);
            documents = addPartOfSpeechDetails(documents, 'Retokenize','none');
            tags = mdl.internalMdl.predict(tokenDetails(documents));
        end
    end

    methods (Hidden,Static)
        function mdl = train(trainDocs, trainLabels, nonEntityTag)
            mdl = hmmEntityModel;

            w = warning ('off','textanalytics:corpusLanguage:UnsupportedScript');
            cleanupObj = onCleanup(@() warning(w));
            % Splits input into sentences so model learns about sentence boundaries
            trainDocs = splitSentences(trainDocs);
            
            % Split labels to be the same size as trainDocs
            trainLabels = mat2cell(trainLabels, 1, doclength(trainDocs));
            
            mdl.internalMdl = textanalytics.internal.entityModel.train(trainDocs, trainLabels, NonEntityTag=nonEntityTag);
            mdl.Tags = categorical(mdl.internalMdl.Tags);
            
            entities = string(mdl.Tags(mdl.Tags~=nonEntityTag));
            % removes B- and I- from tags
            entities = unique(extractAfter(entities, 2));
            mdl.Entities = categorical([entities;nonEntityTag]);
        end
    end

    methods(Access=private)
        % some things we explicitly do not want to support as in wordEmbedding
        cat(d)
        transpose(d)
        ctranspose(d)
        repmat(d,varargin)
        horzcat(varargin)
        vertcat(varargin)
    end

    methods(Access=?thmmEntityModel)
        function tags = getTags(mdl)
            tags = mdl.Tags;
        end
    end
end


