function Output = UBCRM_core(S, Q, D50, D84, ModPhi, option, Phi, Tstar, RL)
% UBCRM_core is a function to isolate the maximum efficency channel using the
% UBCRMH model where bank strength is represented by modified friction
% angle, varying from the basic friction angle, Phi to just under 90
% degrees
%
%%%% INPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   Q       Formative discharge (cumecs)
%   S       Energy gradient (m/m)
%   D50     Characteristic grain size for sediment transport(mm)
%   D84     grain size for bed stability analysis and flow resistance (mm)
%   ModPhi  bank strength as modified friction angle 
%           (varies from Phi to 90 degrees)
%   option  flag to select how much output data you want (optional)
%   Phi     friction angle for bank sediment [degrees] (optional)
%   Tstar   critical Shield's number for D84 (optional)
%   RL      flow resistance law selected (optional)
%               1 = Ferguson 2007, 
%               2 = Strickler Manning,
%               3 = Keulegan
%
%%%% OUTPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%(included in basic output, option = 1)
%   W       solution water surface width [m]
%   d       mean hydraulic depth [m]
%   U       mean flow velocity [m]
%   QB      bed material transport rate [me/s]
%
%(added for extended output, option = 2)
%   P       wetted perimeter [m]
%   R       hydraulic radius [m]
%   Y       trapezoidal channel depth [m]
%   thetaSI side slope angle [degrees]
%   bed_str shear stress acting on the bed [Pa]
%   bank_str shear stres action on the banks [Pa]
%
%(added for verbose output, option = 3 (i.e. save the inputs, too)
%   S       as above under inputs
%   Q       as above under inputs
%   D50     as above under inputs
%   D84     as above under inputs
%   ModPhi  as above under inputs
%   Phi     as above under inputs
%   Tstar   as above under inputs 
%   RL      as above under inputs


if nargin < 6   %set the output option to default of 1 if no value selected
    option = 1;
end
if nargin < 7   %set Phi to default of 35 if no value selected
    Phi = 35;
end

if nargin < 8   %set Tstar to default of 0.025 if no value selected
    Tstar = 0.025;
end
if nargin < 9   %set RL to default of 1 if no value selected
    RL = 1;
end

%%%%%%%%%% STEP 1: CHOOSE AN INITIAL VALUE AND PERFORM INITIAL CALCS %%%%%
    Ptest = 3.0*Q^0.5;  % use hydraulic geometry equation to guess the first bed width

% calculate the local gradient in sediment transport at Ptest
    Y_minus = FindTransp(Ptest*0.99, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
    Y_plus = FindTransp(Ptest*1.01, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
    dQ1(1) = Y_plus - Y_minus;

% now move in direction of gradient towards solution
    if dQ1 > 0
        Ptest = Ptest +0.25*Ptest;
    else
        Ptest = Ptest - 0.25*Ptest;
    end

% perform a second calculation
    Y_minus = FindTransp(Ptest*0.99, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
    Y_plus = FindTransp(Ptest*1.01, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
    dQ1(2) = Y_plus - Y_minus;

%%%%%%%%%% STEP 2: SET UP A WHILE LOOP TO QUICKLY FIND THE MAXIMUM %%%%%%%
    i=2; 
    while dQ1(i-1)/dQ1(i) > 0    % continue loop until we get a reversal of the signs
        if dQ1(i) > 0
            Ptest = Ptest +0.25*Ptest;
        else
            Ptest = Ptest - 0.25*Ptest;
        end
        i = i + 1;
        Y_minus = FindTransp(Ptest*0.99, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
        Y_plus = FindTransp(Ptest*1.01, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
        dQ1(i) = Y_plus - Y_minus;
    end

%%%%%%%%%%%% STEP 3: NOW THAT WE HAVE IDENTIFIED THE BOUNDS WITHIN WHICH
%%%%%%%%%%%% THE SOLUTION LIES, WE SET THE BOUNDS AND FIND THE SOLUTION

% set the gates where P_upper is above the solution and P_lower is below it
    if dQ1(i) < 0
        P_upper = Ptest;      
        P_lower = Ptest - 0.25*Ptest;%move back one step to other side of max
    else
        P_lower = Ptest;
        P_upper = Ptest + 0.25*Ptest; %move forward one step to other side 
    end

% seek the maximum efficiency solution using Newtonian method
    Tol = 0.001;    %set the convergence toleranace, as a proportion of the solution bottom width
    while abs((P_upper - P_lower)/Ptest) > Tol
        %calculate the transport rate for a stable channel on either side of
        %the test bed width
        Y_minus = FindTransp(Ptest*0.99, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
        Y_plus = FindTransp(Ptest*1.01, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
        dQ2 = Y_plus - Y_minus;

        %MOVE THE GATES BASED ON GRADIENT dQ/dP
        if dQ2 < 0
            P_upper = Ptest;
        else
            P_lower = Ptest;
        end
        %recalculate the test bottom width
        Ptest = 0.5*(P_upper + P_lower);
    end
    
%%%%%%%%%%% STEP 4: CALCULATE THE FINAL SOLUTION %%%%%%%%%%%%%%%%%%%%%%%%
    Pbed = Ptest;           %record the solution bed width
    [QB, thetaSI] = ...  %extract the solution transport rate and geometry
        FindTransp(Pbed, S, Q, D50, D84, ModPhi, Phi, Tstar, RL);
    [SI, Y, bed_str, bank_str] = ...   %extract trap depth and stresses
        FindSS(thetaSI, Pbed, S, Q, D84, ModPhi, Phi, Tstar, RL);
    [Qcalc, W, d, U, P, R] = ... %extract channel dimensions
        FindQ(Y, Pbed, thetaSI, S, D84, RL);
    
    if option == 1
       %generate basic output
        Output = [W, d, U, QB];
    elseif option ==2
        %generate extended output including all geometric variables calculated)
        Output = [W, d, U, QB, P, R, Y, thetaSI, bed_str, bank_str];
    elseif option ==3
         %generate verbose output including shear stress calculations
        Output = [W, d, U, QB, P, R, Y, thetaSI, bed_str, bank_str,...
        S, Q, D50, D84, ModPhi, Phi, Tstar, RL];
    end
     
end

%%%%%%%%%%%%%%% END OF MAIN FUNCTION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%% SUB FUNCTIONS LISTED BELOW %%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%% SUB FUNCTION FINDTRANSP, CALLED BY MAIN %%%%%%%%%%%%%%%%%
function [QB, thetaSI] = FindTransp(Pbed, S, Q, D50, D84, ModPhi, Phi, Tstar, RL)
%FindTransp.m is a function that takes a given channel width, slope,
%discharge, bed textures, and bank strength parameters and then solves for
%the bank angle for a critically stable bank, then calculates the sediment
%transport rate associated with that geometry. The function calls FindSS.m
%to calculate the stable channel geometry, which in turn calls FindQ.m to
%ensure that all geometries considered are consistent with continuity
%
%%%% INPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   Pbed    the bottom width of the trapezoidal channel [m]
%   S       energy gradient [m/m]
%   Q       discharge carried by the channel [m3/s]
%   D50     grain size used for sediment transport calculation [mm]
%   D84     grain size used for flow resistance and stability analysis [mm]
%   ModPhi  bank strength parameter, expressed as mod. friction angle 
%           (ranges from phi to <90 degrees)
%   Phi     friction angle for bank and bed sediment
%   Tstar   Critical Shields number for D84
%   RL      choice of flow resistance laws (see FindQ.m for details)
%%%% OUTPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   QB          transport rate for specified conditions [m3/s]
%   thetaSI     angle of the channel banks at critical stability [degrees]
%   YQ          the trapezoid depth associated with specified Q [m]

%determine stable bank angle for the given bed width associated with continuity
Trange = [0.001 ModPhi - 0.001]; 
options = optimset('TolX', 10e-5);
thetaSI = fzero(@(x) FindSS(x, Pbed, S, Q, D84, ModPhi, Phi, Tstar, RL) - 1.0,...
    Trange,options);

%extract the solution shear stress values and critical shear stresses
[SI, Y, bed_str] = FindSS(thetaSI, Pbed, S, Q, D84, ModPhi, Phi, Tstar, RL);

%calculate the sediment transport rate using a sediment transport law
%specified below
QB = TranspRate(D50, bed_str, Pbed);

end

%%%%%%%%%%%%%%% SUB FUNCTION FINDSS, CALLED BY MAIN, FINDTRANSP %%%%%%%%%%%
function [SI, YQ, bed_str, bank_str, bed_crit, bank_crit] = FindSS(theta, Pbed, S, Q, D84, ModPhi, Phi, Tstar, RL)
%FindSS is a function that calculates the ratio of the critical bank
%strength and the shear stress acting onthe bank for a trapezoidal channel
%with known bed width, slope and discharge
%
%%%% INPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   theta   side slope angle [degrees]
%   Pbed    the bottom width of the trapezoidal channel [m]
%   S       energy gradient [m/m]
%   Q       discharge carried by the channel [m3/s]
%   D84     grain size used for flow resistance and stability analysis [mm]
%   ModPhi  bank strength parameter, expressed as mod. friction angle 
%           (ranges from phi to <90 degrees)
%   Phi     friction angle for bank and bed sediment
%   Tstar   Critical Shields number for D84
%   RL      choice of flow resistance laws (see FindQ.m for details)
%%%% OUTPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   SI          bank stability index (critical/applied stress
%   YQ          the trapezoid depth associated with S1 [m]
%   bed_str     shear stress acting on the bank [Pa]
%   bank_str    shear stress acting on the bed [Pa]
%   bed_crit    critical shear stress for bed entrainment [Pa]
%


%%%%%%%%%%%%%%% determine depth associated with continuity %%%%%%%%%%%%%%%%
    Yrange = [0.001*Q^0.5 3.0*Q^0.5];
    options = optimset('TolX', 10e-5);
    %this sets the range of depths within which the solution lies
    
    YQ = fzero(@(x) FindQ(x, Pbed, theta, S, D84, RL) - Q, Yrange, options);

%%%%%%%%%%%%%%% partitiion the shear stress %%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    Pbank =  2 * YQ / sin(theta*pi/180);
    W = Pbed + 2 * YQ / tan(theta*pi/180);
    totstress = 9.8 * 1000 * YQ * S;  
    SFbank = 10 ^ (-1.4026 * log10(Pbed / Pbank + 1.5) + 0.247);    
    %partioning equation
    bed_str = totstress * (1 - SFbank) * (W / (2 * Pbed) + 0.5);        
    %stress acting on the bed
    bank_str = totstress * SFbank * (W + Pbed) * sin(theta*pi/180) / (4 * YQ);  
    %stress acting on the bank
    %this is a set of  empirical equations derrived by Knight and others
    
%%%%%%%%%%%%%%% calculate the critical shear stress values %%%%%%%%%%%%%%
    bed_crit = Tstar*9.8*1650*(D84/1000);
    bank_crit = 9.8*1650*(D84/1000)*...
        Tstar*(tan(ModPhi*pi/180) / tan(Phi*pi/180))*...
        (1 - (sin(theta*pi/180)^2 / sin(ModPhi*pi/180)^2))^0.5;

%%%%%%%%%%%%%%% calculate a bank stability index %%%%%%%%%%%%%%%%%%%%%%%%
    SI = bank_crit/ bank_str;
end

%%%%%%%%%%%%%%% SUB FUNCTION FINDQ, CALLED BY MAIN, FINDSS %%%%%%%%%%%%%%%
function [Q, W, d, U, P, R] = FindQ(Y, Pbed, Theta, S, D84, RL)
%FindQ is a function to solve for the discharge associated with a set of
%input geometry parameters, a bed surface grain size, and a choice of flow
%resistance law to use, all assuming a trapezoidal channel geometry
%
%%% INPUTS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%   Y       flow depth [m]
%   P       bottom width [m]
%   Theta   angle of the channel banks [degrees]
%   S       energy gradient [m/m]
%   RL      choice of resistance law to use 
%           1 = Ferguson 2007, 
%           2 = Strickler Manning,
%           3 = Keulegan
%
%%%% OUTPUT = Q [m3/s]  %%%%%%%%%%%%%%%%%%%%%%%%
    
    %apply default flow resistance law
    if nargin < 6
        RL = 1;
    end
    
    %convert to useful units
    Theta = Theta * pi / 180;   %convert to radians
    D84 = D84 / 1000;           %convert to m
    
    %define constants
    g = 9.81;
    a1 = 6.5;   %constant for RL = 1
    a2 = 2.5;   %constant for RL = 1
    alpha = 5.5;%constant for RL = 2
    
    %calculate geometric parameters
    dW = Y / tan(Theta);    %width of the bank at the base
    dP = sqrt(Y^2 + dW^2);  %length of bank
    Area = Y*(Pbed + dW);      %cross sectional area of flow
    P = Pbed + 2*dP;       %wetted perimeter
    R = Area / P;       %hydraulic radius
    
    %choose the appropriate flow resistance law
    if RL == 1
        %Use Ferguson's continuously varying power law
        Res = a1*a2*(R/D84) / (a1^2 + a2^2*(R/D84)^(5/3))^(1/2);
    elseif RL ==2
        %Manning's relation
        Res = alpha*(R/D84)^(1/6);
    elseif RL == 3
        %Keulegan equation
        Res = (1/0.4)*log10(12.2*R/(D84));
    end
    
    %calculate velocity, U and discharge, Q
    U = Res*sqrt(g*R*S);
    Q = Area * U;
    W = Pbed + 2*dW;
    d = Area / W;
   
end

%%%%%%%%%%%%%%% SUB FUNCTION TRANSPRATE, CALLED BY FINDTRANSP %%%%%%%%%%%%
    function Qb = TranspRate(D50, tau, W)
    %function to calculate sediment transport using Parker equation

    %define constants used in the equation
    rho = 1000; 
    rho_sed = 2650; 
    g = 9.81;

    % calculate dimensionless shear stress
    Tstar = tau/(g*(rho_sed-rho)*(D50/1000));
    Tref=0.0386;
    phi=Tstar/Tref;
    
    %determine which equation to use
    if phi > 1.59
        G=5474*(1-0.853/phi)^4.5;
        %t=1
    elseif phi < 1
        G=phi^14.2;
        %t=2
    else 
        G=exp(14.2*(phi-1)-9.28*(phi-1).^2);
        %t=3
    end
    
    %Calculate the transport rate
    Qb=W*G*(0.0025*(tau/rho)^(3/2)/(g*(rho_sed/rho-1)));

    end