function [  ] = hwk1(  )
%UNTITLED Summary of this function goes here
%   Detailed explanation goes here

% Solve by the Lax-Friedrichs scheme
figure('Position', [100,100,1200,800]);
clf;
subplotInd = 1;
for lambda = [0.8, 1.6]
    for h = [1/10, 1/20, 1/40]
        M = round(4/h);                         % round to the nearest integer
        k = lambda*h;                           %
        N = round(2.4/k);                       % round to the nearest integer
        u0 = setInitialCondition(M);            % Set initial condition
        u = LaxFriedrichs (M, N, u0, lambda);   % Apply the finite difference scheme

        % Compute and print the error at t_N
        uN = exactSolutionATtN(M);
        err = norm(uN - u(N+1,:));
        fprintf('Lax-Friedrichs scheme with lambda = %f, h=1/%d, the error ||u^N-U^N||_2 = %f\n', lambda, round(1/h), err);

        % draw the graph at t_N
        x = linspace(-1,3,M+1);
        subplot(2,3,subplotInd);
        subplotInd = subplotInd+1;
        plot(x, uN, 'b', x, u(N+1,:), 'r');
        legend('Exact solution', 'Numerical solution');
        title(['L-F scheme, \lambda=' num2str(lambda) ', h=1/' int2str(round(1/h))]);
    end
    fprintf('\n');
end

% Solve by the leap frog scheme
figure('Position', [100,100,1200,800]);
clf;
subplotInd = 1;
for lambda = [0.8, 1.6]
    for h = [1/10, 1/20, 1/40]
        M = round(4/h);                         % round to the nearest integer
        k = lambda*h;                           %
        N = round(2.4/k);                       % round to the nearest integer
        u0 = setInitialCondition(M);            % Set initial condition
        u = LeapFrog (M, N, u0, lambda);        % Apply the finite difference scheme

        % Compute and print the error at t_N
        uN = exactSolutionATtN(M);
        err = norm(uN - u(N+1,:));
        fprintf('Leap frog scheme with lambda = %f, h=1/%d, the error ||u^N-U^N||_2 = %f\n', lambda, round(1/h), err);

        % draw the graph at t_N
        x = linspace(-1,3,M+1);
        subplot(2,3,subplotInd);
        subplotInd = subplotInd+1;
        plot(x, uN, 'b', x, u(N+1,:), 'r');
        legend('Exact solution', 'Numerical solution');
        title(['Leapfrog scheme, \lambda=' num2str(lambda) ', h=1/' int2str(round(1/h))]);
    end
    fprintf('\n');
end


% Solve by the backward-time forward-space scheme
figure('Position', [100,100,1200,800]);
clf;
subplotInd = 1;
for lambda = [0.8, 1.6]
    for h = [1/10, 1/20, 1/40]
        M = round(4/h);                         % round to the nearest integer
        k = lambda*h;                           %
        N = round(2.4/k);                       % round to the nearest integer
        u0 = setInitialCondition(M);            % Set initial condition
        u = BTFS (M, N, u0, lambda);        % Apply the finite difference scheme

        % Compute and print the error at t_N
        uN = exactSolutionATtN(M);
        err = norm(uN - u(N+1,:));
        fprintf('BT-FS scheme with lambda = %f, h=1/%d, the error ||u^N-U^N||_2 = %f\n', lambda, round(1/h), err);

        % draw the graph at t_N
        x = linspace(-1,3,M+1);
        subplot(2,3,subplotInd);
        subplotInd = subplotInd+1;
        plot(x, uN, 'b', x, u(N+1,:), 'r');
        legend('Exact solution', 'Numerical solution');
        title(['BT-FS scheme, \lambda=' num2str(lambda) ', h=1/' int2str(round(1/h))]);
    end
    fprintf('\n');
end
end

function [u0] = setInitialCondition(M)
% Set an (M+1)-dim row vector containing the initial condition
x = linspace(-1,3,M+1);        % Define (M+1)-dim row vector containing x coordinates
u0 = zeros(1,M+1);             % Allocate memory for row vector u0
ind = (x>=-0.5) & (x <= 0.5);  % Indices of grid points in [-0.5, 0.5]
u0(ind) = cos(pi*x(ind)).^2;   % Set non-zero values for u0
end

