% code by Brian Hunt with input from Zhixin Lu & Jaideep Pathak
% prediction based on Jaeger reservoir model
% modified by Artur Perevalov


% load res_sample
start = 120;
training_length = 25000;
after_training = 5000;

% filtering if necessary
windowsize = 1;  %parameters for box averaging
b = (1/windowsize)*ones(1,windowsize);
a = 1;

%% defining the data
% just creating some powers of sin
t = [1:50000]/1024;
w1 = 4*2*pi;
w2 = 7.2*2*pi;
data = [sin(w1*t).^5; 2*sin(w2*t).^7; sin((w1+w2)*t).^8]';

% data = filter(b, a, data);

data_norm_coef = 100; %normalising to 1 over this coeff
stdd = std(data(:));
data = data/data_norm_coef; %normalising
mean_d = mean(data);
data = data - ones(length(data),1) * mean_d; %substracting mean values


%% Cutting data for training and prediction
Utr = data(start:start+training_length,:);      % training data
Ute = data(start+training_length+1:start+training_length+after_training,:);  % test data

trans = 100;   % transient time
rng(0);         % seed random number generator
reg = 1e-5; % regularization parameter
rho = 1.1;  % spectral radius
p = 0.095;  % connection probability
N = 2000;   % reservoir size
scale_rho = 0.3; % input strength
alpha =  1; % slowing down parameter

train = size(Utr,1);        % training time
dim = size(Utr,2);          % number of coordinates
scale = scale_rho*ones(1,dim);    % input strength for each coordinate
coord = 1;                 % coordinate to graph

U = Utr;

A = sparse(double(rand(N) < p));        % apply connection probability
A(A~=0) = 2*rand(1,length(find(A)))-1;  % randomize nonzero entries
A = rho*A/abs(eigs(A,1));               % scale spectral radius

Win = (2*rand(N,dim)-1)*diag(scale);    % input weights

X = zeros(train+1, N);                          % preallocate
X(1,:) = 2*rand(1,N)-1;

% initialize reservoir
for n = 1:train    % listening phase
    X(n+1,:) = (1-alpha) * X(n,:) + alpha*tanh(X(n,:)*A' + U(n,:)*Win');   % reservoir with input
end

init = X(train+1,:);        % initial state for prediction
X = X((trans+1):train,:);   % throw away transient
U = U((trans+1):train,:);

%% Here is the magic
Wout = (X'*X + reg*eye(N))\(X'*U);  % fit X*Wout to U with regularization

%% plots
plots = 1; %1 if you need to visualise plots

if plots == 1
    figure(1)
    tvals = (trans+1):train;
    subplot(2,1,1)              % graph training signal versus reservoir output
    plot(tvals, U(:,coord), 'b', tvals, X*Wout(:,coord), 'r'); axis tight
    subplot(2,1,2)              % graph normalized RMS error
    %                                 plot(tvals, sqrt(sum((X*Wout-U).^2,2)./(1)))
    plot(tvals, sqrt(sum((X*Wout-U).^2,2)./(sum((X*Wout).^2,2)+sum(U.^2,2))))

    axis tight
end

test = size(Ute,1);             % testing time
%                             U = Ute - ones(test,1)*shift;   % subtract training mean
U = Ute;

XX = zeros(test, N);
XX(1,:) = init;

for n = 1:(test-1)      % predict using autonomous reservoir with feedback
    XX(n+1,:) = (1-alpha) * XX(n,:) + alpha*tanh(XX(n,:)*A' + (XX(n,:)*Wout)*Win');
end

if plots == 1
    figure(2)
    tvals = 1:test;
    subplot(2,1,1)              % graph test signal versus reservoir prediction
    plot(tvals, U(:,coord), 'b', tvals, XX*Wout(:,coord), 'r'); axis tight
    subplot(2,1,2)              % graph normalized RMS error
%                                 plot(tvals, sqrt(sum((XX*Wout-U).^2,2)./(1)))
    plot(tvals, sqrt(sum((XX*Wout-U).^2,2)./(sum((XX*Wout).^2,2)+sum(U.^2,2))))

    axis tight
end

%% evaluating errors
error_type = 2;

if error_type == 1
    er_vec = zeros(1, dim);  % error vector
    for i = 1:dim
        er_vec(1, i) = sqrt(sum( (U(:,i)-XX*Wout(:,i)).^2 ))/sqrt(after_training);
    end

    aae = sum(er_vec)*data_norm_coef/dim;


elseif error_type == 2
    er_matr = zeros(after_training,dim);
    s = std(U);
    for i = 1:dim
        er_matr(:,i) = (abs(U(:,i)-XX*Wout(:,i)))/s(i);
    end

    er_vec = zeros(after_training,1);

    for i = 1:after_training
        er_vec(i) = sqrt(sum(er_matr(i,:).^2)/dim);
    end

    aae = mean(er_vec);

end

if plots == 1
    if error_type == 2
        figure(3)
        plot(er_vec)
    end
end

prediction = U;
experiment = XX*Wout;

clearvars A after_training coord dim init
clearvars n N out p reg rho scale shift start test train training_length
clearvars trans  Ute Utr Win  i scale_rho Ustd data_norm_coef k fid ans timestep
clearvars tvals a b windowsize plots error_type data
clearvars Wout X U XX mean_d s stdd ww
clearvars alpha t w1 w2 aae c
clearvars er_matr er_vec