% ATOMIC_conv_analysis_multiflight_BWB.m
% grids multiple selected flights on enu coordinates for analysis

%% initialize
clear
close all
printit = false;
warning('off','MATLAB:MKDIR:DirectoryExists');

% matlab script path
path_prog = '/Users/byronb/Documents/MATLAB/2020_ATOMIC';
cd(path_prog);                  % cd to matlab directory
addpath(genpath(path_prog));	% add scripts to path
rehash;

% path to SURFA netCDF files
surfa_nc_path = '/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/surfa';

% path for saving plots from this script
path_plots = '/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/plots/multiFlight';
mkdir(path_plots);

% load drops for all flights as structure 'snd_zinterp'
% with fields: RF, jd, drop_idx, circ_select, lon, lat, z, U, udir, Un, Ue, ta, rh, P, qa
%   drop_select is logical mask for all drops encircling a convective study area
%   drop_idx are drop numbers 1:N for each flight
%   RF is the flight number of each drop
%   jd is the UTC decimal day-of-year start time of each drop
load('/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/drops/raw/ATOMIC_P3_sonde_all_zinterp.mat');

tdk = 273.16;
Rgas = 287.1;
cpa = 1004.67; % air specific heat

% load rad heating profile (tdot)
dfl = '/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/rad_profiles/tdot.mat';
load(dfl,'tdot','play');

%% select profiles for analysis

% set desired flight numbers
ff = [1,3,7,8,10];  ff_str = '1_3_7_8_10';  Qr = -0.55;    % flights circling convective cells
%-- OR --
% ff = [9,11,12];     ff_str = '9_11_12';     Qr = -1.2;     % flights in convectively suppressed regions

% select good drops from all specified flights
xx = ismember(snd_zinterp.RF,ff) & snd_zinterp.circ_select;
RF = snd_zinterp.RF(:,xx);
jd = round(snd_zinterp.jd(xx));
Un = snd_zinterp.Un(:,xx);
Ue = snd_zinterp.Ue(:,xx);
P = snd_zinterp.P(:,xx);
qa = snd_zinterp.qa(:,xx);
rh = snd_zinterp.rh(:,xx);
ta = snd_zinterp.ta(:,xx);
lat = snd_zinterp.lat(xx);
lon = snd_zinterp.lon(xx);
z_ref = snd_zinterp.z;
% thasx2 = (ta + tdk).*(1010./P).^0.287; % air theta, K
thasx = ta + 0.0098*z_ref;  % air theta, C, Johnson definition
N = length(RF);

% smooth profiles for convective analysis with savitsky-golay filter
% unaxs = sgolayfilt(Un,1,25);
% ueaxs = sgolayfilt(Ue,1,25);
unaxs = Un;
ueaxs = Ue;
paxs = sgolayfilt(P,1,25);
qaxs = sgolayfilt(qa,1,25);
rhoa = paxs*100./(Rgas*(ta + tdk).*(1 + 0.61*qaxs/1e3));  % air density
latax = ones(1200,N).*lat;      % 1200xN lat array, using constant lat/lon for each drop
lonax = ones(1200,N).*lon;
zax = ones(1200,N).*z_ref/1e3;  % 1200xN alt array in km, same for each drop

%% compute latitude gradients

% bin mean profiles by latitude from -2° to 17°
[~, Yp,  ~, ~, ~] = binave2(lat,paxs,2,-2,17); Yp = Yp';
[~, Yq,  ~, ~, ~] = binave2(lat,qaxs,2,-2,17); Yq = Yq';
[~, Yun, ~, ~, ~] = binave2(lat,ueaxs,2,-2,17); Yun = Yun';
[~, Yue, ~, ~, ~] = binave2(lat,unaxs,2,-2,17); Yue = Yue';
[~, Yh,  ~, ~, ~] = binave2(lat,rh,2,-2,17); Yh = Yh';
[lat_bins, Yth, ~, ~, ~] = binave2(lat,thasx,2,-2,17); Yth = Yth';
    