function [uN] = exactSolutionATtN(M)
% Compute the exact solution at time t_N = 2.4
x = linspace(-1,3,M+1);              % Define (M+1)-dim row vector containing x coordinates
uN = zeros(1,M+1);                   % Allocate memory for row vector uN
ind = (x>=1.9) & (x <= 2.9);         % Indices of grid points in [-0.5+2.4, 0.5+2.4]
uN(ind) = cos(pi*(x(ind)-2.4)).^2;   % Set non-zero values for uN
end

function [u] = LaxFriedrichs (M, N, u0, lambda)
% Implementation of the Lax-Friedrichs scheme
u = zeros(N+1, M+1);
u(1,:) = u0;
for n=2:(N+1)
    u(n, 2:M) = 0.5*(1-lambda)*u(n-1,3:(M+1)) + 0.5*(1+lambda)*u(n-1,1:(M-1));
    u(n,1) = 0;
    u(n,M+1) = u(n,M);
end
end

function [u] = LeapFrog(M, N, u0, lambda)
% Implementation of the leap frog scheme
u = zeros(N+1, M+1);
u(1,:) = u0;
% Use forward-time central space to compute the solution at t_1
u(2, 2:M) = u(1,2:M) - 0.5*lambda*u(1,3:(M+1)) + 0.5*lambda*u(1,1:(M-1));
u(2,1) = 0;
u(2,M+1) = u(2,M);
% The leap frog scheme starts from t_2
for n=3:(N+1)
    u(n, 2:M) = u(n-2,2:M) - lambda*( u(n-1,3:(M+1)) - u(n-1,1:(M-1)) );
    u(n,1) = 0;
    u(n,M+1) = u(n,M);
end
end

function [u] = BTFS(M, N, u0, lambda)
% Implementation of the backward-time forward-space scheme
u = zeros(N+1, M+1);
u(1,:) = u0;
% Set coefficient matrix. The matrix is the same for all time steps.
B = diag((1-lambda)*ones(M+1,1)) + diag(lambda*ones(M,1),1);
B(M+1,M+1) = 0;
B(M+1,1) = 1;
% time iteration
for n=2:(N+1)
    rhs = [u(n-1,1:M), 0]';
    u(n,:) = (B\rhs)';
end
end
Lax-Friedrichs scheme with lambda = 0.800000, h=1/10, the error ||u^N-U^N||_2 = 0.960910
Lax-Friedrichs scheme with lambda = 0.800000, h=1/20, the error ||u^N-U^N||_2 = 0.932571
Lax-Friedrichs scheme with lambda = 0.800000, h=1/40, the error ||u^N-U^N||_2 = 0.828951

Lax-Friedrichs scheme with lambda = 1.600000, h=1/10, the error ||u^N-U^N||_2 = 71.550798
Lax-Friedrichs scheme with lambda = 1.600000, h=1/20, the error ||u^N-U^N||_2 = 10173.330666
Lax-Friedrichs scheme with lambda = 1.600000, h=1/40, the error ||u^N-U^N||_2 = 2539788099.184201

Leap frog scheme with lambda = 0.800000, h=1/10, the error ||u^N-U^N||_2 = 0.520506
Leap frog scheme with lambda = 0.800000, h=1/20, the error ||u^N-U^N||_2 = 0.187264
Leap frog scheme with lambda = 0.800000, h=1/40, the error ||u^N-U^N||_2 = 0.072600

Leap frog scheme with lambda = 1.600000, h=1/10, the error ||u^N-U^N||_2 = 148015.026415
Leap frog scheme with lambda = 1.600000, h=1/20, the error ||u^N-U^N||_2 = 159297647325.648407
Leap frog scheme with lambda = 1.600000, h=1/40, the error ||u^N-U^N||_2 = 1408866869905907016269824.000000

BT-FS scheme with lambda = 0.800000, h=1/10, the error ||u^N-U^N||_2 = 1926.762333
BT-FS scheme with lambda = 0.800000, h=1/20, the error ||u^N-U^N||_2 = 424099960.804756
BT-FS scheme with lambda = 0.800000, h=1/40, the error ||u^N-U^N||_2 = 133364464863178162176.000000

BT-FS scheme with lambda = 1.600000, h=1/10, the error ||u^N-U^N||_2 = 1.032202
BT-FS scheme with lambda = 1.600000, h=1/20, the error ||u^N-U^N||_2 = 1.074747
BT-FS scheme with lambda = 1.600000, h=1/40, the error ||u^N-U^N||_2 = 1.008332