% -------------------------------------------------------------------------
% INFO:
% This code is meant to compute the mean hillslope length for different
% active network configurations for each of the four large basins
% considered in the paper (Astico, Leogra, Seugne, Umpqua). The input dtm
% of this four basins are provided in the LargeBasins.zip file which should
% be placed (and unzipped) in the same folder of this code file.
% The code perform the computations of the mean for each basin for N active network
% configurations, each configuration corresponding to a different value of total active network length.
% The N values of network length are chosen to be equally spaced in the
% interval [0; M*sqrt(A)], where M is a value that can be set by the user
% and A is the total area of the catchment. The user can set the values of
% N and M in the "Specify inputs" section here below, through the variables
% numEvaluations and maxLOverSqrtAVal respectively.
% The code provides results in the .mat format for each basin separately
% (files results_nameCatchment.mat) and for all the basins together
% (file FinalResultsMeanLargeBasins.mat). An intermediate backup is done
% before every single basin (file BackupResultsMeanLargeBasins.mat).
% Figures plotting the mean for each basin separately are also produced by
% the code and saved in pdf format.
% ALL THE RESULTS PRODUCED BY THE CODE ARE SAVED IN THE FOLDER CONTAINING
% THE INPUT DATASET, NAMED "LargeBasins", THAT CAN BE FOUND IN THE SAME
% REPOSITORY OF THIS CODE AND MUST BE PLACED IN THE SAME FOLDER OF THIS
% CODE.
% -------------------------------------------------------------------------



clear;
close all;
clc;

%Specify inputs
nameDataset = "LargeBasins";
numEvaluations = 15;
maxLOverSqrtAVal = 3;
continuePreviousAnalysis = false;

%Large basins details
indecesLargeBasins = [1 10 19 23];
numbasins = numel(indecesLargeBasins);

%Open the folder with the dataset
restore_path = cd(nameDataset);

%Initialize (or import) general variables with multiple-basins scope
lastOngoingBasin = 1;
hillslope_meanLengths = nan(numbasins,numEvaluations+1);
river_lengths = nan(numbasins,numEvaluations+1);
basins_areaKm2 = nan(numbasins,1);
if (continuePreviousAnalysis)
    load("BackupResultsMeanLargeBasins.mat","hillslope_meanLengths","river_lengths","lastOngoingBasin","basins_areaKm2");
    numEvaluations = size(river_lengths,2)-1;
    numbasins = size(river_lengths,1);
end

%Retrieve info about the basins in the dataset
catchments_list =   ["astico"; "borsoia"; "coweeta12"; "coweeta33"; "coweeta40"; "fernow14"; "fernow16"; "fernow37"; "focobon"; "leogra"; "montecalvello"; "poverty25"; "poverty35"; "turboloEst"; "turboloOvest"; "southForkPotts"; "rietholzback"; "fuciade"; "seugne"; "hubbardBrook"; "valfredda"; "cordevole"; "umpqua"; "switzerland"];
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;          30;       50];


%Loop on the remaining basins of the LSH dataset
for jj=1:numbasins
    bb = indecesLargeBasins(jj);

    %Save the ongoing basin (to eventually run a new analysis starting from this one) and save the general variables as a backup
    lastOngoingBasin = bb;
    save("BackupResultsMeanLargeBasins.mat","hillslope_meanLengths","river_lengths","lastOngoingBasin","basins_areaKm2",'-mat');
    disp("Considering basin "+jj+" out of "+numbasins);

    %Select info of the proper catchment
    nameCatchment = catchments_list(bb);
    dx = dxDTMs(bb);
    disp("  Name catchment: "+nameCatchment);
    disp("     Ongoing preliminaries...");

    %Import basin info (A,Ap,Ad8)
    [A,R] = readgeoraster(nameCatchment+"\DTM.tif");
    [Ap,Rp] = readgeoraster(nameCatchment+"\p.tif");
    [Aad8,Rad8] = readgeoraster(nameCatchment+"\ad8.tif");
    
    %Define the largest contributing catchment in the file
    outlet_id = find_outlet(Aad8);
    if (bb==23) %To remove not suitable meanders and low slopes from the Umpqua river
        outlet_id = sub2ind(size(Aad8),932,1130);
    end
    outlet_id = outlet_id(1);   %In case there're multiple basins with the same area (consider only the first one)
    myCatchment = ContributingCatchment(Ap,outlet_id);
    %Corrections
    A(~myCatchment) = nan;
    Ap(~myCatchment) = nan;
    Aad8(~myCatchment) = nan;
    totpixel = sum(myCatchment,"all");
    areaM2 = totpixel*dx*dx;
    areaKm2 = areaM2/1000/1000;

    %Compute matrix Lval to extract networks
    Lvals = computeAllNetworkLengths(Ap,Aad8,outlet_id,myCatchment,maxLOverSqrtAVal,totpixel);
    Lmax = max(Lvals,[],"all","omitnan");
    checkLmax = (Lmax/sqrt(totpixel)) <= (maxLOverSqrtAVal);
    if (checkLmax)
        disp("        ERROR: matrix Lmax is not suff complete. Program is paused.");
        pause;
    end
    Lvals = Lvals*dx;   %Lvals is in m from here on

    %Evaluate the mean hillslope length for different network lengths
    disp("     Performing analysis...");
    network_Lengths = nan(1,numEvaluations+1);
    hillslope_Means = nan(1,numEvaluations+1);
    network = Lvals<=dx;
    tau = NaN(size(Aad8));
    tau = TauMap(Ap, network, outlet_id, myCatchment, tau);
    network_Lengths(1) = dx/1000;
    hillslope_Means(1) = dx/1000 * mean(tau,"all","omitnan");
    for nn=1:numEvaluations
        Llength = nn*(maxLOverSqrtAVal/numEvaluations)*sqrt(areaM2);
        network = Lvals<Llength;
        tau = NaN(size(Aad8));
        tau = TauMap(Ap, network, outlet_id, myCatchment, tau);
        network_Lengths(nn+1) = max(Lvals(network),[],"all")/1000;
        hillslope_Means(nn+1) = dx/1000 * mean(tau,"all","omitnan");
    end

    %Save single-basin results (network_Lengths [km], hillslope_Means [km] & areaKm2 [km2])
    save("results_"+nameCatchment+".mat","hillslope_Means","network_Lengths","areaKm2",'-mat');

    %Store results in the general variables
    disp("     Storing results...");
    hillslope_meanLengths(jj,:) = hillslope_Means;  %[km]
    river_lengths(jj,:) = network_Lengths;  %[km]
    basins_areaKm2(jj) = areaKm2;  %[km2]

    clc;

