%% INPUT SPECIFICATIONS (TO BE SET BY USER)
% Please, set the values of these variables before running the code. For a
% description of the variables see the supplementary material provided with
% the paper.
    clear;
    close all;
    clc;
    %Do you want the code to report every network increment? suggested=true
    report_iterations = true;
    %Do you want to save the results (as matlab .mat file)? suggested=true
    save_on_file = true;
    %The workspace is saved as backup every "save_after_nIters" iterations.
    %Specify a positive value for "save_after_nIters", please.
    save_after_nIters = 100;
    %Which catchment to consider? (Insert number from 1 to 17, according to the following list)
    % catchments_list = [1="astico"; 2="borsoia"; 3="coweeta12"; 4="coweeta33"; ...
    %                   5="coweeta40"; 6="fernow14"; 7="fernow16"; 8="fernow37"; ...
    %                   9="focobon"; 10="leogra"; 11="montecalvello"; 12="poverty25"; ...
    %                   13="poverty35"; 14="turboloEst"; 15="turboloOvest"; 16="southForkPotts"; ...
    %                   17="rietholzback"];
    catchment_id = 2;
    %Do you want to restart from the end of a previously performed analysis? suggested=false
    restart = false;
    %If starting a new analysis, do you want to compute the info on morphology and hillslopes (or just to analize the
    %network expansion and its effects)? suggested=true
    computeAllInfo = true;
    %If computing the morphology indexes, do you want to preserve any previous (partial) computation of the elliptic
    %eccentricity? suggested=false
    preserve_eccentricity = false;
% END OF INPUT SPECIFICATIONS
%-------------------------------------------------------------------------

% BRIEF DESCRIPTION OF THE RELEVANT VARIABLES DEFINED BY THE CODE AND PROVIDED AS OUTPUTS:
% 1) dx, is the size of the pixel side (in meters) of the DTM;
% 2) network_Lengths, is an array that stores the length of the active
% channel network (L) at every iteration, so the i-th component is the
% length of the i-th considered spatial configuration of the network;
% 3) distributions_w0, is a matrix the i-th column of which represents the
% values of the empirical hillslope length distribution conditioned to the fact that the
% channel network is in its i-th configuration;
% 4) distributions_support, is a matrix the i-th column of which represents the
% hillslope length values at which the hillslope length distribution is evaluated;
% 5) basinMorphologyIndeces, is a matlab structure containing some indexes
% related to the morphology of the catchment.
% 
% For more info, refer to the supplementary information provided with the
% paper.


%% IMPORT DATA, MORPHOLOGY INDEXES AND INITIALIZATION
% Section for importing the DTM (or pevious analysis), computing the
% morphology indexes and initialize all the variables needed by the code.
% Operations are performed automatically, the user is not required to
% modify the following code.
% The section ends at line 395.

%-------------------------------------------------------------------------
% BASIC DATA ARRAY
catchments_list =   ["astico"; "borsoia"; "coweeta12"; "coweeta33"; "coweeta40"; "fernow14"; "fernow16"; "fernow37"; "focobon"; "leogra"; "montecalvello"; "poverty25"; "poverty35"; "turboloEst"; "turboloOvest"; "southForkPotts"; "rietholzback"; "fuciade"; "seugne"; "hubbardBrook"; "valfredda"; "cordevole"];
thresh_lower_lims = [200;      1000;      1000;        1000;        1000;        50;         50;         200;        250;       1000;      1000;           5000;        4000;        300;          300;            500;              10000;          21500000;  3463;     31000000;       50000;       1000000];
thresh_upper_lims = [0;        0;         0;           0;           0;           0;          0;          0;          0;         0;         0;              0;           0;           0;            0;              0;                0;              22000000;  0;        31341400;       501400;      2000000];
%DTM resolutions in m
dxDTMs =            [50;       5;         1;           1;           1;           3;          3;          3;          20;        25;        5;              1;           0.9;         5;            5;              3;                2;              0.5;       26.87;    1;              1;           10];
%DTM area in m^2
AreasDTMs =    1000*[136470;   13790;     120;         310;         410;         140;        160;        370;        8420;      67410;     3750;           250;         340;         630;          460;            730;              3120;           5600;      911100;   31340;          5270;        832290];

% SELECTION OF BASIC DATA REFERRED TO THE PROPER CATCHMENT
nameCatchment = catchments_list(catchment_id);
dx = dxDTMs(catchment_id);
areaCatchment = AreasDTMs(catchment_id);
threshold_lower_lim = thresh_lower_lims(catchment_id);
threshold_startFrom = thresh_upper_lims(catchment_id);

