% Solves for u using domain decomposition (matlab version) method.
% Scans for extrema using partitions and patches (see report on arxiv mentioned in README file if you need to understand the details).
% Constructs the valley network and plots u, V, W_E and the network.

clear all; close all; clc;

global posx posy g num

L = 50; W = 50; ff = 0.06; % L,W are system dimensions. For archaic reasons (that you don't need to worry about), they are in units of HALF \ell, the unit of distance used in the paper. So set them to TWICE the value you want in units of \ell. ff is the fill factor, called 'f' in the paper. The definition of this parameter is the same as in the paper, so it is easy to set.

box = 2; % Do not change! Converts from archaic units of length (\ell/2) to paper-consistent (\ell).
Lf = L/box; Wf = W/box; % These are now in \ell.
Px = 1; Py = 1; % Partition into subdomains. Use L/50 and W/50 for these two numbers, respectively, rounding up if fractions.
num = round(ff*W*L/4); % Total number of scatterers. The 4 here is hard-coded (don't change) due to the use of L and W and the unit thing above.

g = 21.33; % Scatterer height in units of E_0, as defined in the paper.

% Now choose random positions for the scatterers.
posx = randi(Lf,[1,num]);
posy = randi(Wf,[num,1]);
%load blah.mat; % This allows you to load a pre-saved noise realisation instead of generating a new one. 'blah.mat' should contain 'posx' and 'posy', generated for the same L,W,f!

% These are precision setting parameters for solving the PDE. These are sufficient to guarantee convergence and nice accuracy.
NMRsteps = 3; % The maximum number of mesh-refining steps
Hmax = 0.5; % Parameter for the coarse starting triangulation

partx = (0:1:Px)*Lf/Px;
party = (0:1:Py)*Wf/Py;
bounds = zeros(Px*Py,8);
row=1;
for k=1:Px
    for n=1:Py
        bounds(row,:) = [partx(k) partx(k+1) partx(k+1) partx(k) party(n) party(n) party(n+1) party(n+1)];
        row = row+1;
    end
end
blah = ones(Px*Py,1);
shape = [3*blah 4*blah bounds]';

names = char('SD1');
forms = 'SD1';
for k=1:Px*Py-1
    names = char(names,strcat('SD',int2str(k+1)));
    forms = strcat(forms,'+SD',int2str(k+1));
end
names = names';

geom = decsg(shape,forms,names);

left = find(geom(2,:)==0 & geom(3,:)==0);
right = find(geom(2,:)==Lf & geom(3,:)==Lf);
bot = find(geom(4,:)==0 & geom(5,:)==0);
top = find(geom(4,:)==Wf & geom(5,:)==Wf);
ext = [left right bot top];
blah = (1:1:size(geom,2));
int = setdiff(blah,ext);
geom = [geom(:,ext) geom(:,int)];

b = [1 1 1 1 1 1 48 48 49 48].';
blah = ones(1,2*(Px+Py));
[~,b] = meshgrid(blah,b);

if ff==0
    a = 0;
else
    a = @ScatPot2; % This defines the physical potential. Please go into this function (open the .m file provided) and set the scatterer width to the value you want!
end

c = 1;
f = 1;
time = [];

[p,e,t]=initmesh(geom,'Hmax',Hmax);
for step = 1:NMRsteps-1
    [p,e,t]=refinemesh(geom,p,e,t);
end

np = size(p,2);
cp = pdesdp(p,e,t); % Find common points
nc = length(cp);
C = zeros(nc,nc);
FC = zeros(nc,1);
        
in = cell(Px*Py,1);
cn = cell(Px*Py,1);
icn = cell(Px*Py,1);
Kn = cell(Px*Py,1);
Bn = cell(Px*Py,1);
an = cell(Px*Py,1);
fn = cell(Px*Py,1);
en = cell(Px*Py,1);
        
for k=1:Px*Py
    [in{k,1},cn{k,1}] = pdesdp(p,e,t,k);
    icn{k,1} = pdesubix(cp,cn{k,1});
    [K,F] = assempde(b,p,e,t,c,a,f,time,k);
    Kn{k,1} = K(in{k,1},in{k,1});
    d = symamd(Kn{k,1});
    blah = in{k,1};
    in{k,1} = blah(d);
    blah = Kn{k,1};
    Kn{k,1} = chol(blah(d,d));
    Bn{k,1} = K(cn{k,1},in{k,1});
    an{k,1} = Bn{k,1}/Kn{k,1};
    C(icn{k,1},icn{k,1}) = C(icn{k,1},icn{k,1})+K(cn{k,1},cn{k,1})-an{k,1}*an{k,1}';
    fn{k,1} = F(in{k,1});
    en{k,1} = Kn{k,1}'\fn{k,1};
    FC(icn{k,1}) = FC(icn{k,1})+F(cn{k,1})-an{k,1}*en{k,1};
end

u = zeros(np,1);
u(cp) = C\FC;
for k=1:Px*Py
    u(in{k,1}) = Kn{k,1}\(en{k,1}-an{k,1}'*u(cn{k,1}));
end

upde = pdeInterpolant(p,t,u);

if 1 % Plot W_E. x & y are in units of \ell, W_E in units of E_0.
Nx = 2*L;
Ny = 2*W;
x = linspace(0.5,Lf-0.5,2*Nx);
y = linspace(0.5,Wf-0.5,2*Ny);
[X,Y] = meshgrid(x,y);
ue = evaluate(upde,X,Y);
ue = reshape(ue,size(X));
figure; surface(x,y,1./ue); colormap('jet'); shading interp; grid on;
end

if 1 % Plot u. x & y are in units of \ell, u in units of 1/E_0.
Nx = 2*L;
Ny = 2*W;
x = (0:Lf/(2*Nx):Lf);
y = (0:Wf/(2*Ny):Wf);
[X,Y] = meshgrid(x,y);
ue = evaluate(upde,X,Y);
ue = reshape(ue,size(X));
figure; surface(x,y,ue); colormap('jet'); shading interp; grid on;
end

if 1 % Plot V. x & y are in units of \ell, V in units of E_0.
sig = 1/sqrt(2); % NOTE CAREFULLY: this sets scatterer width. The parameter 'sig' is related BUT NOT EQUAL to \sigma from the paper. It's a simple relation, though: sig/sqrt(2) = \sigma. Thus, to get \sigma=1/2, set sig to 1/sqrt(2), and to get \sigma=0.48, set sig to 0.68. Yes, you need to set it here as well as in ScatPot2.m. Annoying, but c'est la vie =)
x = linspace(0,Lf,2*Nx);
y = linspace(0,Wf,2*Ny);
[X,Y] = meshgrid(x,y);
V = 0;
for n=1:num
x0 = posx(n)-0.5;
y0 = posy(n)-0.5;
V = V + g*exp(-((X-x0).^2 + (Y-y0).^2)/sig^2);
end
figure; surface(x,y,V); colormap('jet'); shading interp; grid on;
end

subx = Lf/Px/4;
suby = Wf/Py/4;
numlines = size(geom,2) - 2*(Px+Py);

info = cell(Px*Py+numlines,9);

for k=1:Px*Py
   info{k,1} = [bounds(k,1:2); bounds(k,6:7)];
end

for k=1:numlines
    info{Px*Py+k,1} = [geom(2:3,2*(Px+Py)+k)'; geom(4:5,2*(Px+Py)+k)'];
end

figure; grid on; hold on; % Gradually plot the valley network.

for k=1:Px*Py+numlines
    region = info{k,1};
    x1 = min(region(1,:)); x2 = max(region(1,:));
    y1 = min(region(2,:)); y2 = max(region(2,:));
    if k>Px*Py
        if y1==y2
            y1 = y1 - 2*suby; y2 = y2 + 2*suby;
            x1 = x1 - subx; x2 = x2 + subx;
        elseif x1==x2
            x1 = x1 - 2*subx; x2 = x2 + 2*subx;
            %y1 = y1 - suby; y2 = y2 + suby; % to avoid overlap of patches!
        end
        
        if x1<0, x1=0; end
        if y1<0, y1=0; end
        if x2>Lf, x2=Lf; end
        if y2>Wf, y2=Wf; end
    end
    info{k,1} = [x1 x2 y1 y2];
    Nx = round(4*(x2-x1));
    Ny = round(4*(y2-y1));
    
    nx = (0:Nx+1);
    x = cos(pi*nx/(Nx+1));
    x = x + 1;
    x = x*(x2-x1)/2;
    x = x + x1;
    x = fliplr(x);
    
    ny = (0:Ny+1);
    y = cos(pi*ny/(Ny+1));
    y = y + 1;
    y = y*(y2-y1)/2;
    y = y + y1;
    y = fliplr(y);

    [X,Y] = meshgrid(x,y);
    ue = evaluate(upde,X,Y);
    ue = reshape(ue,size(X));
    uevec = ue.'; uevec = uevec(:);
    
    Dx = zeros(Nx+2,Nx+2);
    zilch = zeros(Nx+2,1);
    for n=1:Nx+2
        uni = zilch; uni(n) = 1;
        Dx(:,n) = fchd(uni).';
    end
    Dx = -2/(x2-x1)*Dx;
    Iy = speye(Ny+2);
    Dx = sparse(kron(Iy,Dx));

    Dy = zeros(Ny+2,Ny+2);
    zilch = zeros(Ny+2,1);
    for n=1:Ny+2
        uni = zilch; uni(n) = 1;
        Dy(:,n) = fchd(uni).';
    end
    Dy = -2/(y2-y1)*Dy;
    Ix = speye(Nx+2);
    Dy = sparse(kron(Dy,Ix));

    uxe = Dx*uevec;
    uye = Dy*uevec;
    uxxe = Dx*uxe;
    uyye = Dy*uye;
    uxye = Dy*uxe;
    uyxe = Dx*uye;
    
    uxe = reshape(uxe,[Nx+2,Ny+2]).';
    uye = reshape(uye,[Nx+2,Ny+2]).';
    uxxe = reshape(uxxe,[Nx+2,Ny+2]).';
    uyye = reshape(uyye,[Nx+2,Ny+2]).';
    uxye = reshape(uxye,[Nx+2,Ny+2]).';
    uyxe = reshape(uyxe,[Nx+2,Ny+2]).';
    
    u = chebfun2(ue,[x1 x2 y1 y2]);
    info{k,2} = u;
    %ux = diffx(u);
    ux = chebfun2(uxe,[x1 x2 y1 y2]);
    info{k,3} = ux;
    %uy = diffy(u);
    uy = chebfun2(uye,[x1 x2 y1 y2]);
    info{k,4} = uy;
    %uxx = diff(u,2,1);
    %uyy = diff(u,2,2);
    %uxy = diff(u,[1,1]);
    %uyx = uxy;
    uxx = chebfun2(uxxe,[x1 x2 y1 y2]);
    uyy = chebfun2(uyye,[x1 x2 y1 y2]);
    uxy = chebfun2(uxye,[x1 x2 y1 y2]);
    uyx = chebfun2(uyxe,[x1 x2 y1 y2]);
    
    extrema = roots(ux,uy);
    
    co = 1e-3;
    if x1==0
        blah = find(extrema(:,1)<x2*co);
        extrema(blah,:) = [];
    end
    if x2==Lf
        blah = find(extrema(:,1)>x2*(1-co));
        extrema(blah,:) = [];
    end
    if y1==0
        blah = find(extrema(:,2)<y2*co);
        extrema(blah,:) = [];
    end
    if y2==Wf
        blah = find(extrema(:,2)>y2*(1-co));
        extrema(blah,:) = [];
    end
    
    keep = [x1+subx, x2-subx y1+suby, y2-suby];
    if x1==0, keep(1) = x1; end
    if y1==0, keep(3) = y1; end
    if x2==Lf, keep(2) = x2; end
    if y2==Wf, keep(4) = y2; end
    info{k,5} = keep;
    
    e = 0.5;
    DerEval = [keep(1)-e keep(2)+e keep(3)-e keep(4)+e];
    if keep(1)==0, DerEval(1) = 0; end
    if keep(3)==0, DerEval(3) = 0; end
    if keep(2)==Lf, DerEval(2) = Lf; end
    if keep(4)==Wf, DerEval(4) = Wf; end
        
    uxx.domain = DerEval;
    uyy.domain = DerEval;
    uxy.domain = DerEval;
    uyx.domain = DerEval;
    
    chuck = find(extrema(:,1)<keep(1)-e | extrema(:,1)>keep(2)+e | extrema(:,2)<keep(3)-e | extrema(:,2)>keep(4)+e);
    extrema(chuck,:) = [];
    info{k,6} = extrema;
    
    if isempty(extrema)
        x = linspace(keep(1),keep(2),2*Nx+1);
        y = linspace(keep(3),keep(4),2*Ny+1);
        [X,Y] = meshgrid(x,y);
        uxe = feval(ux,X,Y);
        uye = feval(uy,X,Y);
        quiver(X,Y,uxe,uye);
        continue;
    end
    
    try 
        D = uxx.*uyy - uxy.*uyx;
    catch
        if DerEval(1)==0
            DerEval(1) = 0.5;
        end
        if DerEval(2)==Lf
            DerEval(2) = Lf-0.5;
        end
        uxx.domain = DerEval;
        uyy.domain = DerEval;
        uxy.domain = DerEval;
        uyx.domain = DerEval;
        chuck = find(extrema(:,1)<DerEval(1) | extrema(:,1)>DerEval(2));
        extrema(chuck,:) = [];
        info{k,6} = extrema;
        if isempty(extrema)
            x = linspace(keep(1),keep(2),2*Nx+1);
            y = linspace(keep(3),keep(4),2*Ny+1);
            [X,Y] = meshgrid(x,y);
            uxe = feval(ux,X,Y);
            uye = feval(uy,X,Y);
            quiver(X,Y,uxe,uye);
            continue;
        end
        try
            D = uxx.*uyy - uxy.*uyx;
        catch
            try
                D = uxx.*uyy - uxy.^2;
            catch
                D = uxx.*uyy - uyx.^2;
            end
        end
    end
    
    De = feval(D,extrema(:,1),extrema(:,2));
    uxxe = feval(uxx,extrema(:,1),extrema(:,2));
    uyye = feval(uyy,extrema(:,1),extrema(:,2));

    indy = find(De<0);
    saddles = extrema(indy,:);
    info{k,7} = saddles;
    
    indy = find(De>0 & uxxe>0 & uyye>0);
    minima = extrema(indy,:);
    info{k,8} = minima;

    indy = find(De>0 & uxxe<0 & uyye<0);
    maxima = extrema(indy,:);
    info{k,9} = maxima;
        
    x = linspace(keep(1),keep(2),2*Nx+1);
    y = linspace(keep(3),keep(4),2*Ny+1);
    [X,Y] = meshgrid(x,y);
    uxe = feval(ux,X,Y);
    uye = feval(uy,X,Y);
    quiver(X,Y,uxe,uye);
end

for k=1:Px*Py+numlines
    keep1 = info{k,5};
    extrema1 = info{k,6};
    saddles1 = info{k,7};
    minima1 = info{k,8};
    maxima1 = info{k,9};
    
    if ~isempty(extrema1)
        net1ext = find(extrema1(:,1)<keep1(1) | extrema1(:,1)>keep1(2) | extrema1(:,2)<keep1(3) | extrema1(:,2)>keep1(4));
    else
        net1ext = [];
    end
    if ~isempty(saddles1)
        net1sad = find(saddles1(:,1)<keep1(1) | saddles1(:,1)>keep1(2) | saddles1(:,2)<keep1(3) | saddles1(:,2)>keep1(4));
    else
        net1sad = [];
    end
    if ~isempty(minima1)
        net1min = find(minima1(:,1)<keep1(1) | minima1(:,1)>keep1(2) | minima1(:,2)<keep1(3) | minima1(:,2)>keep1(4));
    else
        net1min = [];
    end
    if ~isempty(maxima1)
        net1max = find(maxima1(:,1)<keep1(1) | maxima1(:,1)>keep1(2) | maxima1(:,2)<keep1(3) | maxima1(:,2)>keep1(4));
    else
        net1max = [];
    end
    
    if isempty(net1ext)
        continue;
    end
    
    chuckExt = [];
    chuckSad = [];
    chuckMax = [];
    chuckMin = [];
    
    for n=1:Px*Py+numlines
        if n==k
            continue;
        end
        
        keep2 = info{n,5};
        extrema2 = info{n,6};
        saddles2 = info{n,7};
        minima2 = info{n,8};
        maxima2 = info{n,9};
        
        if ~isempty(extrema2)
            main2ext = find(extrema2(:,1)>=keep2(1) & extrema2(:,1)<=keep2(2) & extrema2(:,2)>=keep2(3) & extrema2(:,2)<=keep2(4));
        else
            main2ext = [];
        end
        if ~isempty(saddles2)
            main2sad = find(saddles2(:,1)>=keep2(1) & saddles2(:,1)<=keep2(2) & saddles2(:,2)>=keep2(3) & saddles2(:,2)<=keep2(4));
        else
            main2sad = [];
        end
        if ~isempty(minima2)
            main2min = find(minima2(:,1)>=keep2(1) & minima2(:,1)<=keep2(2) & minima2(:,2)>=keep2(3) & minima2(:,2)<=keep2(4));
        else
            main2min = [];
        end
        if ~isempty(maxima2)
            main2max = find(maxima2(:,1)>=keep2(1) & maxima2(:,1)<=keep2(2) & maxima2(:,2)>=keep2(3) & maxima2(:,2)<=keep2(4));
        else
            main2max = [];
        end
        e = 0.01; % This is a precision parameter for how close two extrema can be togther before we believe that they are really distinct enteties and not the same thing picked up twice. This value gives sufficient accuracy.
        
        if ~isempty(net1ext) && ~isempty(main2ext)
            for j=1:size(net1ext,1)
                d = sqrt((extrema1(net1ext(j),1)-extrema2(main2ext,1)).^2 + (extrema1(net1ext(j),2)-extrema2(main2ext,2)).^2);
                blah = find(d<=e);
                if ~isempty(blah)
                    chuckExt = [chuckExt net1ext(j)];
                end
            end
        end
        
        if ~isempty(net1sad) && ~isempty(main2sad)
            for j=1:size(net1sad,1)
                d = sqrt((saddles1(net1sad(j),1)-saddles2(main2sad,1)).^2 + (saddles1(net1sad(j),2)-saddles2(main2sad,2)).^2);
                blah = find(d<=e);
                if ~isempty(blah)
                    chuckSad = [chuckSad net1sad(j)];
                end
            end
        end
        
        if ~isempty(net1min) && ~isempty(main2min)
            for j=1:size(net1min,1)
                d = sqrt((minima1(net1min(j),1)-minima2(main2min,1)).^2 + (minima1(net1min(j),2)-minima2(main2min,2)).^2);
                blah = find(d<=e);
                if ~isempty(blah)
                    chuckMin = [chuckMin net1min(j)];
                end
            end
        end
        
        if ~isempty(net1max) && ~isempty(main2max)
            for j=1:size(net1max,1)
                d = sqrt((maxima1(net1max(j),1)-maxima2(main2max,1)).^2 + (maxima1(net1max(j),2)-maxima2(main2max,2)).^2);
                blah = find(d<=e);
                if ~isempty(blah)
                    chuckMax = [chuckMax net1max(j)];
                end
            end
        end
    end
    
    extrema1(chuckExt,:) = [];
    saddles1(chuckSad,:) = [];
    minima1(chuckMin,:) = [];
    maxima1(chuckMax,:) = [];
    
    info{k,6} = extrema1;
    info{k,7} = saddles1;
    info{k,8} = minima1;
    info{k,9} = maxima1;
    
end

extrema = [];
saddles = [];
minima = [];
maxima = [];

for k=1:Px*Py+numlines
    extrema = [extrema; info{k,6}];
    saddles = [saddles; info{k,7}];
    minima = [minima; info{k,8}]; 
    maxima = [maxima; info{k,9}];
end

plot(extrema(:,1),extrema(:,2),'ko');
plot(maxima(:,1),maxima(:,2),'bx');
plot(minima(:,1),minima(:,2),'rx');
plot(saddles(:,1),saddles(:,2),'gx');

clear NMRsteps Hmax row blah shape names forms geom SC uni nx ny Ix Iy %upde
clear left right bot top ext int b a c f time step p e t F geom chuck DerEval
clear in cn icn Kn Bn an fn en K d zilch partx party np cp nc C FC bounds De
clear Dx Dy uevec ue uxe uye uxxe uyye uxye uyxe x1 x2 y1 y2 keep region u D
clear ux uy uxx uyy uxy uyx co subx suby indy numlines x y X Y keep1 keep2
clear saddles1 saddles2 minima1 minima2 maxima1 maxima2 extrema1 extrema2
clear main2ext main2sad main2max main2min net1ext net1sad net1max net1min
clear chuckExt chuckSad chuckMax chuckMin

keep = [];
for k=1:size(info,1)
    keep = [keep; info{k,5}];
end

e = 0.01; % This is a step size. This value gives sufficient accuracy. 
C = cell(length(saddles),2);
for k=1:length(saddles)
    x0 = saddles(k,1); y0 = saddles(k,2);
    for j=1:2
        traj = zeros(3,10001);
        traj(1,1) = x0; traj(2,1) = y0;
        for n = 2:10001
            
            part = find(keep(:,1)<=traj(1,n-1) & keep(:,2)>=traj(1,n-1) & keep(:,3)<=traj(2,n-1) & keep(:,4)>=traj(2,n-1));
            part(2:end) =  [];
            u = info{part,2};
            if n==2
                traj(3,n-1) = feval(u,traj(1,n-1),traj(2,n-1));
            end
            ux = info{part,3};
            uy = info{part,4};
            
            theta = atan2(feval(uy,traj(1,n-1),traj(2,n-1)),feval(ux,traj(1,n-1),traj(2,n-1)));
            
            if j==2 && n==2
                traj(1,n) = traj(1,n-1) + e*cos(theta);
                traj(2,n) = traj(2,n-1) + e*sin(theta);
            else
                traj(1,n) = traj(1,n-1) - e*cos(theta);
                traj(2,n) = traj(2,n-1) - e*sin(theta);
            end
            
            traj(3,n) = feval(u,traj(1,n),traj(2,n));
            
            dist2min = (traj(1,n) - minima(:,1)).^2 + (traj(2,n) - minima(:,2)).^2;
            dist2min = sqrt(dist2min);
            indy = find(dist2min<e);
            dist2max = (traj(1,n) - maxima(:,1)).^2 + (traj(2,n) - maxima(:,2)).^2;
            dist2max = sqrt(dist2max);
            indy2 = find(dist2max<e);
            if ~isempty(indy) || traj(1,n)<e || traj(1,n)>Lf-e || traj(2,n)<e || traj(2,n)>Wf-e
                break;
            elseif ~isempty(indy2)
                warning('Aaaaaah! Terminating at a maximum!')
                break;
            end
            if n==10001
                warning('Trajectory used up max length & did not reach extremum!')
            end
        end
        traj(:,n+1:end) = [];
        C{k,j} = traj;
    end
end

for k=1:length(C)
    traj = C{k,1};
    plot3(traj(1,:),traj(2,:),traj(3,:),'b.');
    traj = C{k,2};
    plot3(traj(1,:),traj(2,:),traj(3,:),'r.');
end

clear keep part u ux uy dist2min dist2max indy indy2
%clear info extrema saddles minima maxima
