clear all
close all
clc

% choose trajectory durations.  Tau = 10s, 5s, 3s
tau = [10,5,3];

% corresponding filenames.
fname = {'Figure5_Tau10s.mat',...
'Figure5_Tau5s.mat',...
'Figure5_Tau3s.mat'};

% each .mat file contains 
% a vector t, which contains the timepoints in seconds from the full simulation
% a matrix mov2D_psi which is a set of 2D wavefunctions, sampled at regular intervals over the full simulation. 
% The time between sampled wave functions is [max value of t]/[size of third dimension of mov2D_psi - 1].

 % timepoints to plot
ttt = [.66,.66,.66];   %time / tspiral

% centering for the zoomed in regions.
x1_ar = 1.0e-04 * [0  -.05 -.05];   
y1_ar = 1.0e-04 * [-.3 -.3 -.3];

% set up where to put the plots.
Bshift = 55;
Ashift = 65;

% colormap limits
clims1 = [0.2 Inf];
clims2 = clims1;

%fontsizes
f1 = 22;  % fontsize for timing labels.   
f2 = 20;  % fontsize for colormap numbers.  
f3 = 26;  % fontsize for colormap labels.    

h3 = figure(103);
set(h3, 'position', [1       1    150*4+5 150*length(ttt)+20+Bshift]);
set(h3,'color','w');

pctr = 0;

