classdef dbm < handle
    properties (Constant)
        n_nodes_deep_levels = [config.getInstance.dbm_h1_nodes config.getInstance.dbm_h2_nodes]
        learning_rate = config.getInstance.dbm_learning_rate
        
        % equilibriums
        n_mean_field_echoes = config.getInstance.dbm_n_mean_field_echoes
        n_gibbs_echoes = config.getInstance.dbm_n_gibbs_echoes
        n_use_echoes = config.getInstance.dbm_use_echoes
    end

    properties (SetAccess = immutable)
        lang
        n_visible_nodes
        n_levels
    end

    properties (SetAccess = public)
        activities
        biases
        weights
        n_training_iterations
    end
   
    methods
        function obj = dbm(lang)
            obj.n_levels = size(dbm.n_nodes_deep_levels, 2) + 1;
            obj.lang = lang;
            obj.n_visible_nodes = lang.n_input_nodes;
            obj.activities = cell(obj.n_levels, 1);
            obj.biases = cell(obj.n_levels, 1);
            obj.weights = cell(obj.n_levels-1, 1);
            obj.n_training_iterations = 0;

            obj.reset;
            
            disp("DBM parameters:")
            disp("Number of hidden nodes in level h1: " + dbm.n_nodes_deep_levels(1))
            disp("Number of hidden nodes in level h2: " + dbm.n_nodes_deep_levels(2))
            disp("Learning rate: " + dbm.learning_rate)
            disp("Initialise weights with random small numbers: " ...
                + config.getInstance.dbm_initial_non_zero_weights)
            disp("n_mean_field_echoes: " + dbm.n_mean_field_echoes)
            disp("n_gibbs_echoes: " + dbm.n_gibbs_echoes)
            disp("use_echoes: " + dbm.n_use_echoes)
        end

        function load_trained_dbm(obj, n_instances)
            % 1st layer activities and biases
            obj.activities{1} = zeros(1, obj.n_visible_nodes, 'double');
            obj.biases{1} = zeros(1, obj.n_visible_nodes, 'double');
            
            % 1st layer weights
            dbm_filename = fullfile(config.getInstance.path_dbms, obj.lang.name ...
                + "_dbm_layer1_" + n_instances + "." + config.getInstance.dbm_format);
            obj.weights{1} = readmatrix(dbm_filename);

            for i = 2:obj.n_levels
                % i-th layer activities and biases
                obj.activities{i} = zeros(1, dbm.n_nodes_deep_levels(i-1), 'double');
                obj.biases{i} = zeros(1, dbm.n_nodes_deep_levels(i-1), 'double');
                
                % i-th layer weights
                if i < obj.n_levels
                    dbm_filename = fullfile(config.getInstance.path_dbms, obj.lang.name ...
                        + "_dbm_layer" + i + "_" + n_instances + "." ...
                        + config.getInstance.dbm_format);
                    obj.weights{i} = readmatrix(dbm_filename);
                end
            end

            % update training iterations counter
            obj.n_training_iterations = n_instances;
        end

        function reset(obj)
            obj.activities{1} = zeros(1, obj.n_visible_nodes, 'double');
            obj.biases{1} = zeros(1, obj.n_visible_nodes, 'double');

            if config.getInstance.dbm_initial_non_zero_weights
                obj.weights{1} = rand(obj.n_visible_nodes, dbm.n_nodes_deep_levels(1)) / 1000;
            else
                obj.weights{1} = zeros(obj.n_visible_nodes, dbm.n_nodes_deep_levels(1), 'double');
            end

            for i = 2:obj.n_levels
                obj.activities{i} = zeros(1, dbm.n_nodes_deep_levels(i-1), 'double');
                obj.biases{i} = zeros(1, dbm.n_nodes_deep_levels(i-1), 'double');
                if i < obj.n_levels
                    obj.weights{i} = zeros( ...
                        dbm.n_nodes_deep_levels(i-1), dbm.n_nodes_deep_levels(i), 'double');
                end
            end

            obj.n_training_iterations = 0;
        end

        function learn(obj, train_batch_data)
            for batch = 1:size(train_batch_data, 3)
                for step = 1:size(train_batch_data, 1)
                    % Populate input level and spread it up
                    input_data = train_batch_data(step, :, batch);
                    obj.spread_up_input(input_data, 0)

                    % Hebbian learning phase
                    obj.hebbian_learning(dbm.learning_rate)

                    % Dreaming phase
                    obj.resonate(1)

                    % Anti-Hebbian learning phase
                    obj.hebbian_learning(-dbm.learning_rate)
                end
            end

            % clean up activities
            obj.clean_up_activities

            % update training iterations
            obj.n_training_iterations = obj.n_training_iterations ...
                + size(train_batch_data, 3) * size(train_batch_data, 1);

            % save trained dbm to disk
            obj.save
            comment("TRAINED!!!!!!", "else")

            % clean big data from memory
            clear train_batch_data
        end

        function comprehend(obj, input_vector)
            if config.getInstance.dbm_compute_similarities
                comment("comprehending", 'header')
            end
            obj.use(input_vector, 1:obj.lang.n_input_nodes)
        end
        
        function produce(obj, input_vector)
            if config.getInstance.dbm_compute_similarities
                comment("producing", 'header')
            end
            obj.use(input_vector, 1:obj.lang.n_sound_nodes)
        end
        
        function similarity = measure_similarity_between_dbm_states(obj, a1, a2)
            similarity = zeros(1, 3);
            for level = 1:obj.n_levels
                if norm(a1{level}) ~= 0 && norm(a2{level}) ~= 0 
                    similarity(level) =  dbm.cosine_similarity(a1{level}, a2{level});
                end
            end
        end

        function table_similarity = measure_phonological_similarity(obj, input_utterances, unclamped)
            % input: is a vector for every utterance
            % output: create 3 tables with activities {2}, {3}, 
            % and a combination of them (after spreading)

            table_similarity{1} = zeros(size(input_utterances, 1), size(input_utterances, 1));
            table_similarity{2} = zeros(size(input_utterances, 1), size(input_utterances, 1));
            table_similarity{3} = zeros(size(input_utterances, 1), size(input_utterances, 1));
            
            second_levels_per_utterance = zeros(size(input_utterances, 1), ...
                obj.n_nodes_deep_levels(1));
            third_levels_per_utterance = zeros(size(input_utterances, 1), ...
                obj.n_nodes_deep_levels(2));
            both_levels_per_utterance = zeros(size(input_utterances, 1), ...
                obj.n_nodes_deep_levels(1) + obj.n_nodes_deep_levels(2));
            
            % spread and remember the hidden levels
            for i=1:size(input_utterances, 1)
                input_vector = input_utterances(i, :);
                obj.use(input_vector, unclamped{i,1})
                second_levels_per_utterance(i,:) = obj.activities{2};
                third_levels_per_utterance(i,:) = obj.activities{3};
                both_levels_per_utterance(i,:) = [obj.activities{2} obj.activities{3}];
            end

            % compare activities{2} vectors between each other for similarity
            for i=1:size(input_utterances, 1)
                for j=1:size(input_utterances, 1)
                    v1 = second_levels_per_utterance(i,:);
                    v2 = second_levels_per_utterance(j,:);
                    table_similarity{1}(i, j) = dbm.cosine_similarity(v1, v2);

                    v1 = third_levels_per_utterance(i,:);
                    v2 = third_levels_per_utterance(j,:);
                    table_similarity{2}(i, j) = dbm.cosine_similarity(v1, v2);

                    v1 = both_levels_per_utterance(i,:);
                    v2 = both_levels_per_utterance(j,:);
                    table_similarity{3}(i, j) = dbm.cosine_similarity(v1, v2);
                end
            end
        end

        function io_similarity = measure_io_similarity(obj, ...
                input_utterances, full_utterances, unclamped)
            % create similarity table with tested output in rows and all
            % the inputs in columns
            io_similarity = zeros( ...
                size(input_utterances, 1), size(input_utterances, 1));
            for i=1:size(input_utterances, 1)
                input_vector = input_utterances(i, :);
                obj.use(input_vector, unclamped{i,1})
                for j=1:size(input_utterances, 1)
                    io_similarity(i, j) = dbm.cosine_similarity( ...
                        obj.activities{1}, full_utterances(j, :));
                end
            end
        end

        function draw_trained(obj)
            if obj.n_visible_nodes == lang1.n_input_nodes
                draw.dbm(obj.activities, obj.weights, draw.lang1_trained_dbm);
            elseif obj.n_visible_nodes == lang2.n_input_nodes
                draw.dbm(obj.activities, obj.weights, draw.lang2_trained_dbm);
            else
                error("Attempting to draw DBM for unknown language with " ...
                    + obj.n_visible_nodes + " input nodes")
            end
        end

        function draw(obj)
            if obj.n_visible_nodes == lang1.n_input_nodes
                draw.dbm(obj.activities, obj.weights, draw.lang1_dbm);
            elseif obj.n_visible_nodes == lang2.n_input_nodes
                draw.dbm(obj.activities, obj.weights, draw.lang2_dbm);
            else
                error("Attempting to draw DBM for unknown language with " ...
                    + obj.n_visible_nodes + " input nodes")
            end
        end
    end

    
    methods (Access = private)
        function spread_up_input(obj, input_data, stochastic)
            if config.getInstance.dbm_compute_similarities
                comment("learning spread up similarity", 'header')
            end
            
            % populate input level
            obj.activities{1} = input_data;
            
            % zero deep levels
            for level = 2:obj.n_levels
                obj.activities{level} = zeros(1, size(obj.activities{level}, 2), "double");
            end

            % spread up
            for e = 1:dbm.n_mean_field_echoes
                if config.getInstance.dbm_compute_similarities
                    prev_activities = obj.activities;
                end

                % do the spreading
                for level = 2:obj.n_levels
                    obj.spread_to_level(level, stochastic);
                end
            
                % measure similarity to know about the equilibrium state
                if config.getInstance.dbm_compute_similarities
                    similarity = obj.measure_similarity_between_dbm_states( ...
                        prev_activities, obj.activities);
                    comment("learning spread up similarity, step " + e + ": " ...
                        + sprintf('%f %f %f', similarity(:)), "else")
                end
            end
            if config.getInstance.dbm_compute_similarities
                disp(obj.activities)
            end
        end

        function hebbian_learning(obj, learning_rate)
            for level = 1:obj.n_levels
                obj.biases{level} = obj.biases{level} + ...
                    learning_rate * obj.activities{level};
            end

            for level = 1:obj.n_levels-1
                obj.weights{level} = obj.weights{level} + ...
                    learning_rate * obj.activities{level}.' * obj.activities{level+1};
            end
        end
        
        function resonate(obj, stochastic)
            if config.getInstance.dbm_compute_similarities
                comment("resonating", 'header')
            end
            for e = 1:dbm.n_gibbs_echoes
                if config.getInstance.dbm_compute_similarities
                    prev_activities = obj.activities;
                end

                % do the spreading
                obj.spread_to_level1(1:obj.lang.n_input_nodes);
                obj.spread_to_level(obj.n_levels, stochastic);
                for level = 2:obj.n_levels-1
                    obj.spread_to_level(level, stochastic);
                end
                 
                % measure similarity to know about the equilibrium state
                if config.getInstance.dbm_compute_similarities
                    similarity = obj.measure_similarity_between_dbm_states( ...
                        prev_activities, obj.activities);
                    comment("learning resonate similarity, step " + e + ": " ...
                        + sprintf('%f %f %f', similarity(:)), "else")
                end
            end          
        end

        function spread_to_level (obj, level, stochastic)
            if level == 1
                error("call spread_to_level1 to spread to level 1")
            elseif level == obj.n_levels
                obj.activities{level} = ...
                    obj.activities{level-1} * obj.weights{level-1} + obj.biases{level};
            else
                obj.activities{level} = ...
                    obj.activities{level+1} * obj.weights{level}.' + obj.biases{level} ...
                    + obj.activities{level-1} * obj.weights{level-1};
            end

            obj.activities{level} = dbm.sigmoid(obj.activities{level});
            if stochastic
                obj.activities{level} = binornd(1, obj.activities{level});
            end
        end
        
        function spread_to_level1 (obj, unclamped)
            activities_1 = obj.activities{2} * obj.weights{1}.' + obj.biases{1};
            obj.activities{1}(unclamped) = activities_1(unclamped);
        end

        function clean_up_activities(obj)
            obj.activities{1} = zeros(1, obj.n_visible_nodes, 'double');
            for i = 2:obj.n_levels
                obj.activities{i} = zeros(1, dbm.n_nodes_deep_levels(i-1), 'double');
            end
        end

        function use(obj, input_vector, unclamped)
            % populate input level
            obj.activities{1} = input_vector;

            % zero deep levels
            for level = 2:obj.n_levels
                obj.activities{level} = zeros(1, size(obj.activities{level}, 2), "double");
            end

            % spread up
            for e = 1:dbm.n_use_echoes
                if config.getInstance.dbm_compute_similarities
                    prev_activities = obj.activities;
                end
                
                % spread to hidden levels
                for level = 2:obj.n_levels
                    obj.spread_to_level(level, 0);
                end
                % spread to visible level
                obj.spread_to_level1(unclamped);
               
                % measure similarity to know about the equilibrium state
                if config.getInstance.dbm_compute_similarities
                    similarity = obj.measure_similarity_between_dbm_states( ...
                        prev_activities, obj.activities);
                    comment("using spread up similarity, step " + e + ": " ...
                        + sprintf('%f %f %f', similarity(:)), "else")
                end
            end
        end

        function save(obj)
            for i = 1:obj.n_levels-1
                filename = fullfile(config.getInstance.path_dbms, obj.lang.name ...
                    + "_dbm_layer" + i + "_" + obj.n_training_iterations ...
                    + "." + config.getInstance.dbm_format);
                writematrix(obj.weights{i}, filename)
            end
        end
    end

    methods (Static)
        function s = sigmoid(z)
            s = 1.0 ./ (1.0 + exp(-z));
        end

        function similarity = cosine_similarity (v1, v2)
            similarity = 1 - pdist2(v1, v2, 'cosine');
        end
    end
end
