%  pwp_rk4.m (PWP_UH version 1.0, March 2017)
%{ 
   From PWP_2.m (URI PWP version 1.4) (see their comments below), but with a 4th order 
   Runge-Kutta integration scheme for the time stepping

   Usage: Enter parameters in pwp_rk4_params.m file and run in the 
          Matlab command window.

   Note: the stir.m function, called by grad_mix.m was corrected to stir 
         also density, as in the original Fortran version of the program. 
         PWP_2.m V 1.4 was originally calculating density from the stirred  
         temperature and salinity, yielding density inversions in some 
         particular cases, and consequently giving negative Richardson 
         gradients which the grad_mix.m routine wasn't able to handle.

         Also fixed the application of the diffus.m function, which
         was incorrectly being called by reference.

   Fernando Santiago-Mandujano, Roger Lukas (University of Hawaii), 
   June 2016
   

   Time step with Runge-Kutta method:

	H0 is the value for the current time step. HT is the tendency of H.  DT is 
	the time step.  F(H) is the function to compute the tendencies.

	H0 = H
	KSUM = 0

	FIRST ITERATION
	HT = F(H)
	KSUM = KSUM + HT
	H = H0 + 0.5 * DT * HT

	SECOND ITERATION
	HT = F(H)
	KSUM = KSUM + 2.0 * HT
	H = H0 + 0.5 * DT * HT

	THIRD ITERATION
	HT = F(H)
	KSUM = KSUM + 2.0 * HT
	H = H0 + 1.0 * DT * HT

	FOURTH ITERATION
	HT = F(H)
	KSUM = KSUM + HT

	H = H0 + DT * KSUM / 6.0
%}

%{
  Comments from MATLAB URI PWP version 1.4 (8 Oct 2001)

  This main program is used to drive the pwp subroutine (pwpgo)
  with a daily cycle of heating, or, an arbitrary (observed) time
  series of surface flux data.

  Written by Jim Price, april 27, 1989. Please direct any
  comments or questions to Jim Price, WHOI, Woods Hole, MA
  02543, USA, tel. 508-289-2526.

  Version 2:  changed the gradient Ri relaxation method.
  Version 3:  clarified the bulk Ri relaxation.
  Version 4:  fixed an error that disabled background diffusion.

  Last editted on 20 Jan, 1999 by JFP.

  Translated for MATLAB 5.X by Peter Lazarevich and Scott Stoermer,
  GSO/URI, 11 July 2001.
%}

%  The following matlab subroutines are used:
%   pwp_rk4_params.m 
%   interp_forcing.m
%   sw_g.m
%   sw_dens0.m
%   sw_smow.m
%   sw_aonb.m
%   sw_cp.m
%   sw_alpha.m
%   sw_beta.m
%   sw_ptmp.m
%   sw_adtg.m
%   pwp_diag.m
%   pwpgo.m
%   interp_forcing.m
%   absorb.m
%   remove_si.m
%   mix5.m
%   rot.m
%   bulk_mix.m
%   grad_mix.m
%   stir.m
%   comp_tend.m
%   diffus.m

%  -- Initialize the user-defined variables. --

global f dt dz days depth dt_save lat g  rkz beta1 beta2  Method  InstabilityMatrix richard_grad

global cpw rb rg  % this is needed here (FS-M)
global kcount

% Reads file with parameters and file names
pwp_rk4_params

g = sw_g(lat,0); % Gravity (m/s^2)
f = 2*7.29E-5*sin(lat*pi/180); %  Coriolis parameter

%  -- Load the air/sea flux data. --

disp(['loading ',flx])
filein = fopen(flx,'r');
ff = textscan(filein,'%s %f %f %f %f %f %f %f\n');
fclose(filein);
tti = datenum(ff{1},'yyyymmddHHMMSS.FFF');
flux.time		= tti;     
flux.sw			= ff{2}; % Incoming short wave radiation (W/m^2)
flux.lw			= ff{3}; % Outgoing long wave radiation (W/m^2)
flux.qlat		= ff{4}; % Upward latent heat flux (W/m^2)
flux.qsens		= ff{5}; % Upward sensible heat flux (W/m^2)
flux.tx			= ff{6}; % Zonal wind stress (N/m^2)
flux.ty			= ff{7}; % Meridional wind stress (N/m^2)
flux.precip	        = ff{8}; % Precipitation (m/s)
clear ff

