% ATOMIC_conv_analysis_halo_BWB_Pgrid.m
% vertical grid is pressure in 20 mb bins

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

% edit this to choose flight and run
% loads snd_p with fields: jd, lon, lat, z, U, udir, Un, Ue, ta, rh, P, qa
%   jd is the UTC decimal day-of-year start time of each drop
load('/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/halo_drops/mat/halo_20200213_all_pinterp.mat');
ff_str = '2020_02_13';

% 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
path_plots = '/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/plots/multiFlight_Pgrid_halo';
mkdir(path_plots);

tdk = 273.16;
Rgas = 287.1;
cpa = 1004.67; % air specific heat J/kg/K

% load rad heating profile (tdot)
dfl = '/Users/byronb/Raw_Data_Files.noindex/ATOMIC_2020/rad_profiles/tdot.mat';
load(dfl,'tdot','play');
tdot = flip(tdot);
play = flip(play/100);  % Pa to mb
p_bin = 250:20:1010; p_bin = p_bin'; % 38 level fixed pressure bin edges 250-1010 mb
tdot_bin = interval_avg(play,tdot,p_bin);
figure; plot(tdot,play,tdot_bin(:,2),tdot_bin(:,1)+10,'o'); grid;
set(gca,'ydir','reverse'); ylim([250,1015]);
ylabel('Pressure (mb)'), xlabel('Heating rad (°C/day)');
title('ATOMIC, Binned Heating Rate Profile');
if printit
    pname = fullfile(path_plots,['Heating_Rate_Profile.png']);
    print(pname,'-dpng');
end
tdot = tdot_bin(:,2);

%% select circles for analysis

jd = floor(snd_p.jd);
Un = snd_p.Un;
Ue = snd_p.Ue;
z = snd_p.z;
qa = snd_p.qa;
rh = snd_p.rh;
ta = snd_p.ta;
lat = snd_p.lat;
lon = snd_p.lon;
P_ref = snd_p.P;
% thasx2 = (ta + tdk).*(1010./P).^0.287; % air theta, K
thasx = ta + 0.0098*z;  % air theta, C, Johnson definition
N = length(jd);  % total number of drops to be used

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

%% compute latitude gradients
xmin = floor(min(lat)); xmax = ceil(max(lat));
% bin mean profiles by latitude from ~12° to ~15°
[~, Yq,  ~, ~, ~] = binave2(lat,qaxs,1,xmin,xmax); Yq = Yq';
[~, Yun, ~, ~, ~] = binave2(lat,unaxs,1,xmin,xmax); Yun = Yun';
[~, Yue, ~, ~, ~] = binave2(lat,ueaxs,1,xmin,xmax); Yue = Yue';
[~, Yh,  ~, ~, ~] = binave2(lat,rh,1,xmin,xmax); Yh = Yh';
[lat_bins, Yth, ~, ~, ~] = binave2(lat,thasx,1,xmin,xmax); Yth = Yth';
    