for pp = 1:length(ttt)
    load(fname{pp})  % load file.

    %% Constants
    % 
    a = 100.44*5.29e-11;           % approx scattering length for Rb87 in meters.
    m = 1.5e-25;                % m = mass of Rb87 in kg
    hbar = 1.05457e-34;         % Planck's constant in J*s

    % need to calculate the local speed of sound from the BEC parameters and
    % the atomic density.
    % the local speed of sound is used to scale the velocity profile in column 3.
    % a global speed of sound is used to scale the velocity profile in column 4.

    % Start with the 3D experimental parameters

    omegar = 2*pi*8;    % radial trap freq is 8 Hz. 
    omegaz = 2*pi*90;           % axial trap freq is 90 Hz.
    omega = (omegar^2*omegaz)^(1/3);  % geometric mean.

    % 
    lr = sqrt(hbar/m/omegar);  % H.O. lengths
    lz = sqrt(hbar/m/omegaz);
    l = (lr^2*lz)^(1/3);
    % 
    u = (15*N*a/l)^0.4 * 0.5 * hbar * omega;  % 3D chemical potential 
    % 
    g = 4*pi*hbar^2*a/m;       % 3D  interaction prefactor
    % 
    c0 = sqrt(u/m);            % bulk speed of sound.

    % 2D effective parameters for numerics
    N2D = (15*N*a/l)^0.6/16 * sqrt(pi)/a * sqrt(2 + sqrt(2)) * sqrt(hbar*omega^3/m/omegar^4);
    omegar2D = 2^(-.25) * omegar;
    lz2D = sqrt(pi*hbar/m/omega) * sqrt(2)*sqrt(2+sqrt(2)) * (15*N*a/l)^(-0.2)/sqrt(2*pi); 

    g2D = g/sqrt(2*pi)/lz2D;

    G2D = g2D*N2D;             % nonlinear prefactor in 2D.

    %
    %   Set up grids and time steps
    %   All units are in mks.
    %

    %% space
    %
    xmax    = 120e-6;             % x grid size in meters
    ymax    = xmax;
    nxy      = 512;             % number of grid points  

    nmid  = floor(nxy/2);
    v = 0:1:nxy-1;
    dx = xmax/(nxy-1);
    dy = ymax/(nxy-1);
    [x,y] = meshgrid(v,v);
    x = -xmax/2 + x*dx;
    y = -ymax/2 + y*dy;
    xar = -xmax/2 + v*dx;
    yar = -ymax/2 + v*dy;

    %% time
    %
    tspiral = ts;                %time for spiral
    dt = t(2)-t(1);
    nt = length(t);
    ncycle = floor(nt/500);

    %% make figure

    [sy, sx, sz] = size(mov2D_psi);

    pctr = pctr+1;

    for qq = round(ttt(pp)*ts/dt/ncycle)+1    
        % atomic density
        n =  abs(mov2D_psi(:,:,qq)).^2;
        nmax = max(max(n));
        
        % set up masks to set regions nearly zero density to zero. 
        % this just makes the plots less messy.
        phmask1 = ones(size(n)).*(n/nmax > 1e-6);
        phmask2 = ones(size(n)).*(n/nmax > 1e-3);
        vmask = ones(size(n)).*(n/nmax > 5e-4);
        vmask1 = phmask2;
        %
        n = n.*phmask2;
        %
        ph = (angle(mov2D_psi(:,:,qq))+pi()).*phmask1;
        ph1 = ph.*phmask1;

        % calculate velocity profile from the phase.
        [gx, gy] = gradient(ph1,dx,dx);
        vx = gx*hbar/m.*phmask2;
        vy = gy*hbar/m.*phmask2;
        vx1 = vx;
        vy1 = vy;
        vv = sqrt(vx1.^2 + vy1.^2);

        % mask the phase a bit more.
        ph2 = ph1.*phmask2;

        % find the area that will be set to black in each plot.
        n(n == 0) = NaN;
        vmask1(vmask1 == 0) = NaN;
        ph2(ph2 == 0) = NaN;
        
        imAlpha=ones(size(n));
        imAlpha(isnan(n))=0;
        imAlpha1=ones(size(ph2));
        imAlpha1(isnan(ph2))=0;
        imAlpha2=ones(size(vmask1));
        imAlpha2(isnan(vmask1))=0;

        %% smoothing algorithm
        % method is for each point where v_x or v_y > 0.001, take the average
        % of the legit data points around that point and replace the value with
        % that data point.  note: have a copy so that I don't end up
        % accidentally making streaks.    

        clear cy
        clear cx
        [cy, cx] = find(abs(vy1) > 0.001);    % look at vy
        
        for cc = 1:length(cx)
            vxr = cx(cc)-1:cx(cc)+1;  
            vyr = cy(cc)-1:cy(cc)+1;

            vytest = (vy1(vyr,vxr))*1e6;   % smoothing vy
            vylegit = find(abs(vytest) ~= 0 & abs(vytest) < 1e3);
            vynew = mean(vytest(vylegit));    % in um/s
            vy(cy(cc),cx(cc)) = vynew*1e-6;   % in m/s
        end

        clear cy
        clear cx
        [cy, cx] = find(abs(vx1) > 0.001);      % look at vx

        for cc = 1:length(cx)
            vxr = cx(cc)-1:cx(cc)+1;  
            vyr = cy(cc)-1:cy(cc)+1;

            vxtest = (vx1(vyr,vxr))*1e6;    % smoothing vx
            vxlegit = find(abs(vxtest) ~= 0 & abs(vxtest) < 1e3);
            vxnew = mean(vxtest(vxlegit));    % in um/s
            vx(cy(cc),cx(cc)) = vxnew*1e-6;   % in m/s
        end

        %% velocity field requires both magnitude and direction.
        
        v = sqrt(vx.^2 + vy.^2);        % magnitude of the velocity

        % gets rid of long arrows on quiver plot.
        vx = vx.*vmask.*(v<.3*c0);
        vy = vy.*vmask.*(v<.3*c0);   % unscaled

        % downsample for quiver plot.
        s = 12*nxy/256;
        %
        nn = downsample(downsample(n,s)',s)';
        xx = downsample(downsample(x,s)',s)';
        yy = downsample(downsample(y,s)',s)';
        vxx = downsample(downsample(vx,s)',s)';
        vyy = downsample(downsample(vy,s)',s)';

        x1 = x1_ar(pctr);
        y1 = y1_ar(pctr);
        xmax1 = 2.5e-5;

        % make figure
        figure(103)
        % plot density profile
        axes('Units', 'pixels', 'position', [5, 20+150*(length(ttt)-pctr)+5+Bshift, 145, 145]) 
        imagesc(xar-x1, yar-y1 , n/nmax,'AlphaData',imAlpha)
        set(gca,'color',0*[1 1 1]);
        set(gca, 'CLim', [0 1])
        axis image
        axis xy
        xlim([-xmax1 xmax1])
        ylim([-xmax1 xmax1])
        set(gca,'XTickLabel','')
        set(gca,'YTickLabel','')
        hold on

        str = ['\bf{$\tau_\mathrm{s} = ' num2str(tspiral) ' \, \mathrm{s}$}'];
        text(-xmax1+dx*4, xmax1 - dx*20,str, 'Interpreter', 'latex', 'fontsize', f1, 'color', 'black')

        hold off

        % add label / colorbar at the bottom of each column.
        if pp == 1
            cax = colorbar('SouthOutside');
            set(cax, 'Units', 'pixels', 'position', [15, Ashift, 125, 10], 'fontsize', f2);
            cax.Label.Interpreter = 'latex';
            cax.TickLabelInterpreter = 'latex';
            cax.Label.String = '{$n/n_\mathrm{max}$}';
            cax.Label.FontSize = f3;
        end

        % plot phase
        axes('Units', 'pixels', 'position', [150+5, 20+150*(length(ttt)-pctr)+5+Bshift, 145, 145]) % 121 is for example
        imagesc(xar-x1, yar-y1, ph2/2/pi,'AlphaData',imAlpha1)
        set(gca, 'CLim', [0 1])
        axis xy
        axis image
        xlim([-xmax1 xmax1])
        ylim([-xmax1 xmax1])
        set(gca,'color',0*[1 1 1]);
        set(gca,'XTickLabel','')
        set(gca,'YTickLabel','')

        % add label / colorbar at the bottom of each column.
        if pp == 1
            cax = colorbar('SouthOutside');
            set(cax, 'Units', 'pixels', 'position', [150+15, Ashift, 125, 10], 'fontsize', f2);
            cax.Label.Interpreter = 'latex';
            cax.TickLabelInterpreter = 'latex';
            cax.Label.String = '{$\phi/2\pi$}';
            cax.Label.FontSize = f3;
        end

        % plot speed profile
        axes('Units', 'pixels', 'position', [150*2+5, 20+150*(length(ttt)-pctr)+5+Bshift, 145, 145]) % 121 is for example
        imagesc(xar-x1, yar-y1, v./sqrt(g2D*n*N2D*sqrt(2)/m),'AlphaData',imAlpha2)
        set(gca, 'CLim', [0 6])
        axis xy
        axis image
        xlim([-xmax1 xmax1])
        ylim([-xmax1 xmax1])
        set(gca,'color',0*[1 1 1]);
        set(gca,'XTickLabel','')
        set(gca,'YTickLabel','')

        if pp == 1
            cax = colorbar('SouthOutside');
            set(cax, 'Units', 'pixels', 'position', [150*2+15, Ashift, 125, 10], 'fontsize', f2);
            cax.Label.Interpreter = 'latex';
            cax.TickLabelInterpreter = 'latex';
            cax.Label.String = '{$v/c_\mathrm{local}$}';
            cax.Label.FontSize = f3;
        end

        % plot zoomed in speed profile
        axes('Units', 'pixels', 'position', [150*3+5, 20+150*(length(ttt)-pctr)+5+Bshift, 145, 145]);
        imagesc(xar-x1, yar-y1, v/c0.*(v/c0<1),'AlphaData',imAlpha2)
        axis xy
        axis image
        set(gca, 'CLim', [0 .5])
        set(gca,'color',0*[1 1 1]);
        set(gca,'XTickLabel','')
        set(gca,'YTickLabel','')

        % now add quiver plot to show the direction of flow
        scalefactor = .00003;   % scales the length of the arrows.
        quiver(xx-x1, yy-y1, vxx/c0*scalefactor, vyy/c0*scalefactor, 'AutoScale', 'off');
        hax = gca; %get the axis handle
        axis xy;
        axis image %plot the quiver to see the dimensions of the plot
        imagesc(hax.XLim,hax.YLim,v/c0.*(v/c0<1), 'AlphaData',imAlpha2); %plot the image within the axis limits
        hold on; %enable plotting overwrite
        quiver(xx-x1, yy-y1, vxx/c0*scalefactor, vyy/c0*scalefactor, 'color', 'white', 'linewidth', 1, 'AutoScale', 'off')
        axis xy
        axis image
        hold off

        xlim([-xmax1 xmax1])
        ylim([-xmax1 xmax1])
        set(gca, 'CLim', [0 .5])
        set(gca,'color',0*[1 1 1]);
        set(gca,'XTickLabel','')
        set(gca,'YTickLabel','')

        % add label / colorbar at the bottom of each column.
        if pp == 1
            cax = colorbar('SouthOutside');
            set(cax, 'Units', 'pixels', 'position', [150*3+15, Ashift, 125, 10], 'fontsize', f2);
            cax.Label.Interpreter = 'latex';
            cax.TickLabelInterpreter = 'latex';
            cax.Label.String = '{$v/c_0$}';
            cax.Label.FontSize = f3;
        end

    end
end
%%
% print(h3,'-dpng', ['Figure_SimsImages_v0'])
% print(h3,'-dsvg', ['Figure_SimsImages_v0'])
% print(h3,'-dpdf', ['Figure_SimsImages_v0'])