%  -- Interpolate the air/sea flux variables at dt resolution. --

%  Set nmet equal to the number of time increments using a resolution of
%  dt (seconds).
nmet 	= days*8.64E4/dt;
time	= flux.time(1)+(0:(nmet-1))*dt/8.64E4;

% Check record length

if time(end) > flux.time(end)
    time = time(time < flux.time(end));
    nmet = length(time);
    disp(['Met input shorter than # of days selected, truncating run to ', ...
         num2str(nmet*dt/8.64E4) ' day(s)'])
    disp('Pausing ..Hit Enter to continue')
    pause
end

% Check the time-resolution of the inertial period

if dt > 1/10*2*pi/f
    ans1 = input('Time step, dt, too large to accurately resolve the =inertial period. Is this okay? (y/n)','s');
    if ans1 == 'n'
        disp(['Please restart the program with a new dt <= ',num2str(2*pi/(10*f)),' sec'])
        return
    end
end


%------- Change made by Thibault GUIGNES on April 28th 2007 -------%
% if we have two identical values of flux.time, the interpolation
% does not work so we supress one of the identical values for the
% met data

identical=find(diff(flux.time)==0);
flux.time(identical,:)=[];
flux.sw(identical,:)=[];
flux.lw(identical,:)=[];
flux.qlat(identical,:)=[];
flux.qsens(identical,:)=[];
flux.tx(identical,:)=[];
flux.ty(identical,:)=[];
flux.precip(identical,:)=[];

qi		= interp1(flux.time,flux.sw,time,'nearest');
qo		= interp1(flux.time,flux.lw+flux.qlat+flux.qsens,time,'nearest');
tx		= interp1(flux.time,flux.tx,time,'nearest');
ty		= interp1(flux.time,flux.ty,time,'nearest');
precip	= interp1(flux.time,flux.precip,time,'nearest');

Qnet=qi-qo;

disp('Met data loading complete')

%  -- Load initial T,s profile data. --

disp(['loading ',fts])

ff = load(fts);

profile.z = ff(:,1);
profile.T = ff(:,2);
profile.s = ff(:,3);
% profile.d = sw_dens0(profile.s,profile.T);

% check for density inversions. The program won't run correctly if the
% initial temperature and salinity profiles have density inversions.
d = sw_dens0(profile.s,profile.T);
l = find(diff(d)<0);
if ~isempty(l)
   disp('The initial temperature and salinity profiles have density inversions. ')
   ans2 = input(' Do you want me to eliminate the data points causing the inversion? (y/n) ','s');
	if (ans2 == 'n')
	    disp('Please restart the program with initial temperature and salinity profiles without density inversions')
	    return
	else
          while 1
	    profile.z(l+1) = [];  % eliminate density inversions
            profile.T(l+1) = [];
            profile.s(l+1) = [];
            d = sw_dens0(profile.s,profile.T);
            l = find(diff(d)<0);
            if isempty(l)
		   break
	    end
	  end
	end
end

clear ff

%  -- Interpolate the profile variables at dz resolution. --

global nz z T s d

%  Set nz equal to the number of depth increments + 1 using a resolution of dz.

nz = 1+depth/dz;
z = ((0:nz-1)*dz)';

% Check record length

if z(end) > profile.z(end)
    z = z(z<=profile.z(end));
    nz = length(z);
    disp(['Profile input shorter than depth to run selected, truncating to ', ...
           num2str(z(end)) ' meters'])
    disp('Pausing .. Hit Enter to continue')
    pause
end

% Check the depth-resolution of the profile file

profile_increment =(profile.z(end)-profile.z(1))/(length(profile.z)-1);
if dz < profile_increment/5
    ans1 = input('Depth increment, dz, is much smaller than profile resolution. Is this okay? (y/n)','s');
    if ans1 == 'n'
        disp(['Please restart the program with a new dz >= ', ...
               num2str(profile_increment/5)])
        return
    end
end

% eliminate decreasing depths and repeated values for interpolation to work
[zz,l] = sort(profile.z,1,'ascend');
profile.z = profile.z(l);
profile.T = profile.T(l);
profile.s = profile.s(l);
l = find(diff(profile.z)==0); 
profile.z(l) = [];  % eliminate repeated depths
profile.T(l) = [];
profile.s(l) = [];


