clear, close, clc

%% Exercise 7.1
a = [0.25, -0.2, 0.1, 0.05];
b = [0.6, 0.3, -0.05];
N = 100;
sigma_u = 1;
sigma_d = 0.1;

u_id = randn(N, 1) * sqrt(sigma_u);
[y, y_noiseless] = generate_dataset(a, b, u_id, sigma_d, N);


%% Exercise 7.2

theta_id = cell(8, 1);
for i=1:8
    theta_id{i} = compute_arx(y, u_id, i);
end

%% Exercise 7.3

N_val = 100;
sigma_u = 1;
sigma_d_val = 0.1;

u_val = randn(N_val, 1) * sqrt(sigma_u);
[y_val, y_val_noiseless] = generate_dataset(a, b, u_val, sigma_d_val, N_val);

%% Exercise 7.4

plot(y_val_noiseless, "--black", 'LineWidth', 2);
title('Predicted validation output')
xlabel('Timestep')
ylabel('y')
hold on
plot(y_val, "blacko")
y_val_pred = cell(8, 1);
y_id_pred = cell(8, 1);
for i=1:8
    a = theta_id{i}(1:i);
    b = theta_id{i}(i+1:end);
    
    y_val_pred{i} = inference(a, b, u_val);
    y_id_pred{i} = inference(a, b, u_id);

    plot(y_val_pred{i}, 'color', rand(1,3))
end

%% Exercise 7.5

err_id = zeros(8, 1);
err_val = zeros(8, 1);
for i=1:8
    err_id(i) = mean((y_id_pred{i} - y_noiseless) .^ 2);
    err_val(i) = mean((y_val_pred{i} - y_val_noiseless) .^ 2);
end
figure
plot(err_id, '--redo')
hold on
plot(err_val, '--greeno')

%% Exercise 7.6 and 7.7 I'm just too lazy

function [y, y_noiseless]=generate_dataset(a, b, u, sigma_d, N)
    y = zeros(N, 1);
    y_noiseless = zeros(N, 1);

    d_id = randn(N, 1) * sqrt(sigma_d);

    for i=1:N
        for j=1:length(a)
            if i - j > 0
                y(i) = y(i) + a(j) * y(i-j);
            end
        end
        for j=1:length(b)
            if i - j > 0
                y(i) = y(i) + b(j) * u(i-j);
            end
        end
        y_noiseless(i) = y(i);
        y(i) = y(i) + d_id(i);
    end
end

function theta=compute_arx(y, u, order)
    % Build the phi matrix
    if order == 1
        phi = [y(1:end-1) u(1:end-1)];
    else
        phi_y = toeplitz(y(1:end-1), zeros(order, 1));
        phi_u = toeplitz(u(1:end-1), zeros(order, 1));
        phi = [phi_y, phi_u];
    end
    
    % Identify the system
    theta = phi \ y(2:end);
end

function y_pred=inference(a, b, u)
    N = length(u);
    y_pred = zeros(N, 1);
    for i=1:N
        for j=1:length(a)
            if i - j > 0
                y_pred(i) = y_pred(i) + a(j) * y_pred(i - j);
            end
        end
        for j=1:length(b)
            if i - j > 0
                y_pred(i) = y_pred(i) + b(j) * u(i - j);
            end
        end
    end
end