clc
clear all


a=-15;
b=15;
M = (b-a)*(64);           % Number of grids 
echo_T = 2.0;           % The time for a single echo
epsilon=0.5;       % epsilon 0.5 0.75 1.0
epsilon
tau=1/200;
N=echo_T./tau;


L = b-a;          % The length of the spatial interval

h=L/M;
% Initialize the spatial grid and the initial conditions
t1 = a : h : b - h;
[x, y] = meshgrid(t1);
%%%%%%%%%%%%%%%%%%%%initial values
phi1 = exp(-(x.^2 + y.^2) / 2);
phi2 = exp(-((x - 1).^2 + y.^2) / 2);
%phi2 = zeros(M);
%%%%%%%%%%%%%%%%%%%%electric potential
const = 4 * pi / sqrt(3);
% honeycomb lattice potential
V = cos(-const * x) + cos(const * (1 / 2 * x + sqrt(3) / 2 * y)) + ...
    cos(const * (1 / 2 * x - sqrt(3) / 2 * y));
A1 = zeros(size(V));
A2 = zeros(size(V));

k = 2*pi/L * [0 : 1 : M / 2 - 1, - M / 2 : 1 : -1]; % Wave number
[mu_1, mu_2] = meshgrid(k);
sigma1=[0,1;1,0];
sigma2=[0,-1i;1i,0];
sigma3=[1,0;0,-1];
I_2=eye(2);
% Initialize the Q1 and Q2 matrices
Q = zeros(2, 2, length(k),length(k));
Q1 = zeros(2, 2, length(k),length(k));
Q2 = zeros(2, 2, length(k),length(k));

for i = 1:length(k)
    for j=1:length(k)
        % Initialize the Q1 and Q2 matrices Gammal
        Gammal =mu_1(i,j)* epsilon*sigma1+ mu_2(i,j)* epsilon*sigma2 +sigma3;
        Q(:, :, i,j)=expm(-1i * tau / epsilon * Gammal);
        Gammal_inv=pinv(Gammal);
        %Initialize the Q1 and Q2 matrices Q1 Q2
        Q1(:, :, i,j) = -1i * epsilon * Gammal_inv * (I_2 - Q(:, :, i,j));
        Q2(:, :, i,j) = -1i * epsilon *tau* Gammal_inv + epsilon^2 * Gammal_inv * Gammal_inv * (I_2 - Q(:, :, i,j));
    end
end

Phi = zeros(2, M, M);
Phi(1,:,:) = phi1;
Phi(2,:,:) = phi2;

G = zeros(2, M, M);
G_hat = zeros(2, M, M);
G_hat_pre = zeros(2, M, M);


Phi_hat=zeros(2,M,M);
T=0;
plot2DTE(a, b, M, Phi,T,epsilon);

for n=1:N
    Phi_hat(1,:,:) = fft2(squeeze(Phi(1,:,:)));
    Phi_hat(2,:,:) = fft2(squeeze(Phi(2,:,:)));
% %     Calculate the nonlinear terms (vectorization instead of looping)
    phi1 = squeeze(Phi(1,:,:));
    phi2 = squeeze(Phi(2,:,:));
    G(1,:,:) = (V.*phi1 - (A1 - 1i*A2).*phi2)/epsilon;
    G(2,:,:) = (V.*phi2 - (A1 + 1i*A2).*phi1)/epsilon;
    
% %     Fourier transform of the nonlinear term
    G_hat(1,:,:) = fft2(squeeze(G(1,:,:)));
    G_hat(2,:,:) = fft2(squeeze(G(2,:,:)));

% %     Update the formula
    if n==1
        for i = 1:M
            for j = 1:M
                Phi_hat(:,i,j) = Q(:,:,i,j)*Phi_hat(:,i,j)-Q1(:,:,i,j)*G_hat(:,i,j).*1i;
            end
        end
        G_hat_pre=G_hat;
        Phi(1,:,:) = ifft2(squeeze(Phi_hat(1,:,:)));
        Phi(2,:,:) = ifft2(squeeze(Phi_hat(2,:,:)));
        continue;
    end
    for i = 1:M
            for j=1:M
                Phi_hat(:,i,j) = Q(:,:,i,j)*Phi_hat(:,i,j)-Q1(:,:,i,j)*G_hat(:,i,j).*1i-Q2(:,:,i,j)*((G_hat(:,i,j)-G_hat_pre(:,i,j))./tau).*1i;
            end
    end
    G_hat_pre=G_hat;
    Phi(1,:,:) = ifft2(squeeze(Phi_hat(1,:,:)));
    Phi(2,:,:) = ifft2(squeeze(Phi_hat(2,:,:)));
end
T=echo_T;
plot2DTE(a, b, M, Phi,T,epsilon);
for echo=1:3
    for n=1:N
        Phi_hat(1,:,:) = fft2(squeeze(Phi(1,:,:)));
        Phi_hat(2,:,:) = fft2(squeeze(Phi(2,:,:)));
    % %     Calculate the nonlinear terms (vectorization instead of looping)
        phi1 = squeeze(Phi(1,:,:));
        phi2 = squeeze(Phi(2,:,:));
        G(1,:,:) = (V.*phi1 - (A1 - 1i*A2).*phi2)/epsilon;
        G(2,:,:) = (V.*phi2 - (A1 + 1i*A2).*phi1)/epsilon;
        
    % %     Fourier transform of the nonlinear term
        G_hat(1,:,:) = fft2(squeeze(G(1,:,:)));
        G_hat(2,:,:) = fft2(squeeze(G(2,:,:)));
    
    % %     Update the formula
        for i = 1:M
                for j=1:M
                    Phi_hat(:,i,j) = Q(:,:,i,j)*Phi_hat(:,i,j)-Q1(:,:,i,j)*G_hat(:,i,j).*1i-Q2(:,:,i,j)*((G_hat(:,i,j)-G_hat_pre(:,i,j))./tau).*1i;
                end
        end
        G_hat_pre=G_hat;
        Phi(1,:,:) = ifft2(squeeze(Phi_hat(1,:,:)));
        Phi(2,:,:) = ifft2(squeeze(Phi_hat(2,:,:)));
    end
    T=T+echo_T;
    plot2DTE(a, b, M, Phi,T,epsilon);
end