T	= interp1(profile.z,profile.T,z);
s	= interp1(profile.z,profile.s,z);
d	= sw_dens0(s,T);
T0 = T;
d0 = d;

dQsnet = 0; % time-integrated surface heat flux


global u v absrb

%  -- Load the u,v profile data 
%  -- Initialize additional profile variables at dz resolution. --
%  u and v, 	east and north current
%  absrb,			absorbtion fraction

disp(['loading ',fuv])

ff = load(fuv);

speed_profile.z = ff(:,1);
speed_profile.u = ff(:,2);
speed_profile.v = ff(:,3);

clear ff


u = interp1(speed_profile.z,speed_profile.u,z);
v = interp1(speed_profile.z,speed_profile.v,z);

absrb = absorb(beta1,beta2);

disp('Initial conditions loading complete')

%  -- Interpolate evaporation minus precipitaion at dt resolution. --

evap = (0.03456/(86400*1000))*interp1(flux.time,flux.qlat,time,'nearest');
% evap = (1/rho)*(latent heat flux/Lv); Lv = latent heat of vaporization = 2.5 *10^6 (Joules/kg)
emp = evap + precip;  % evap calculated from qlat which is positive
                          % upward

%  -- Specify a simple "background" diffusion to be applied to the profiles. --

dstab = dt*rkz/dz^2;

if dstab > 0.5
    disp('Warning, this value of rkz will be unstable')
end

%  -- Define the variables to be saved. --

pwpo.dt         = dt;
pwpo.dz         = dz;
pwpo.lat        = lat;
pwpo.z		    = z;
pwpo.ti         = [];
pwpo.t			= [];
pwpo.s			= [];
pwpo.d			= [];
pwpo.u			= [];
pwpo.v			= [];
pwpo.mld 		= [];  % mixed layer depth
pwpo.mli        = [];  % mixed layer index
pwpo.ri         = [];  % Richardson number
pwpo.rig        = [];  % Gradient Richardson number
pwpo.n2         = [];  % Buoyancy frequency squared
pwpo.dvdz2      = [];  % Shear squared
pwpo.Lo         = [];  % Buoyancy (Obukhov) length scale
pwpo.bflux      = [];  % Buoyancy flux
pwpo.tau        = [];  % Wind stress
pwpo.dPE        = [];  % Change of potential energy
pwpo.dQ         = [];  % Change in heat content in water column
pwpo.dQsnet     = [];  % Time-integrated surface heat flux
Kcount          = 0;   % Gradient Richardson no. iterations counter

%  -- Step through the PWP model. --

disp(['STATUS (out of ' int2str(nmet) ' steps):'])

InstabilityMatrix=zeros(length(z),nmet);
richard_grad(1:nmet)=0;