% latitude gradient plots
pp = find(z_ref==20);    % pick 20m surface level
figure;
subplot(4,1,1); plot(lat_bins,Yp(pp,:),'o-','linewidth',2); grid;
ylabel('P (mb)'); title(['Mean Surface Gradients: ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')] );
subplot(4,1,2); plot(lat_bins,Yq(pp,:),'o-','linewidth',2); grid;
ylabel('q (g/kg)');
subplot(4,1,3); plot(lat_bins,Yun(pp,:),'o-',lat_bins,Yue(pp,:),'o-','linewidth',2); grid;
ylabel('U (m/s)'); legend('Un','Ue','location','best');
subplot(4,1,4); plot(lat_bins,Yth(pp,:),'o-','linewidth',2); grid;
ylabel('Theta (C)'); xlabel('Latitude');
if printit
    pname = fullfile(path_plots,['Lat_Surface_Gradients_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

%% regrid drops and compute regression coefficients

% convert lat/lon to enu grid points on N/E axes in meter units
% select axes origin for each flight as mean of lat/lon for good drops on
% that flight.
xE = NaN(size(lon)); xN = xE;
wgs84 = wgs84Ellipsoid('meters');  % specifies ellipsoid and units...
for ii = 1:length(ff) % regrid each flight to the center point for that flight
    fff = RF==ff(ii);
    lon0 = mean(lon(fff));
    lat0 = mean(lat(fff));
    [xE(fff),xN(fff),~] = geodetic2enu(lat(fff),lon(fff),0,lat0,lon0,0,wgs84);
end
    
% compute regression coefficients w/respect to xE/xN coordinates in meters
an = NaN(1200,3); ae=an; aq=an; ap=an; ath=an; % arrays for coefficients
paxm = NaN(1200,1); qaxm = paxm; unaxm = paxm; ueaxm = paxm; thaxm = paxm; rhoaxm = paxm; % mean profiles
for dd = 1:1200
    kk = isfinite(unaxs(dd,:));  % select finite Un values at level dd
    aa = [ones(size(latax(dd,kk)))',xE(kk)',xN(kk)'];
    cc = aa\unaxs(dd,kk)';  % regression coefficients for Un on xE and xN
    an(dd,:) = cc;
    kk = isfinite(ueaxs(dd,:));  % Ue
    aa = [ones(size(latax(dd,kk)))',xE(kk)',xN(kk)'];
    cc = aa\ueaxs(dd,kk)';
    ae(dd,:) = cc;
    kk = isfinite(qaxs(dd,:));  % q
    aa = [ones(size(latax(dd,kk)))',xE(kk)',xN(kk)'];
    cc = aa\qaxs(dd,kk)';
    aq(dd,:) = cc;
    kk = isfinite(thasx(dd,:)); % Theta
    aa = [ones(size(latax(dd,kk)))',xE(kk)',xN(kk)'];
    cc = aa\thasx(dd,kk)';
    ath(dd,:) = cc;
    kk = isfinite(paxs(dd,:));  % P
    aa = [ones(size(latax(dd,kk)))',xE(kk)',xN(kk)'];
    cc = aa\paxs(dd,kk)';
    ap(dd,:) = cc;

    % compute means for this level
    rhoaxm(dd) = nanmean(rhoa(dd,:));
    paxm(dd)   = nanmean(paxs(dd,:));
    qaxm(dd)   = nanmean(qaxs(dd,:));
    unaxm(dd)  = nanmean(unaxs(dd,:));
    ueaxm(dd)  = nanmean(ueaxs(dd,:));
    thaxm(dd)  = nanmean(thasx(dd,:));
end 

%% compute profiles

% Divergence and Vorticity
% dv = dUn/dN + dUe/dE, cv = dUn/dE - dUe/dN
% winds are 'from' so negate
dv = -(an(:,3) + ae(:,2));  % horizontal divergence
cl = -(an(:,2) - ae(:,3));  % vertical vorticity

% compute vertical velocity
dpr = diff(paxm); dpr = 100*[dpr(1);dpr];  % Delta press in Pa
dzr = 10;                                  % Delta z in m
wi = zeros(1200,1); wi2 = wi; dwi = wi;
for ii = 2:1200
    wi(ii) = wi(ii-1) - dv(ii)*dzr;       % Vertical velocity, 10m spacing between levels.
    wi2(ii) = wi2(ii-1) + dpr(ii)*dv(ii); % Vertical velocity, pressure integral
end
wi2 = wi2./rhoaxm/9.83;  % pressure weighted w

% taper wi to zero at 100mb level (16.2 km)
wim2 = wi2(1200);
dwi = wim2./16.2.^2.*(z_ref/1000).^2; % taper to zero at 100 mb or 16.2 km
wi2b = wi2 - dwi;

wix = wi2;  % use

% compute vertical velocity by integrating horizontal divergence
% wi = zeros(1200,1); wi2 = wi;
% for ii = 2:1200
%     % integral of divergence = vert. vel.
%     wi(ii) = wi(ii-1) - dv(ii)*10; % Vertical velocity, 10m spacing between levels.
%     wi2(ii) = wi2(ii-1) - dv(ii)*10*rhoaxm(ii); % Vertical velocity, pressure coordinate
% end
% wi2 = wi2./rhoaxm;

% smoothed mean q and theta profiles over all valid drops
qaxms = sgolayfilt(qaxm,1,55);
thaxms = sgolayfilt(thaxm,1,25);

% vertical gradient of specific humidty
dqdz = diff(qaxms)/10; dqdz = [dqdz;0];

% vertical gradient of theta
dthdz = diff(thaxms)/10; dthdz = [dthdz;0];

% vertical gradient of water vapor density
aa = 0.1; dqdz = filter(aa,[1 aa-1],dqdz);
dQdz = rhoaxm.*dqdz;

% smooth mean wind component profiles
unaxms = sgolayfilt(unaxm,1,25);
ueaxms = sgolayfilt(ueaxm,1,25);

% vertical Un, Ue derivatives
dundz = diff(unaxms)/10; dundz = [dundz;0];
duedz = diff(ueaxms)/10; duedz = [duedz;0];

% mean f*v profiles
fun = nanmean((unaxs.*sin(pi/180*latax)),2)*1.45e-4;
fue = nanmean((ueaxs.*sin(pi/180*latax)),2)*1.45e-4;

% here Q1 is actually Q1/Cp & Q2 isQ2/Le*Le/Cp, Q/Cp is deg/s units,  2.5 = Le/Cp
Q1 = (-ath(:,3).*unaxm - ath(:,2).*ueaxm + wix.*dthdz);%.*(paxm/paxm(1)).^0.287;
Q2 = -(-aq(:,3).*unaxm - aq(:,2).*ueaxm + wix.*dqdz)*2.5;

% total integrals
Q1b = cpa*nansum(rhoaxm.*dzr.*Q1);
Qr = -1 + 1.35e-3*Q1b;  % from CF analysis of Johnson, DYNAMO paper
Q2b = cpa*nansum(rhoaxm.*dzr.*Q2);
Qrb = cpa*nansum(rhoaxm.*dzr.*Qr)/24/3600;
Iq = cpa*nansum(rhoaxm.*dzr.*(Q1-Q2-Qr/24/3600)); % Q1-Q2-Qr = Hs+Hl
wib = nansum(dzr.*wix)./z_ref(1200)*100;  % mean w in cm/s
Prec = (Q1b-Qrb-10)/2500/1e3*24*3600;      % precipitation, mm/day

Ints = [Q1b;Q2b;Qrb;Iq;wib;Prec];

disp('Q1b      Q2b      Qrb      Iq       wi2      Precp');
fprintf('%06.2f   %06.2f   %06.2f   %06.2f   %06.5f  %06.2f\n', Ints);
    
%% extract surface heat fluxes from gridded SURFA fields

out = NaN(length(ff),10);
for ii = 1:length(ff)
    jj = find(RF==ff(ii));
    stjd = jd(jj(1));
    out(ii,:) = SURFA_avg(stjd,surfa_nc_path,lon,lat,jj);
end
mout = mean(out,1);

% Integral of DeltaQ, W/m2
I_MSE = zeros(1200,1);
for ii = 2:1200
    %  W/m2 = W/m2 + 10m * W s/kg/K * kg/m3 * K/s
    I_MSE(ii) = I_MSE(ii-1) + 10*cpa*rhoaxm(ii-1)*(Q1(ii-1)-Q2(ii-1)-Qr/24/3600);
end
MSE0 = mout(9) + mout(10);  % surface MSE = hs0 + hl0 from SURFA
    
% compare SURFA heat fluxes to COARE 3.5
% A=[usr tau hsb hlb hbb hsbb hlwebb tsr qsr zot zoq Cd Ch Ce  L zet dter dqer tkt Urf Trf Qrf RHrf UrfN Rnl Le rhoa UN U10 U10N Cdn_10 Chn_10 Cen_10 RF Qs Evap T10 Q10 RH10];
%     1   2   3   4   5   6    7      8   9  10  11  12 13 14 15 16   17   18   19  20  21  22  23   24   25 26  27  28  29  30    31     32     33   34 35  36  37  38   39
C35 = coare35vn(mout(1),10,mout(2),2,mout(3),2,mout(4),mout(5),mout(6),mout(7),mean(lat),2000,0,NaN,NaN);
disp(['Sensible Heat: SURFA = ',sprintf('%2.2f',mout(9)),'   COARE = ',sprintf('%2.2f',C35(3))]);
disp(['Latent Heat:   SURFA = ',sprintf('%3.2f',mout(10)),'  COARE = ',sprintf('%3.2f',C35(4))]);

Ints = [Ints;mout(9);mout(10);C35(3);C35(4)];
save(fullfile(path_plots,['Integrals_Flights_',ff_str,'.mat']),'Ints');
    
%% plots

% divergence profile
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(-(an(:,3) + ae(:,2)), paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200 1010]);
ylabel('P (mb)'); xlabel('Divergence (s^{-1})'); title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Divergence_Profile_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% curl profile
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot((ae(:,3) - an(:,2)), paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200 1010]);
ylabel('P (mb)'); xlabel('Curl (s^{-1})'); title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Curl_Profile_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% N/E q gradient profiles
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(aq(:,3),paxm,'-',aq(:,2),paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200 1010]);
ylabel('P (mb)'); xlabel('Q gradients g/kg/m'); legend(h1,'North grad','East grad','location','best');
title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Q_Gradients_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% maps of drop sonde locations
zz = isfinite(lon);
figure; plot(lon(zz),lat(zz),'bo-',lon,lat,'rx',[-180,-140],[0,0],'k--','linewidth',2); grid;
xlim([min(lon(zz))-1,max(lon(zz))+1]);
xlabel('Longitude (deg)'); ylabel('Latitude (deg)'); title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Dropsonde_Map_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end
figure; plot(xE/1000,xN/1000,'bd','linewidth',2); hold on;
plot([-1000,1000],[0,0],'k-',[0,0],[-1000,1000],'k-'); grid;
axis([min(xE/1000)-100,max(xE/1000)+100,min(xN/1000)-100,max(xN/1000)+100]);
xlabel('Kilometers East'); ylabel('Kilometers North');
title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Dropsonde_xGrid_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% W and q profiles
wi_P = -9.8.*rhoaxm.*wi;  % omega vertical velocity in Pa/s, Wi_P = -g * rohoa * w
figure;
subplot(1,2,1); plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(wi*100,paxm,'r-',wi2*100,paxm,'b-',wi2b*100,paxm,'b--','linewidth',2); grid;
%     h1 = plot(wi_P,paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200 1010]);
ylabel('P (mb)'); xlabel('Vertical Velocity (cm/s)');
legend(h1,'z-int','P-int','tapered P-int','location','best');
title('ENRR  NOAA G-4');
subplot(1,2,2); plot(qaxm,paxm,'-','linewidth',2); grid;
set(gca,'yDir','reverse'); ylim([200 1010]);
xlabel('Specific Humidity (g/kg)');
title(['Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['W_q_Profiles_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% density weighted wind components
rhoun = rhoaxm.*unaxm; rhoue = rhoaxm.*ueaxm;
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(aq(:,3).*rhoun,paxm,'-',aq(:,2).*rhoue,paxm,'-',-wi.*dQdz,paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200 1010]); xlim([-1e-4 2e-4]);
ylabel('P (mb)'); xlabel('dQ/dt  g/m^2/s');
legend(h1,'North grad','East grad','W*dq/dz'); title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['dQ_dt_Profiles_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% Q1/Q2 profiles
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(Q2*3600*24,paxm,'-',Q1*3600*24,paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200,1010]);
ylabel('P (mb)'); xlabel('Heating Rate (C/day)');
legend(h1,'Q2','Q1','location','best'); title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Q1_Q2_Profiles_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% MSE profile =  MSE0 - I_MSE
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(MSE0-I_MSE, paxm,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([200,1010]);
ylabel('P (mb)'); xlabel('Moist Static Energy Flux (W/m^2)');
title(['ENRR  NOAA G-4 Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['MSE_Flux_Profile_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

% dU/dt terms
figure;
subplot(1,2,1);
plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(-ap(:,3)./rhoaxm*1e2,paxm,'-',fue,paxm,'-',-an(:,3).*unaxm,paxm,'-',-an(:,2).*ueaxm,paxm,'-','linewidth',2);
hold off; set(gca,'yDir','reverse'); ylim([200,1010]); xlim([-2 2]*1e-4); grid;
xlabel('dUn/Dt terms'); ylabel('P (mb)');
legend(h1,'-dP/dn/\rho','-f*ue','un*Dundn','ue*Dunde','location','best');
title(['ENRR  NOAA G-4 ']);
subplot(1,2,2);
plot([0,0],[200,1200],'k-'); hold on;
h2 = plot(-ap(:,2)./rhoaxm*1e2,paxm,'-',-fun,paxm,'-',-ae(:,3).*unaxm,paxm,'-',-ae(:,2).*ueaxm,paxm,'-','linewidth',2);
hold off; set(gca,'yDir','reverse'); ylim([200,1010]); xlim([-2 2]*1e-4); grid;
xlabel('dUe/Dt terms'); ylabel('P (mb)');
legend(h2,'-dP/de/\rho','f*un','un*Duedn','ue*Dende','location','best');
title(['ENRR Flights ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['dU_dt_Terms_RF_',ff_str,'.png']);
    print(pname,'-dpng');
end

%      close all;