% latitude gradient plots
pp = find(P_ref==1000);    % pick surface level
figure;
subplot(3,1,1); plot(lat_bins,Yq(pp,:),'o-','linewidth',2); grid;
ylabel('q (g/kg)'); title(['Mean Surface N-S Gradients: Halo Circles ',replace(ff_str,'_',' ')] );
subplot(3,1,2); plot(lat_bins,Yun(pp,:),'o-',lat_bins,Yue(pp,:),'o-','linewidth',2); grid;
ylabel('U (m/s)'); legend('Un','Ue','location','best');
subplot(3,1,3); plot(lat_bins,Yth(pp,:),'o-','linewidth',2); grid;
ylabel('Theta (C)'); xlabel('Latitude');
if printit
    pname = fullfile(path_plots,['Lat_Surface_Gradients_HaloCircles_',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...
lon0 = nanmean1(lon);
lat0 = nanmean1(lat);
[xE,xN,~] = geodetic2enu(lat,lon,0,lat0,lon0,0,wgs84);

    
% compute regression coefficients w/respect to xE/xN coordinates in meters
an = NaN(38,3); ae=an; aq=an; az=an; ath=an; % arrays for coefficients
zaxm = NaN(38,1); qaxm = zaxm; unaxm = zaxm; ueaxm = zaxm; thaxm = zaxm; rhoaxm = zaxm; % mean profiles
for dd = 1:38
    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(z(dd,:));  % z
    aa = [ones(size(latax(dd,kk)))',xE(kk)',xN(kk)'];
    cc = aa\z(dd,kk)';
    az(dd,:) = cc;

    % compute means for this level
    rhoaxm(dd) = nanmean(rhoa(dd,:));
    qaxm(dd)   = nanmean(qaxs(dd,:));
    zaxm(dd)   = nanmean(z(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 = -100*20;                      % Delta press in Pa
dzr = diff(zaxm); dzr = [155;dzr];  % Delta z in m
wi2 = zeros(38,1); dwi = wi2;
for ii = 2:38
    wi2(ii) = wi2(ii-1) - dpr*dv(ii); % Vertical velocity, omega
end
% wi2 = wi2./rhoaxm/9.83;  % pressure weighted w m/s

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

wix = wi2;  % use omega

% vertical gradient of specific humidty, g/kg/Pa
dqdp = diff(qaxm)/dpr; dqdp = [dqdp;0];

% vertical gradient of theta deg/mb
dthdp = diff(thaxm)/dpr; dthdp = [dthdp;0];

% vertical gradient of water vapor density, g/m3/mPa
dQdp = rhoaxm.*dqdp;

% vertical Un, Ue derivatives, m/s/mPa
dundp = diff(unaxm)/dpr; dundp = [dundp;0];
duedp = diff(ueaxm)/dpr; duedp = [duedp;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 K/s units,  2.5 = Le/Cp
Q1 = (-ath(:,3).*unaxm - ath(:,2).*ueaxm + wix.*dthdp);
Q2 = -(-aq(:,3).*unaxm - aq(:,2).*ueaxm + wix.*dqdp)*2.5;

% radiation
Qr = mean(tdot);     % mean heating rate, K/day units
Qrp = tdot/24/3600;  % heating rate profile, K/s units
% Qr = -1 + 1.35e-3*Q1b;  % from CF analysis of Johnson, DYNAMO paper


%% total integrals
Q1b = -cpa/9.85*nansum(dpr*Q1);
% Qrb = -cpa/9.85*30*dpr*Qr/24/3600;
Qrb = -cpa/9.85*nansum(dpr*Qrp);
Q2b = -cpa/9.85*nansum(dpr*Q2);
Iq = -cpa/9.85*nansum(dpr*(Q1-Q2-Qrp)); % Q1-Q2-Qr = Hs+Hl
wib = -nansum(dpr*wix)/800/100/100*3600;       % mean omega in mb/hr
Prec = (Q1b-Qrb-10)/2500/1e3*24*3600;          % precipitation, mm/day

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

disp('Q1b      Q2b      Qrb      Iq       wib     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);
% MSE0 = mout(9) + mout(10);  % surface MSE = hs0 + hl0 from SURFA

MSE0 = 190; % from COARE run

% Integral of DeltaQ, W/m2
I_MSE = zeros(38,1);
for ii = 2:38
    %  W/m2 = W/m2 + kg/m/s2 * W*s/kg/K * s2/m * K/s
    I_MSE(ii) = I_MSE(ii-1) - dpr*cpa/9.85*(Q1(ii-1)-Q2(ii-1)-Qrp(ii-1));
end

    
% 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)), P_ref,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([260 1010]);
ylabel('P (mb)'); xlabel('Divergence (s^{-1})'); title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Divergence_Profile_Circles_',ff_str,'.png']);
    print(pname,'-dpng');
end

% curl profile
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot((ae(:,3) - an(:,2)), P_ref,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([260 1010]);
ylabel('P (mb)'); xlabel('Curl (s^{-1})'); title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Curl_Profile_Circles_',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),P_ref,'-',aq(:,2),P_ref,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([260 1010]);
ylabel('P (mb)'); xlabel('Q gradients g/kg/m'); legend(h1,'North grad','East grad','location','best');
title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Q_Gradients_Circles_',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(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Dropsonde_Map_Circles_',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(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Dropsonde_xGrid_Circles_',ff_str,'.png']);
    print(pname,'-dpng');
end

% W and q profiles
figure;
subplot(1,2,1); plot([0,0],[260,1010],'k-'); hold on;
h1 = plot(wi2/100*3600,P_ref,'b-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([260 1010]);
ylabel('P (mb)'); xlabel('Omega mb/hr');
title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
subplot(1,2,2); plot(qaxm,P_ref,'-','linewidth',2); grid;
set(gca,'yDir','reverse'); ylim([260 1010]);
xlabel('Specific Humidity (g/kg)');
title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['W_q_Profiles_Circles_',ff_str,'.png']);
    print(pname,'-dpng');
end

% Q1/Q2 profiles
figure; plot([0,0],[200,1200],'k-'); hold on;
h1 = plot(Q2*3600*24,P_ref,'-',Q1*3600*24,P_ref,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([260,1010]);
ylabel('P (mb)'); xlabel('Heating Rate (C/day)');
legend(h1,'Q2','Q1','location','best'); title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['Q1_Q2_Profiles_Circles_',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, P_ref,'-','linewidth',2); grid;
hold off; set(gca,'yDir','reverse'); ylim([260,1010]);
ylabel('P (mb)'); xlabel('Moist Static Energy Flux (W/m^2)');
title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['MSE_Flux_Profile_Circles_',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(fue,P_ref,'-',-an(:,3).*unaxm,P_ref,'-',-an(:,2).*ueaxm,P_ref,'-','linewidth',2);
hold off; set(gca,'yDir','reverse'); ylim([260,1010]); grid;
xlabel('dUn/Dt terms'); ylabel('P (mb)');
legend(h1,'-f*ue','un*Dundn','ue*Dunde','location','best');
title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
subplot(1,2,2);
plot([0,0],[200,1200],'k-'); hold on;
h2 = plot(-fun,P_ref,'-',-ae(:,3).*unaxm,P_ref,'-',-ae(:,2).*ueaxm,P_ref,'-','linewidth',2);
hold off; set(gca,'yDir','reverse'); ylim([260,1010]); grid;
xlabel('dUe/Dt terms'); ylabel('P (mb)');
legend(h2,'f*un','un*Duedn','ue*Dende','location','best');
title(['ATOMIC Halo Circles ',replace(ff_str,'_',' ')]);
if printit
    pname = fullfile(path_plots,['dU_dt_Terms_Circles_',ff_str,'.png']);
    print(pname,'-dpng');
end

%      close all;