% start at m=2 until nmet-1 to be able to apply Runge-Kutta integration
for m = 2:nmet-1

    % Mixed layer depth calculation method chosen as a parameter
    if Method==1;
        ml_index = find(diff(d)>0.0001*dz,1,'first');% density gradient difference 
                                          % method : drho/dz > 1*10^-4 (kg.m-3).m-1
    elseif Method==2
        ml_index = 5/dz + find(d(round(5/dz):end)-d(round(5/dz))>0.002,1,'first'); 
            % Delta criteria density method 5m reference : delta = 0.01 kg.m-3
    elseif Method==3
        ml_index = find(abs(diff(T))>0.005*dz,1,'first'); % temperature gradient 
                                        % difference method : 0.02 degreeC.m-1
    elseif Method==4
        ml_index = find(-T+T(1)>0.1,1,'first'); % Delta criteria temperature method 
                                      %surface reference : delta = 0.8 degree C
    end
    if isempty(ml_index)
        error_text = ['Error in mixed layer depth calculation: the mixed layer is =too deep!', ...
             sprintf('\n') ...
            'Possible reasons:' sprintf('\n') ...
            '1) temperature inversion at surface is too strong' sprintf('\n') ...
            '2) wind is too strong' sprintf('\n') ...
            '3) evaporation is too high'];
        error([ sprintf('\n') error_text])
    end


    %  Get the depth of the surface mixed-layer.
    ml_depth = z(ml_index+1);

    % calculate diagnostic variables
    pwp_diag
 
    %  -- Store variables. --=09
    if (mod(m,dt_save) == 0 && m>1)
        pwpo.ti(:,end+1)        = time(m-1); % m-1 because it's values 
           % from the previous iteration, and RK4 is calculated for m
        pwpo.t(:,end+1)			= T;
        pwpo.s(:,end+1)			= s;
        pwpo.d(:,end+1)			= d;
        pwpo.u(:,end+1)			= u;
        pwpo.v(:,end+1)			= v;
        pwpo.mld(end+1)           = ml_depth;
        pwpo.ri(:,end+1)        = richardson_number;
        pwpo.rig(:,end+1)       = richardson_gradient_number;
        pwpo.dvdz2(:,end+1)     = Shear_squared;
        pwpo.n2(:,end+1)        = N_squared;
        pwpo.Lo(end+1)            = Lo;
        pwpo.bflux(end+1)         = Bflux;
        pwpo.tau(end+1)           = Tau;
        pwpo.dPE(end+1)         = dPE;  % Change of potential energy
        pwpo.dQ(end+1)          = dQ;  % Change in heat content in water column
        pwpo.dQsnet(end+1)      = dQsnet;  % Time-integrated surface heat flux  		
    end
        pwpo.ml_index=InstabilityMatrix ;

    %  Display running status
    if mod(m, 10) == 0
        disp([int2str(m), ' (' sprintf('%2.1f',100*m/nmet), '%)'])
    end

    
    rk = 1;   % rk=1 will use Runge-Kutta procedure, rk=0 is the original Euler 
              %      iteration (for testing purposes)
              
    if rk ~=1    % Euler iteration procedure
                    pwpgo(qi(m),qo(m),emp(m),tx(m),ty(m),m);
    else         % Runge-Kutta iteration
        
        % Apply pwp model for this time step
        % Apply four iterations corresponding to 4th order Runge_Kutta integration

        % store the original model state and initialize tendencies 
        To = T; so = s; uo = u; vo = v;
        Ta = T; sa = s; ua = u; va = v;
        Tb = T; sb = s; ub = u; vb = v;
        Tt = zeros(size(T)); st = Tt; ut = Tt; vt = Tt;

        % First iteration
    
        % Use forcing at step m
        
        pwpgo(qi(m),qo(m),emp(m),tx(m),ty(m),m);
        Ta = T; sa = s; ua = u; va = v;
        [Ta,sa,ua,va,Tb,sb,ub,vb,Tt,st,ut,vt] = comp_tend(Ta,sa,ua,va, ...
                                Tb,sb,ub,vb,Tt,st,ut,vt,To,so,uo,vo,1,0.5);
        Kcount = Kcount + kcount;

        % Second iteration
      
        % Use forcing at step (m)+dt/2
        newt = time(m) + (time(m+1)-time(m))/2;
        [qi_i,qo_i,emp_i,tx_i,ty_i] = interp_forcing(time(m),qi(m),qo(m), ...
         emp(m),tx(m),ty(m),time(m+1),qi(m+1),qo(m+1),emp(m+1),tx(m+1),ty(m+1), ...
         newt);
        T = Ta; s = sa; u = ua; v = va;
        pwpgo(qi_i,qo_i,emp_i,tx_i,ty_i,m);
        Ta = T; sa = s; ua = u; va = v;
        [Ta,sa,ua,va,Tb,sb,ub,vb,Tt,st,ut,vt] = comp_tend(Ta,sa,ua,va, ...
                                Tb,sb,ub,vb,Tt,st,ut,vt,To,so,uo,vo,2,0.5);

        Kcount = Kcount + kcount;

        % Third iteration

        % Use forcing at step (m)+dt/2
        T = Ta; s = sa; u = ua; v = va;
        pwpgo(qi_i,qo_i,emp_i,tx_i,ty_i,m);
        Ta = T; sa = s; ua = u; va = v;
        [Ta,sa,ua,va,Tb,sb,ub,vb,Tt,st,ut,vt] = comp_tend(Ta,sa,ua,va, ...
                                Tb,sb,ub,vb,Tt,st,ut,vt,To,so,uo,vo,2,1);
        Kcount = Kcount + kcount;

        % Fourth iteration

        % Use forcing at step m+1
        T = Ta; s = sa; u = ua; v = va;
        pwpgo(qi(m+1),qo(m+1),emp(m+1),tx(m+1),ty(m+1),m+1);
        Ta = T; sa = s; ua = u; va = v;
        [Ta,sa,ua,va,Tb,sb,ub,vb,Tt,st,ut,vt] = comp_tend(Ta,sa,ua,va, ...
                                Tb,sb,ub,vb,Tt,st,ut,vt,To,so,uo,vo,1,1);
                            
        T = To; s = so; u = uo; v = vo;
        
        % Apply all tendencies and get final values
        T = T + Tt./6;
        s = s + st./6;
        u = u + ut./6;
        v = v + vt./6;
    end
    
    Kcount = Kcount + kcount;
    
    %  Apply a "background" diffusion if rkz is non-zero.
    % NOTE: function diffus wasn't doing anything in the original Matlab version of
    %       the program (URI PWP version 1.4) because Matlab functions don't accept calls by  
    %       reference, the way in which the function was being called was  
    %       expecting for the second argument to be called by reference,
    %       e.g. diffus(dstab,T); and diffus(dstab,s);
    %       
    %       (FS-M, 4/23/09)

    if rkz > 0
        T = diffus(dstab,T);
        s = diffus(dstab,s);
        d = sw_dens0(s,T);
        u = diffus(dstab,u);
        v = diffus(dstab,v);
    end
    
