% SYMREGRESS  A simple example of the symbolic regression. It uses
% genetic algorithm for minimization and reverse Polish notation inside
% chromosomes. It is compatible both with MATLAB and GNU Octave.

% (C) 2015 Alexey Voskov
function symregress
% If launched in GNU Octave -- set gnuplot as graphics toolkit
% (fltk is slow and qt gives bad quality of the plot;
%  tested in GNU Octave 4.0)
a = ver;
names = [a(:).Name];
if ~isempty(strfind(names, 'Octave'))
    disp('GNU Octave detected. Setting gnuplot as graphics toolkit');
    graphics_toolkit('gnuplot');
elseif ~isempty(strfind(names, 'MATLAB'))
    disp('MATLAB detected');
end
% Create initial set of points
example_id = 4;
switch example_id
    case 1
        xv = -3:0.05:3;
        yv = xv.^2 - 2*xv + 3 + randn(size(xv));
        chrom_len = 25;
    case 2
        xv = -3:0.05:3;
        yv = 1./(1 + abs(xv - 0.3).^2.5);
        chrom_len = 25;
    case 3
        xv = -3:0.05:3;
        yv = 4.3*xv.^3 - 10*xv.^2 + 1.5 + 3*randn(size(xv));
        chrom_len = 25;
    case 4
        xv = -2.5:0.1:2.5;
        yv = 2*xv.^4 + 2*xv.^3 - 12*xv.^2 - 8*xv + 16;
        chrom_len = 35;
end
% Generate pool of chromosomes
pool_size = chrom_len * 30;
slen = (pool_size / 15);
pool_genes = {};
pool_rss = [];
for i=1:pool_size;
    chrom = newChromosome(chrom_len);
    rss = calcRSS(xv, yv, chrom, chrom_len);
    pool_genes = [pool_genes, {chrom}];
    pool_rss = [pool_rss, rss];
end
% Run the iterations of the genetic algorithm
rss_old = Inf;
rss_history = [];
iter = 1;
while true && iter < 2000
    % Selection
    [pool_rss, ind] =  sort(pool_rss);
    pool_rss = pool_rss(1:slen);
    pool_genes = pool_genes(ind(1:slen));
    rss_history = [rss_history, pool_rss(1)];    
    % Add random chromosome
    pool_rss(slen) = NaN;
    while isinf(pool_rss(slen)) || isnan(pool_rss(slen))
        pool_genes{slen} = newChromosome(chrom_len);
        pool_rss(slen) = calcRSS(xv, yv, pool_genes{slen}, chrom_len);
    end
    % Show the information about iteration
    fprintf('ITERATION %d| RSS: %.3f (WORST RSS: %.3f)\n', iter, pool_rss(1), pool_rss(end-1));
    iter = iter + 1;
    % If new approximation has been reached -- show it
    if pool_rss(1) < rss_old
        close all;
        rss_old = pool_rss(1);
        yv_calc = calcChromosome(pool_genes{1}.genes, xv);
        % Update the plot with the points and regression model
        subplot(121);
        plot(xv, yv_calc, 'k-', 'LineWidth', 2);
        hold on;
        plot(xv, yv, 'bo', 'LineWidth', 2);
        hold off;
        xlabel('x', 'FontSize', 14); ylabel('y', 'FontSize', 14);
        % Update the plot with the RSS change history
        subplot(122);
        semilogy(rss_history, 'LineWidth', 2);
        xlabel('Generation', 'FontSize', 14);
        ylabel('RSS', 'FontSize', 14);
        % Type new formula obtained by genetic algorithm
        fprintf('----- New formula\n');
        fprintf('Reverse Polish (postfix) notation\n');
        disp(chromosomeToPostfixString(pool_genes{1}.genes));
        fprintf('Classic (infix) notation\n');        
        disp(chromosomeToInfixString(pool_genes{1}.genes));                
        drawnow;
    end
    % Generation of the rest of pool
    for i=1:round(pool_size - slen)        
        ind1 = randi(slen, 1, 1);
        ind2 = randi(slen, 1, 1);
        chrom = mixChromosomes(pool_genes{ind1}, pool_genes{ind2});
        rss = calcRSS(xv, yv, chrom, chrom_len);
        pool_genes = [pool_genes, {chrom}];
        pool_rss = [pool_rss, rss];
    end
