%------------------------------------------------------------
% Tony Hyun Kim
% 6.638 PS2 Problem 1b
%   Propagation of solitons using NSE
%------------------------------------------------------------
clear all; close all;

%------------------------------------------------------------
% SETTINGS
%------------------------------------------------------------

% Time
Nt = 2^12;                  % Length of signal
t  = 200*linspace(-1,1,Nt); % Time vector [tau]
dt = t(2)-t(1);             % Sample time [tau]
Fs = 1/dt;                  % Sampling frequency [1/tau]

% Space
dz   = 0.001;
Zmax = 3;
Nz   = floor(Zmax/dz);
z = linspace(0,Zmax,Nz);

% Initial pulse
N = 1;
a = 1/sqrt(2);
ui = N*sech(a*t);

%------------------------------------------------------------
% COMPUTATION
%   Nothing more to touch
%------------------------------------------------------------

% Compute the sampled frequency grid
f = Fs/2*linspace(-1,1,Nt);
w = 2*pi*f;

% Also precompute the D/2 operator
D2 = exp(1i/2*dz*w.^2);

% Allocate solution space and initialize
%   Note that Uf will store fftshift'ed versions of fft results
U  = zeros(Nz,Nt); % Time domain
Uf = zeros(Nz,Nt); % Frequency domain

U(1,:)  = ui;
Uf(1,:) = fftshift(dt*fft(ui));
clear ui;

% Propagate solution
%   u1: ifft after applying first half of D
%   u2: apply N
%   u3: ifft after applying latter half of D
for k = 2:Nz
    u0 = U(k-1,:);
    u1 = 1/dt*ifft(ifftshift(D2.*Uf(k-1,:)));
    u2 = exp(-1i*dz*abs(u0).^2).*u1;
    u2f = fftshift(dt*fft(u2));
    u3 = 1/dt*ifft(ifftshift(D2.*u2f));

    U(k,:) = u3;
    Uf(k,:) = fftshift(dt*fft(U(k,:)));
end
clear u0 u1 u2 u2f u3 D2;

% Note that envelope amplitude is deduced by computing the magnitude
%   of the complex vector (i.e. keep the real part in exponent).
U = abs(U);
Uf = abs(Uf);

p1_DR; % Display results