end

% Get MLD and save array for the last iteration

% Mixed layer depth calculation method chosen as a parameter
if Method==1;
        ml_index = find(diff(d)>0.0001*dz,1,'first');% density gradient difference 
                                                 % method :1*10^-4 (kg.m-3).m-1
elseif Method==2
        ml_index = 5/dz + find(d(round(5/dz):end)-d(round(5/dz))>0.002,1,'first'); 
            % Delta criteria density method 5m reference : delta = 0.01 kg.m-3
elseif Method==3
        ml_index = find(abs(diff(T))>0.005*dz,1,'first'); % temperature gradient 
                                        % difference method : 0.02 degreeC.m-1
elseif Method==4
        ml_index = find(-T+T(1)>0.1,1,'first'); % Delta criteria temperature method 
                                      %surface reference : delta = 0.8 degree C
end
if isempty(ml_index)
        error_text = ['Error reached in mixed layer depth calculation: the mixed layer is =too deep!', ...
             sprintf('\n') ...
            'Possible reasons:' sprintf('\n') ...
            '1) tempearature inversion at surface is too strong' sprintf('\n') ...
            '2) wind is too strong' sprintf('\n') ...
            '3) evaporation is too high'];
        error([ sprintf('\n') error_text])
end


%  Get the depth of the surface mixed-layer.
ml_depth = z(ml_index+1);

% Get diagnostics
pwp_diag

    %  -- Store variables. --=09
if mod(m,dt_save) == 0
        pwpo.ti(:,end+1)        = time(nmet-1); % m-1 
                        % because RK4 is calculated for m-1
        pwpo.t(:,end+1)			= T;
        pwpo.s(:,end+1)			= s;
        pwpo.d(:,end+1)			= d;
        pwpo.u(:,end+1)			= u;
        pwpo.v(:,end+1)			= v;
        pwpo.mld(end+1)        = ml_depth;
        pwpo.ri(:,end+1)        = richardson_number;
        pwpo.rig(:,end+1)       = richardson_gradient_number;
        pwpo.dvdz2(:,end+1)     = Shear_squared;
        pwpo.n2(:,end+1)        = N_squared;
        pwpo.Lo(end+1)         = Lo;        
        pwpo.bflux(end+1)      = Bflux;
        pwpo.tau(end+1)        = Tau;
        pwpo.dPE(end+1)         = dPE;  % Change of potential energy
        pwpo.dQ(end+1)          = dQ;  % Change in heat content in water column
        pwpo.dQsnet(end+1)      = dQsnet;  % Time-integrated surface heat flux
end

pwpo.z = z;


% save pwpo structure and input data structures in fout file
eval(['save ' fout ' pwpo flux profile speed_profile dt dz days depth'])
disp(['Model Run Completed - Selected variables saved in file: ',fout,'.mat'])

pwpo