end    
end

function rss = calcRSS(xv, yv, chrom, maxlen)
% Calculate sum of squares for the chromosome. If the chromosome
% is longer than maxlen, RSS will be multiplied by 
% numel(chrom.genes) / maxlen coefficients
    rss = sum((yv - calcChromosome(chrom.genes, xv)).^2);
    if numel(chrom.genes) > maxlen
        rss = rss * numel(chrom.genes) / maxlen;
    end
end

function out = newChromosome(len)
    genes_pool = struct(...
        'gene',       {'#', 'X', 'U-', '+', '-', '*', '/', '^'}, ...
        'stackDepth', {0,   0,    1,    2,   2,   2,  2,  2},...
        'deltaStack', {1,   1,    0,   -1,  -1,   -1, -1, -1});
    function gene = newGene(stackDepth)
        % Select an appropriate gene
        done = false;    
        while ~done
            ind = randi([1 numel(genes_pool)], 1, 1);
            if genes_pool(ind).stackDepth <= stackDepth
                done = true;
            end
        end
        % Generate output
        gene = genes_pool(ind);
        if gene.gene == '#'
            gene.gene = randn(1,1) * 5;
        end
    end

    genes = {}; stack = [];
    stackDepth = 0;    
    i = 1;
    while i <= len || (i > len && stackDepth ~= 1)
        gene = newGene(stackDepth);
        stackDepth = stackDepth + gene.deltaStack;
        genes = [genes, gene.gene];
        stack = [stack, stackDepth];
        i = i + 1;
        if strcmp(gene.gene, 'X') && rand(1,1) < 0.5
            genes = [genes, randi([2 4], 1)]; stackDepth = stackDepth + 1;
            stack = [stack, stackDepth];
            genes = [genes, '^']; stackDepth = stackDepth - 1;
            stack = [stack, stackDepth];
        end
    end
    out = struct('genes', {genes}, 'stack', {stack});
end

function m = mutateChromosome(c)
    m = c;
    ops = {'+', '-', '*', '/', '^'};
    inds = find(rand(1,numel(m.genes)) > 0.85);
    for i=inds
        gene = m.genes{i};
        if ~ischar(gene)
            p = rand(1,1);
            if p <= 0.33
                m.genes{i} = m.genes{i} + 0.01*randn(1,1);
            elseif p <= 0.66
                m.genes{i} = m.genes{i} + 0.1*randn(1,1);                
            else
                m.genes{i} = m.genes{i} + randn(1,1);                
            end
        elseif any(strcmp(ops, gene)) 
            m.genes{i} = ops{randi(5,1,1)};
        end
    end
end

function cout = mixChromosomes(c1, c2)
    c1_maxstack = max(c1.stack);
    c2_maxstack = max(c2.stack);
    ci_maxstack = min([c1_maxstack, c2_maxstack]);
    ci_stlev = randi(ci_maxstack,1,1);
        
    ind1v = find(c1.stack == ci_stlev); ind1 = ind1v(randi(numel(ind1v),1,1));
    ind2v = find(c2.stack == ci_stlev); ind2 = ind2v(randi(numel(ind2v),1,1));
        
    cout = struct('genes', {[c1.genes(1:ind1), c2.genes((ind2+1):end)]}, ...
        'stack', {[c1.stack(1:ind1), c2.stack((ind2+1):end)]});
    
    cout = mutateChromosome(cout);
    
end

function str = chromosomeToPostfixString(genes)
    str = '';
    for i = 1:numel(genes)
        if isnumeric(genes{i})
            str = sprintf('%s%8.4f ', str, genes{i});
        else
            str = sprintf('%s%8s ', str, genes{i});
        end
        if mod(i, 8) == 0
            str = sprintf('%s\n', str);
        end
    end
    str = sprintf('%s\n', str);
end