end

%Save general variables with multiple-basins scope
save("FinalResultsMeanLargeBasins.mat","hillslope_meanLengths","river_lengths","basins_areaKm2",'-mat');


%Compute statistics and plot
expectedXvals = 0:(maxLOverSqrtAVal/numEvaluations):maxLOverSqrtAVal;
scalingFactors = 1./(numbasins*sqrt(basins_areaKm2));
actualXvals = scalingFactors'*river_lengths;

colors = [0 0.4470 0.7410; 0.8500 0.3250 0.0980; 0.9290 0.6940 0.1250; 0.4940 0.1840 0.5560; 0.4660 0.6740 0.1880];

for jj=1:numbasins
    bb = indecesLargeBasins(jj);
    %Retrieving catchment data
    nameCatchment = catchments_list(bb);
    dx = dxDTMs(bb);
    areaKm2 = basins_areaKm2(jj);
    netLKm = river_lengths(jj,:);
    maxX = max(netLKm,[],"all");
    xval = 0:(maxX/100):maxX;
    scaledXval = xval/sqrt(areaKm2);
    emp = hillslope_meanLengths(jj,:);
    norm = emp(1);
    horton = 0.5 * areaKm2 ./ xval;
    cf = areaKm2;
    new = sqrt(xval.^2 + cf*ones(size(xval))) - xval;
    cf2 = norm^2;
    new2 = sqrt(xval.^2 + cf2*ones(size(xval))) - xval;
    %Plotting
    fig = figure();
    hold on;
    xlabel("$L/\sqrt{A}$","Interpreter","latex");
    ylabel("$\langle\ell\rangle/\langle\ell\rangle_0$","Interpreter","latex");
    xticks([0 0.5 1 1.5 2 2.5 3]);
    xticklabels([0 0.5 1 1.5 2 2.5 3]);
    plot(actualXvals,emp/norm,'Color',colors(1,:),'LineStyle',"none",'Marker','o','MarkerSize',5,'MarkerFaceColor',[1 1 1]);
    plot(xval/sqrt(areaKm2),horton/norm,'Color',colors(4,:),'LineStyle','--');
    plot(xval/sqrt(areaKm2),new/norm,'Color',colors(3,:),'LineWidth',2.2);
    plot(xval/sqrt(areaKm2),new2/norm,'Color',colors(2,:),'LineWidth',2.2,'LineStyle',':');
    xlim([0, maxLOverSqrtAVal]);
    ylim([0, 1.5]);
    % legend(["Horton"; "Our model"; "Mean LHS (DyNET "+numbasins+" basins)"], 'Location','best');
    legend([nameCatchment+" "+round(areaKm2,1)+" km2; dx="+dx+" m"], "Location","best");
    exportgraphics(fig,"Means"+nameCatchment+"Plot.pdf","ContentType","vector","Append",false);
end

cd(restore_path);



disp("Program finished succesfully");










% AUXILIARY FUNCTIONS

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 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

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 Lvals = computeAllNetworkLengths(Ap,Aad8,outlet_id,myCatchment,maxLOverSqrtAVal,totpixel)

    % INITIALIZING OUTPUT VARIABLE
    Lvals = nan(size(Ap));
    % 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;
    % Vector of the lengths of all the considered networks
    network_Lengths = nan(1,numel(Ap));
    network_Lengths(1) = 1;
    % Vector of the branching indeces of all the considered networks
    branching_Index = nan(1,numel(Ap));
    branching_Index(1) = 0;
    % Other counters
    num_innerBranchings = 0;
    num_activeSources = 1;
    maxAAA = Aad8(outlet_id);
    printedPercentStatus = 25;
    % %To avoid infinite looping at the end
    % isnanAp = sum(isnan(Ap),"all");
    % checkEveryNIters = 400;

    while ~isempty(moving_sources)
        %Updating iteration counter k, network pixels T, and report
        k = k + 1;
        T=T+numel(moving_sources);
        currPercent = 100*(maxAAA-Aad8(moving_sources(1)))/maxAAA;
        if (currPercent>=printedPercentStatus-0.5)
            % disp("   Operation status: "+printedPercentStatus+"%");
            printedPercentStatus = 25+printedPercentStatus;
        end

        %Filling the output matrix
        Lvals(moving_sources) = network_Lengths(k);

        % if (mod(k,checkEveryNIters)==0)
        %     mynan = sum(isnan(Lvals),"all");
        %     disp("mynan: "+mynan+";    isnanAp: "+isnanAp);
        %     if (mynan<=isnanAp)
        %         break;
        %     end
        % end

        if (network_Lengths(k)/sqrt(totpixel)>maxLOverSqrtAVal)
            break;
        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,myCatchment,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)>0)     %This if statement has been swapped with the next one on 20/12/24
                %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
            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
        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,myCatchment,true);
                %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);
    
    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