global step;
step = 0.5;       % Step for the hillslope distribution along \ell

if (restart)

    disp("LOADING THE PREVIOUS ANALYSIS BACKUP FOR CATCHMENT: "+upper(nameCatchment));
    threshold_lower_lim_backup = threshold_lower_lim;
    if (save_on_file)
        restore_path = cd();
        % mkdir("Results_"+nameCatchment);
        cd("Results_"+nameCatchment);
        % mkdir("Output_files");
        cd("Output_files");
    end
    load("Workspace_Backup.mat");
    %SETTING PLOT ENVIRONMENT
    figure;
    ax=axes();
    MapPlot(ax,A);
    figure;
    ax=axes();
    MapPlot(ax,Ap);
    figure;
    ax=axes();
    MapPlot(ax,Aad8);
    figure;
    ax=axes();
    geo_net_pic = zeros(size(Aad8));
    geo_net_old = geo_net;
    geo_net_pic(geo_net_old) = 2;
    %Computing again the geomorphic network (w.r.t. a new lower threshold
    %maybe)
    threshold_lower_lim = threshold_lower_lim_backup;
    clear threshold_lower_lim_backup;
    geo_net = Aad8 >= threshold_lower_lim;
    geo_T = sum(geo_net,"all");
    %Still plotting
    new_geo = geo_net & (~geo_net_old);
    geo_net_pic(new_geo) = 1;
    MapPlot(ax,geo_net_pic);
    %Enlarge loaded variables according to the new value of geo_T
    increase = geo_T - numel(network_Lengths);
    network_Lengths = [network_Lengths nan(1,increase)];
    branching_Index = [branching_Index nan(1,increase)];
    thresholdArea_alongIterations = [thresholdArea_alongIterations nan(1,increase)];
    topographicWetnessIndex_Zanetti = [topographicWetnessIndex_Zanetti nan(1,increase)];
    distributions_support = [distributions_support nan(size(distributions_support,1),increase)];
    distributions_alpha = [distributions_alpha nan(size(distributions_alpha,1),increase)];
    distributions_w0 = [distributions_w0 nan(size(distributions_w0,1),increase)];
    distributions_w = [distributions_w nan(size(distributions_w,1),increase)];
    clear increase;
    disp(" ");