function str = chromosomeToInfixString(genes)
    function strGene = addBrackets(gene)
        if isa(gene, 'double')
            if gene >= 0
                strGene = sprintf('%.12f', gene);
            else
                strGene = sprintf('(%.12f)', gene);
            end
        elseif strcmp(gene, 'x')
            strGene = 'x';            
        else
            strGene = ['(', gene, ')'];
        end
    end

    function strGene = geneAbs(gene)
        if isa(gene, 'double')
            strGene = sprintf('%.12f', abs(gene));
        elseif strcmp(gene, '-x')
            strGene = 'x';
        else
            strGene = ['abs(', gene, ')'];
        end        
    end

    function strGene = geneSign(gene)
        if isa(gene, 'double')
            if gene > 0; strGene = '1';
            elseif gene < 0; strGene = '-1';
            else strGene = '0';
            end
        else
            strGene = ['sign(', gene, ')'];
        end                
    end

    function stack = binOp(stack, opfunc, opname)
        st_arg1 = stack{end-1};
        st_arg2 = stack{end};                
        if isa(st_arg1, 'double') && isa(st_arg2, 'double')
            stack = [stack(1:end-2), opfunc(st_arg1, st_arg2)];
        else
            if opname == '^' && isa(st_arg1, 'double');
                st_arg1 = abs(st_arg1);
            end
            if opname == '^'
                stack = [stack(1:end-2), [geneAbs(st_arg1), opname, addBrackets(st_arg2)]];
            elseif opname == '/'
                b = ['((', geneAbs(st_arg2), '+1e-10)*', geneSign(st_arg2), '+1e-10)'];
                stack = [stack(1:end-2), [addBrackets(st_arg1), opname, b]];
            else
                stack = [stack(1:end-2), [addBrackets(st_arg1), opname, addBrackets(st_arg2)]];                
            end
        end        
    end

    stack = {};
    for i=1:numel(genes)
        gene = genes{i};
        switch gene
            case 'X'
                stack = [stack, 'x'];
            case 'U-'
                st_end = stack{end};
                if isa(st_end, 'double')
                    stack = [stack(1:end-1), -st_end];
                elseif strcmp(st_end, 'x')
                    stack = [stack(1:end-1), '-x'];                    
                elseif numel(st_end) > 1 && st_end(1) == '-'
                    stack = [stack(1:end-1), st_end(2:end)];
                else
                    stack = [stack(1:end-1), ['-', addBrackets(st_end)]];
                end
            case 'abs'
                stack = [stack(1:end-1), geneAbs(stack{end})];
            case '+'
                stack = binOp(stack, @(x,y) (x + y), '+');
            case '-'
                stack = binOp(stack, @(x,y) (x - y), '-');
            case '*'
                stack = binOp(stack, @(x,y) (x .* y), '*');
            case '/'
                stack = binOp(stack, @(x,y) (x ./ y), '/');
            case '^'
                stack = binOp(stack, @(x,y) (abs(x) .^ y), '^');                
            otherwise
                stack = [stack, gene];
        end
    end
    for i=1:numel(stack)
        if isa(stack{i}, 'double')
            stack{i} = sprintf('%.12f',stack{i});
        end
    end
    if numel(stack) ~= 1
        error('Invalid chromosome');
    end
    str = [stack{:}];
end

function ans = calcChromosome(genes, x)
    stack_top = 0;
    stack = cell(1,2*numel(genes));
    for i=1:numel(genes)
        gene = genes{i};
        switch gene
            case 'X'
                stack_top = stack_top + 1;                
                stack{stack_top} = x; 
            case 'U-'
                stack{stack_top} = -stack{stack_top};
            case 'abs'
                stack{stack_top} = abs(stack{stack_top});
            case '+'
                stack_top = stack_top - 1;
                stack{stack_top} = stack{stack_top} + stack{stack_top + 1};
            case '-'
                stack_top = stack_top - 1;
                stack{stack_top} = stack{stack_top} - stack{stack_top + 1};
            case '*'
                stack_top = stack_top - 1;
                stack{stack_top} = stack{stack_top} .* stack{stack_top + 1};
            case '/'
                stack_top = stack_top - 1;
                b = stack{stack_top + 1};
                b = (abs(b) + 1e-10).*sign(b) + 1e-10;                
                stack{stack_top} = stack{stack_top} ./ b;
            case '^'
                stack_top = stack_top - 1;
                stack{stack_top} = abs(stack{stack_top}) .^ stack{stack_top + 1};
            otherwise
                stack_top = stack_top + 1;
                stack{stack_top} = gene;
        end
    end
    if stack_top ~= 1
        error('Invalid chromosome');
    end
    ans = stack{1};
end