else

    %IMPORT AND PRELIMINARY OPERATIONS:
    %LOADING DIGITAL TERRAIN MODELS
    disp("READING THE RASTER FILE FOR CATCHMENT: "+upper(nameCatchment));
    [A,R] = readgeoraster("Catchments\"+nameCatchment+"\DTM.tif");
    [Ap,Rp] = readgeoraster("Catchments\"+nameCatchment+"\p.tif");
    [Aad8,Rad8] = readgeoraster("Catchments\"+nameCatchment+"\ad8.tif");
    disp(" ");
    
    %DEFINING THE CATCHMENT AND THE GEOMORPHIC NETWORK
    disp("IDENTIFING THE LARGEST CATCHMENT AND THE GEOMORPHIC NETWORK");
    geo_net = Aad8 >= threshold_lower_lim;
    geo_T = sum(geo_net,"all");
    outlet_id = find_outlet(Aad8);
    threshold_upper_lim = Aad8(outlet_id);
    myCatchment = ContributingCatchment(Ap,outlet_id);
    %Corrections
    A(~myCatchment) = nan;
    Ap(~myCatchment) = nan;
    Aad8(~myCatchment) = nan;
    totpixel = sum(myCatchment,"all");
    areaCheck = totpixel * dx*dx;
    inputAreaBackup = -1;
    if (abs(areaCheck-areaCatchment)>dx*dx)
        disp("   WARNING: input area seems to differ from the measured value.");
        disp("            Swithcing to the measured value: "+areaCheck);
        inputAreaBackup = areaCatchment;
        areaCatchment = areaCheck;
    end
    disp(" ");

    if (computeAllInfo)
        %MORPHOLOGY INDECES OF THE CATCHMENT
        disp("ANALYZING THE MORPHOLOGY OF THE CATCHMENT: "+upper(nameCatchment));
    
        %Gravelius coefficient (from Sassolat, 2018)
        disp("   Computing GRAVELIUS COEFFICIENT");
        warningGravelius_resolution = false;
        if (abs(dx-sqrt(areaCatchment)/10)>dx)
            disp("   WARNING: resolution of DTM may not be good for the Gravelius coefficient.");
            disp("            Used resolution: "+dx);
            disp("            Best resolution: "+sqrt(areaCatchment)/10);
            warningGravelius_resolution = true;
        end
        perimeter_map = false(size(myCatchment));
        perimeter_length = 0;
        for row=1:size(myCatchment,1)
            for col=1:size(myCatchment,2)
                if (myCatchment(row,col))
                    if ((row==1)||(row==size(myCatchment,1))||(col==1)||(col==size(myCatchment,2)))
                        perimeter_map(row,col) = true;
                        perimeter_length = perimeter_length + dx;
                    else
                        if (~myCatchment(row+1,col+1))
                            perimeter_map(row,col) = true;
                        end
                        if (~myCatchment(row,col+1))
                            perimeter_map(row,col) = true;
                            perimeter_length = perimeter_length + dx;
                        end
                        if (~myCatchment(row-1,col+1))
                            perimeter_map(row,col) = true;
                        end
                        if (~myCatchment(row-1,col))
                            perimeter_map(row,col) = true;
                            perimeter_length = perimeter_length + dx;
                        end
                        if (~myCatchment(row-1,col-1))
                            perimeter_map(row,col) = true;
                        end
                        if (~myCatchment(row,col-1))
                            perimeter_map(row,col) = true;
                            perimeter_length = perimeter_length + dx;
                        end
                        if (~myCatchment(row+1,col-1))
                            perimeter_map(row,col) = true;
                        end
                        if (~myCatchment(row+1,col))
                            perimeter_map(row,col) = true;
                            perimeter_length = perimeter_length + dx;
                        end
                    end
                end
            end
        end
        graveliusCoefficient = 0.5*perimeter_length/sqrt(pi*areaCatchment);

        %Elliptic eccentricity (from Betterle, Botter 2021)
        disp("   Computing ELLIPTIC ECCENTRICITY");
        maxdist = 0;
        tmpdist = 0;
        rowstart = 1;
    
        if (preserve_eccentricity)
            orig_path = cd();
            cd("Results_"+nameCatchment);
            cd("Output_files");
            load("Workspace_Backup.mat","maxdist","row");
            cd(orig_path);
            rowstart = row;
            disp("      Restarting from row "+rowstart+" out of "+size(myCatchment,1));
            disp("         Inherited a maxdist = "+maxdist);
        end
    
        for row=rowstart:size(myCatchment,1)
            if (mod(row,floor(size(myCatchment,1)/20))==0)
                disp("      Completed "+row+" rows out of "+size(myCatchment,1));
            end
            for col=1:size(myCatchment,2)
                if (perimeter_map(row,col))
                    for row2=row:size(myCatchment,1)
                        dr2 = (row2-row)^2;
                        dc = floor( sqrt( max(maxdist^2-dr2, 0) ) );
                        ssStop = col-dc;
                        ssStart = col+dc;
                        for col2=1:ssStop
                            if (perimeter_map(row2,col2))
                                tmpdist = sqrt( dr2 + (col2-col)^2 );
                                if (tmpdist>maxdist)
                                    maxdist = tmpdist;
                                end
                            end
                        end
                        for col2=ssStart:size(myCatchment,2)
                            if (perimeter_map(row2,col2))
                                tmpdist = sqrt( dr2 + (col2-col)^2 );
                                if (tmpdist>maxdist)
                                    maxdist = tmpdist;
                                end
                            end
                        end
                    end
                end
            end
        end
        majorSemiaxis_ellipse = maxdist*dx/2;
        minorSemiaxis_ellipse = areaCatchment/pi/majorSemiaxis_ellipse;
        aspectRatio_ellipse = minorSemiaxis_ellipse/majorSemiaxis_ellipse;
        eccentricity_ellipse = sqrt( 1-(aspectRatio_ellipse)^2 );
    
        %Circular elongation ratio (from Schumm, 1956, p. 612)
        disp("   Computing SCHUMM ELONGATION RATIO");
        elongationRatio_Schumm = 0.5*maxdist*dx/sqrt(areaCatchment/pi);
    
        %Rectangular aspect ratio
        disp("   Computing RECTANGULAR ASPECT RATIO");
        majorSide_rectangle = maxdist*dx;
        minorSide_rectangle = areaCatchment/majorSide_rectangle;
        aspectRatio_rectangle = minorSide_rectangle/majorSide_rectangle;
    
        %Convergence angle (similar to Castelltort, 2009, p. 1 eq. 1)
        disp("   Computing CASTELLTORT CONVERGENCE ANGLE");
        convergenceAngle_Castelltort_degrees = atand(aspectRatio_rectangle/2);
        convergenceAngle_Castelltort_radians = atan(aspectRatio_rectangle/2);
    
        basinMorphologyIndeces.convergenceAngle_Castelltort_degrees = convergenceAngle_Castelltort_degrees;
        basinMorphologyIndeces.convergenceAngle_Castelltort_radians = convergenceAngle_Castelltort_radians;
        basinMorphologyIndeces.aspectRatio_rectangle = aspectRatio_rectangle;
        basinMorphologyIndeces.elongationRatio_Schumm = elongationRatio_Schumm;
        basinMorphologyIndeces.graveliusCoefficient = graveliusCoefficient;
        basinMorphologyIndeces.eccentricity_ellipse = eccentricity_ellipse;
        disp(" ");

    end
    
    disp("ADDITIONAL PRELIMINARIES");
    %SETTING PLOT ENVIRONMENT
    figure;
    ax=axes();
    MapPlot(ax,A);
    figure;
    ax=axes();
    MapPlot(ax,Ap);
    figure;
    ax=axes();
    MapPlot(ax,Aad8);
    figure;
    ax=axes();
    geo_net_pic = zeros(size(Aad8));
    geo_net_pic(geo_net) = 1;
    MapPlot(ax,geo_net_pic);
    
    % SAVING THE VARIABLE TOTPIXEL
    if (save_on_file)
        if (computeAllInfo)
            restore_path = cd();
            mkdir("Results_"+nameCatchment);
            cd("Results_"+nameCatchment);
            mkdir("Output_files");
            cd("Output_files");
            save("totpixel.mat","totpixel",'-mat');
        else
            restore_path = cd();
            mkdir("Results_"+nameCatchment+"_onlyBranching");
            cd("Results_"+nameCatchment+"_onlyBranching");
            mkdir("Output_files");
            cd("Output_files");
            save("totpixel.mat","totpixel",'-mat');
        end
    end

    % INITIALIZING VARIABLES FOR THE ANALYSIS
    tot_sources = outlet_id;
    % List of the channel sources that are currently active (only outlet at
    % the beginning, when no network is developed yet)
    moving_sources = outlet_id;
    % Other auxiliary variables related to sources
    is_moving = true;
    upstreams = [];
    distances = [];
    related_old_source = [];
    branching = false;
    % Num of pixels of the current network
    T=0;
    % Iteration counter
    k=0;
    % Map of hillslope distances from the river network
    tau = NaN(size(Aad8));
    % Vector of the lengths of all the considered networks
    network_Lengths = nan(1,geo_T);
    network_Lengths(1) = 1;
    % Vector of the branching indeces of all the considered networks
    branching_Index = nan(1,geo_T);
    branching_Index(1) = 0;
    % Other counters
    num_innerBranchings = 0;
    num_activeSources = 1;
    % Vector of the threshold areas used to extract all the considered networks
    thresholdArea_alongIterations = nan(1,geo_T);
    % Vector of the TWI for all the considered networks
    topographicWetnessIndex_Zanetti = nan(1,geo_T);


    if (computeAllInfo)
        %Computing local slopes of pixels
	    dr = [0 -1 -1 -1 0 1 1 1];  %to get downstream pixel
	    dc = [1 1 0 -1 -1 -1 0 1];
        locSlope = nan(size(Aad8));
        locSlope(outlet_id) = 0;
        warningTWI_negativeSlope = false;
        for row=1:size(myCatchment,1)
            for col=1:size(myCatchment,2)
                if (myCatchment(row,col)&&isnan(locSlope(row,col)))
                    for looo=1:8
                        if (Ap(row,col)==looo)
                            heightDiff = A(row,col) - A(row+dr(looo),col+dc(looo));
                            if (heightDiff<0)
                                warningTWI_negativeSlope = true;
                            end
                            lengthDiff = dx;
                            if (mod(looo,2)==0)
                                lengthDiff = sqrt(2)*lengthDiff;
                            end
                            locSlope(row,col) = heightDiff/lengthDiff;
                        end
                    end
                end
            end
        end
        locSlope(outlet_id) = min(locSlope,[],'all');
    
    end

    disp(" ");

end

% END OF IMPORT DATA, MORPHOLOGY INDEXES AND INITIALIZATION (the section
% started at about line 50).
%-------------------------------------------------------------------------





%% NETWORK ANALYSIS AND HILLSLOPE DISTRIBUTIONS
% This sections contains the main loop of the code over all of the spatial
% configurations of the channel network in between the minimal one (only
% outlet is channelized) and the maximal one (selected to be beyond the
% observed maximal drainage density of the catchment).
% Operations are performed automatically, the user is not required to
% modify the following code.
% The section ends at line 585.

%STARTING NETWORK ANALYSIS
disp("COMPUTING HILLSLOPE LENGTH DISTRIBUTIONS:");
disp(" ");
while ~isempty(moving_sources)
    %Updating iteration counter k, network pixels T, and report
    k = k + 1;
    T=T+numel(moving_sources);
    if (report_iterations)
        disp("Starting iteration "+k);
        disp("   Considered area threshold (px): "+Aad8(moving_sources(1)));
        disp("   Current amount of px in the network: T="+T+" out of Tmax="+geo_T);
    end
    thresholdArea_alongIterations(k) = Aad8(moving_sources(1));
    if (computeAllInfo)
        %Computing the TWI for the current network
        avgSlopeAtSources = sum(locSlope(moving_sources),"all")/numel(moving_sources);
        topographicWetnessIndex_Zanetti(k) = log(dx*dx*thresholdArea_alongIterations(k)/perimeter_length/avgSlopeAtSources);
        %Creating the corresponding network (for the given moving_sources)
        network = Aad8 > Aad8(moving_sources(1));
        %Computing the hillslope lengths (along drainage directions)
        if (k==1)
            tau = TauMap(Ap,network, outlet_id, myCatchment, tau);
        else
            tau = TauMap(Ap,network, outlet_id, sources_catchment, tau);
        end
        %Computing w0 from tau (w0 is hillslope lengths distribution
        %related to the whole catchment)
        [tauValsW0, w0Vals, normalization0] = OmegaDistribution(tau);
        if (sum(isnan(tauValsW0),"all")==numel(tauValsW0))
            disp("   WARNING: an error may've occured in the pdf computation, during iteration "+k);
        end
        vectDenAlpha = w0Vals*normalization0;
    
        %Computing the contributing catchment of the moving_sources
        sources_catchment = false(size(Aad8));
        for i=1:numel(moving_sources)
            partial_catchment = ContributingCatchment(Ap,moving_sources(i));
            sources_catchment = sources_catchment | partial_catchment;
        end
        %Computing w from tau (w is travel times distribution
        %related to just the contributing catchment of the moving sources)
        tauNaN = tau;
        tauNaN(~sources_catchment) = NaN;
        [tauValsW, wVals, normalization] = OmegaDistribution(tauNaN);
        vectNumAlpha = wVals*normalization;
        %Computing alpha as function of tau
        tauValsAlpha = tauValsW0;
        alphaVals = AlphaDistributionFast(vectNumAlpha, vectDenAlpha);
    
        if (k==1)
            distributions_support = nan(numel(vectDenAlpha),geo_T);
            distributions_alpha = nan(numel(vectDenAlpha),geo_T);
            distributions_w0 = nan(numel(vectDenAlpha),geo_T);
            distributions_w = nan(numel(vectDenAlpha),geo_T);
        end
        distributions_support(1:numel(tauValsW0),k) = tauValsW0;
        distributions_alpha(1:numel(alphaVals),k) = alphaVals;
        distributions_w0(1:numel(w0Vals),k) = w0Vals;
        distributions_w(1:numel(wVals),k) = wVals;

    end

    %Computing the next moving_sources and updating the tot_sources
    for i=1:numel(tot_sources)
        if (is_moving(i))
            [partial_upstreams, partial_dist] = find_upstreams(tot_sources(i),Ap,geo_net,true);
            upstreams = [upstreams partial_upstreams];
            distances = [distances partial_dist];
            related_old_source = [related_old_source i*ones(1,numel(partial_upstreams))];
        end
    end
    branching = false;
    is_moving = false(1,numel(tot_sources));
    maxA = max(Aad8(upstreams),[],"all");
    mask = Aad8(upstreams)==maxA;
    moving_sources = upstreams(mask);
    network_Lengths(k+1) = network_Lengths(k) + sum(distances(mask));
    moving_sources_backup = moving_sources;
    toReplace = related_old_source(mask);
    added = 0;
    num_it = numel(tot_sources);
    for j=1:num_it
        mask2 = toReplace==j;
        replacing = moving_sources(mask2);
        if (numel(replacing)>1)
            disp("   Branching occurring at source "+j+" out of "+num_it);
            branching = true;
            added = added + numel(replacing) - 1;
            num_innerBranchings = num_innerBranchings+1;
            num_activeSources = num_activeSources + numel(replacing) - 1;
        end
        if (numel(replacing)>0)
            %tot_sources(j+added) = replacing;
            tot_sources = [tot_sources(1:j+added-1) replacing tot_sources(j+added+1:end)];
            is_moving = [is_moving(1:j+added-1) true(size(replacing)) is_moving(j+added+1:end)];
            moving_sources(mask2) = [];
            toReplace(mask2) = [];
        end
    end
    mask2 = toReplace<0;
    replacing = moving_sources(mask2);
    if (numel(replacing)>0)
        disp("   Either branching(s) at inner position(s) or recovering old branch(es).");
        branching = true;
        for klm = 1:numel(replacing)
            aaaaaaaaaa = Aad8(replacing(klm));
            %Recover downstream pixel (senior network pixel) and all its
            %upstreams
	        % get downstream pixel
            senior_direction = Ap(replacing(klm));
            senior = 1;
            if (senior_direction==1)
                senior = replacing(klm) + size(Aad8,1);
            elseif (senior_direction==2)
                senior = replacing(klm) + size(Aad8,1)-1;
            elseif (senior_direction==3)
                senior = replacing(klm) - 1;
            elseif (senior_direction==4)
                senior = replacing(klm) - size(Aad8,1)-1;
            elseif (senior_direction==5)
                senior = replacing(klm) - size(Aad8,1);
            elseif (senior_direction==6)
                senior = replacing(klm) - size(Aad8,1)+1;
            elseif (senior_direction==7)
                senior = replacing(klm) + 1;
            elseif (senior_direction==8)
                senior = replacing(klm) + size(Aad8,1)+1;
            else
                disp("   ERROR in finding downstream pixel at line 259...");
            end
            [senior_ups, senior_dist] = find_upstreams(senior,Ap,geo_net,false);
            %Count number of its upstreams that have higher contributing
            %area than aaaaaaaaaaa, if the number is greater than zero then
            %it's inner branching
            AAAAA_seniorMask = Aad8(senior_ups)>aaaaaaaaaa;
            if (sum(AAAAA_seniorMask,'all')>0)   %If an inner branching actually occurs (not just jumping from a branch to another)
                num_innerBranchings = num_innerBranchings + 1;
                num_activeSources = num_activeSources + 1;
            end
        end
        tot_sources = [tot_sources replacing];
        is_moving = [is_moving true(size(replacing))];
        moving_sources(mask2) = [];
        toReplace(mask2) = [];
    end
    if ((numel(moving_sources)>0)||(numel(toReplace)>0))
        disp("   ERROR in updating sources!");
    end
    moving_sources = moving_sources_backup;
    upstreams(mask) = [];
    related_old_source = -ones(size(upstreams));
    branching_Index(k+1) = 1 - 1/(num_innerBranchings+num_activeSources);

    %Saving the results on file (maybe)
    if (save_on_file)
        if (mod(k,save_after_nIters)==0)
            last_backup_at_iter = k;
            save("Workspace_Backup.mat",'-mat');
        end
    end

end

distributions_w0 = distributions_w0/step;
distributions_w = distributions_w/step;

if (save_on_file)
    save("Workspace_Final.mat",'-mat');
    cd(restore_path);
end

% END OF NETWORK ANALYSIS AD HILLSLOPE DISTRIBUTIONS (the section
% started at line 405).
%-------------------------------------------------------------------------


% END OF MAIN PROGRAM
%-------------------------------------------------------------------------









%-------------------------------------------------------------------------
%-------------------------------------------------------------------------
%-------------------------------------------------------------------------
%AUXILIARY FUNCTIONS

function alpha = AlphaDistributionFast(num, den)

    alpha = zeros(numel(den),1);
	alpha(1:numel(num)) = num./den(1:numel(num));
	alpha(isnan(alpha)) = 0;

end


function catchment = ContributingCatchment(d8, outletId)
	% d8 = matrice con la mappa delle direzioni di drenaggio, metodo d8, i.e. 1 = est, 2 = nord-est, 3 = nord, etc.
	% outletId = indice del pixel che funge da sezione di chiusura (i.e. il pixel pi a monte nel reticolo idrografico, cio
	%			quello con l'area contribuente minore tra quelli che appartengono al reticolo

	% ritorna la matrice catchment con true se il pixel drena attraverso outletId, falso altrimenti

	catchment = false(size(d8));
	catchment(outletId) = true;

	rows = size(catchment, 1);
	cols = size(catchment, 2);

	queueR = NaN(1, 10); % queue of row indeces of pixels to parse
	queueC = NaN(1, 10); % queue of col indeces of pixels to parse
    [or, oc] = ind2sub(size(d8), outletId);
	queueR(1) = or;
	queueC(1) = oc;
	queuei = 0; % index in the queue of last parsed pixel
	queuen = 1; % how many pixels in the queue

	dir = [5 6 7 8 1 2 3 4];

	while queuei < queuen
		queuei = queuei + 1;

		% gather all pixels around
		rr = queueR(queuei) + [0 -1 -1 -1 0 1 1 1];
		cc = queueC(queuei) + [1 1 0 -1 -1 -1 0 1];
		
		% check each valid pixel around, if it drains in the current pixel add it to the queue to be parsed next
		for n = 1:numel(rr)
            rrn = rr(n);
            ccn = cc(n);
            if (rrn > 0 && ccn > 0 && rrn <= rows && ccn <= cols)
                if (d8(rrn,ccn) == dir(n))
		            catchment(rrn,ccn) = true;
			        queuen = queuen + 1;
			        queueR(queuen) = rrn;
			        queueC(queuen) = ccn;
                end
            end
		end
	end
end


function id = find_outlet(Ad8)

    maxV = max(Ad8,[],"all");
    totEl = size(Ad8,1)*size(Ad8,2);
    id = [];
    for i=1:totEl
        if (Ad8(i)==maxV)
            id = [id, i];
        end
    end
    if (numel(id)>1)
        disp("WARNING: given catchment seems to have multiple outlets.");
    end

end


function [upstreams, distances] = find_upstreams(id,p,network,only_network)

    if ((size(p,1)==size(network,1))&&(size(p,2)==size(network,2)))
        sz1=size(p,1);
        sz2=size(p,2);
        totEl = sz1*sz2;
        if ((id>0)&&(id<totEl+1))
            %Consider all the valid neighbours
            if (id==1)
                upstreams = [id+1 id+sz1+1 id+sz1];
                distances = [1 sqrt(2) 1];
            elseif (id==sz1)
                upstreams = [id-1 id+sz1-1 id+sz1];
                distances = [1 sqrt(2) 1];
            elseif (id<sz1)
                upstreams = [id-1 id+sz1-1 id+sz1 id+sz1+1 id+1];
                distances = [1 sqrt(2) 1 sqrt(2) 1];
            elseif (id==totEl)
                upstreams = [id-1 id-sz1-1 id-sz1];
                distances = [1 sqrt(2) 1];
            elseif (id==totEl-sz1+1)
                upstreams = [id+1 id-sz1+1 id-sz1];
                distances = [1 sqrt(2) 1];
            elseif (id>totEl-sz1+1)
                upstreams = [id-1 id-sz1-1 id-sz1 id-sz1+1 id+1];
                distances = [1 sqrt(2) 1 sqrt(2) 1];
            else
                ch = mod(id,sz1);
                if (ch==0)
                    upstreams = [id-sz1-1 id-sz1 id-1 id+sz1-1 id+sz1];
                    distances = [sqrt(2) 1 1 sqrt(2) 1];
                elseif (ch==1)
                    upstreams = [id-sz1+1 id-sz1 id+1 id+sz1+1 id+sz1];
                    distances = [sqrt(2) 1 1 sqrt(2) 1];
                else
                    upstreams = [id-sz1-1 id-sz1 id-sz1+1 id-1 id+1 id+sz1-1 id+sz1 id+sz1+1];
                    distances = [sqrt(2) 1 sqrt(2) 1 1 sqrt(2) 1 sqrt(2)];
                end
            end
            %Maybe exclude the ones that are not in the network
            if (only_network)
                for i=1:numel(upstreams)
                    if (network(upstreams(i))==0)
                        upstreams(i)=-1;
                    end
                end
                mask = upstreams==-1;
                upstreams(mask) = [];
                distances(mask) = [];
            end
            %Check if they're actually upstreams
            for i=1:numel(upstreams)
                if (p(upstreams(i))==1)
                    ch2=upstreams(i)+sz1;
                elseif (p(upstreams(i))==2)
                    ch2=upstreams(i)+sz1-1;
                elseif (p(upstreams(i))==3)
                    ch2=upstreams(i)-1;
                elseif (p(upstreams(i))==4)
                    ch2=upstreams(i)-sz1-1;
                elseif (p(upstreams(i))==5)
                    ch2=upstreams(i)-sz1;
                elseif (p(upstreams(i))==6)
                    ch2=upstreams(i)-sz1+1;
                elseif (p(upstreams(i))==7)
                    ch2=upstreams(i)+1;
                elseif (p(upstreams(i))==8)
                    ch2=upstreams(i)+sz1+1;
                else
                    ch2=-1;
                    disp("WARNING: found invalid entry in matrix p at find_upstreams.");
                end
                if (ch2==id)
                else
                    upstreams(i)=-1;
                end
            end
            mask = upstreams==-1;
            upstreams(mask) = [];
            distances(mask) = [];
        else
            disp("ERROR: Index id is out of range at find_upstreams.");
        end
    else
        disp("ERROR: Incompatible dimensions of inputs p and network at find_upstreams.");
    end

end


function MapPlot(ax, mapmatrix)
	hold(ax, "on");
	ax.YDir = "reverse";
	axis(ax, "equal");
	axis(ax, "off");

	colorbar(ax);

	imagesc(mapmatrix);
end


function [tau, omega, normalization] = OmegaDistribution(tauMap)

    %step = 0.5;
    global step;
	tau = step:step:max(tauMap(:))+2*step;

	omega = zeros(numel(tau), 1);

	% for t = 1:numel(tau)
	% 	omega(t) = sum(tauMap == tau(t), "all");
    % end

    normalization = 0;
    fine = size(tauMap,1)*size(tauMap,2);
    for ind = 1:fine
        t = tauMap(ind);
        if (t>0)
            insertAt = floor(t/step)+1;
		    omega(insertAt) = omega(insertAt)+1;
            normalization = normalization + 1;
        end
    end

    % normalization = sum(tauMap > 0, "all");

	omega = omega / normalization;
    
end


function tau = TauMap(d8, network, outlet_id, mask, tau)
	% d8 = matrice con la mappa delle direzioni di drenaggio, metodo d8, i.e. 1 = est, 2 = nord-est, 3 = nord, etc.
	% network = matrice binaria con 0 = versante, 1 = canale.

	% assume distanza di drenaggio = 1 (o sqrt(2)) se pixel versante, oppure 0 se canale.

	% ritorna la matrice tau con la mappa delle distanze di drenaggio

	%tau is the map of already computed distances that you may want to preserve (at least in part)
    % and mask is the binary matrix specifies which part of the tau you'd like to recompute
    % If no previous tau matrix is avaliable from previous iterations then pass tau = NaN(size(d8))

    tau(mask)=nan;

    [outlet_r, outlet_c] = ind2sub(size(tau), outlet_id);

    for r = 1:size(tau,1)
        for c = 1:size(tau,2)
            if (mask(r,c))
		        tau = calcTauPixel(tau, network, d8, r, c, outlet_r, outlet_c);
            end
	    end
	end
end

function tau = calcTauPixel(tau, network, d8, r, c, outlet_r, outlet_c)
	% return value if it was already calculated
	if ~isnan(tau(r,c))
		return;
    end

	% no tau if no drainage direction
    myd8 = d8(r,c);
	if ~(myd8 > 0 && myd8 < 9)
		tau(r,c) = NaN;
		return
	end

	% get downstream pixel
	dr = [0 -1 -1 -1 0 1 1 1];
	dc = [1 1 0 -1 -1 -1 0 1];

	rr = r + dr(myd8);
	cc = c + dc(myd8);

	% tau is null if it drains out of bounds
	if rr < 1 || cc < 1 || rr > size(tau, 1) || cc > size(tau, 2)
		if network(r,c) == true
			tau(r,c) = 0;
		else
			tau(r,c) = 1;
		end
		return;
    end

	% calculate tau of downstream pixel if necessary
    if (r==outlet_r && c==outlet_c)
        tau(rr, cc) = 0;
    else
        if isnan(tau(rr, cc))
		    tau = calcTauPixel(tau, network, d8, rr, cc, outlet_r, outlet_c);
        end
	end

	% this tau is downstream either tau + 1 or tau + sqrt(2) (if pixel is not
    % network), depending on whether myd8 is diagonal or not.
	if network(r,c) == true
		tau(r,c) = tau(rr, cc) + 0;
    else
        if (mod(myd8,2)==0)
		    tau(r,c) = tau(rr, cc) + sqrt(2);
        else
		    tau(r,c) = tau(rr, cc) + 1;
        end
	end
end