% MATLAB CODE MIMICING module_intialize_les.F from WRF
% the purpose of this script is to provide an easy to use interface to
% refine the model grids, topography, and sounding data without having to
% run the IDEAL.EXE

close all
clear all
scrz=get(0,'screensize');
addpath('/uufs/chpc.utah.edu/common/home/u0559900/MATLAB_SCRIPTS/'); %we call some scripts from here



%% figure toggle
fgtg=1; % set to zero to surpress figures

%% First we mimic the functionality of "read_sounding"

fname='zhong_input_sounding';
%fname='input_sounding';

% read the first line which contains surface pressure, temp, & mixing rat.
[p_surf,th_surf,qv_surf]=textread(fname,'%f %f %f',1);
%read the rest of the file formatted as hgt, theta,mix, u, v
[h_input,th_input,qv_input,u_input,v_input]=textread(fname,'%f %f %f %f %f','headerlines',1);
rho_input=zeros(size(h_input));
%get the number of input levels from the sounding
nl=length(h_input);


if fgtg==1
    
    
    figure(1)
    clf
    plot(th_input,h_input);
    
    
end

%% now replicate WRF "get_sounding" which computes the profile from the input data

%%%%%%%%%%%%%%%%%
%%% CONSTANTS %%%
%%%%%%%%%%%%%%%%%

rd=287; % dry air gas constant
cp=7*rd./2; % constant pressure
cv=cp-rd; % constant volume
cvpm=-cv./cp;
r_v=461.6; % water vapor
rvovrd=r_v./rd;
p1000mb=1000.*10^2; % MSLP
g=9.81; % gravitational acceleration m/s^2

%% Start calculations

%convert mixing ratio from g/kg to g/g
qv_input= (10^-3).*qv_input;

%convert surface P from hPa to Pa
p_surf=(10^2).*p_surf;

qvf=1+(rvovrd.*(qv_surf.*.001));
% surface density
rho_surf=1./((rd./p1000mb).*th_surf.*qvf.*((p_surf./p1000mb).^cvpm));
% surface "pi" value
pi_surf=(p_surf./p1000mb).^(rd./cp);

%% now intergrate hydrostatically up the column to produce the moist sounding

% start with the surface to the lowest model level;
qvf=1 + rvovrd.*qv_input(1);
qvf1=1 + qv_input(1); % used to adjust density for water vapor content
rho_input(1)=rho_surf;
dz=h_input(1); % note to make sure that the first sounding level is not at height zero... but somewhere above the ground

for it=1:10 % number of iterations to refine taylor expansion (hydrostatic) guess
    pm_input(1)=p_surf - .5.*dz.*(rho_surf + rho_input(1)).*g.*qvf1; % moist pressure (PM) of this level
    %now adjust rho_input within this loop based on the new pressure and the given temperature
    rho_input(1)= 1./((rd./p1000mb).*th_input(1).*qvf.*((pm_input(1)./p1000mb).^cvpm));
    %disp(pm_input);
    % now repeat to converge on a reasonable solution
end

% now integrate by moving up the column

for k=2:nl
    rho_input(k)= rho_input(k-1); %assign the density from the adjacent lower layer
    dz=h_input(k)-h_input(k-1); % use the local dz measurement;
    qvf1=.5.*(2+(qv_input(k)+qv_input(k-1)));
    qvf = 1 + rvovrd*qv_input(k);
    for it=1:10
        pm_input(k)=pm_input(k-1) - (.5.*dz.*(rho_input(k) + rho_input(k-1)).*g.*qvf1);
        rho_input(k)= 1./((rd./p1000mb).*th_input(k).*qvf.*((pm_input(k)./p1000mb).^cvpm));
    end
end


% quick test figure;

if fgtg==1
    figure(1);clf
    subplot(2,1,1)
    plot(pm_input./100,h_input,'-o','linewidth',2);
    xlabel('Pressure (hPa)');
    ylabel('Height (m)');
    set(gca,'xgrid','on','ygrid','on')
    
    subplot(2,1,2)
    plot((pm_input./rho_input')./rd,pm_input./100,'linewidth',2);
    hold on;
    xlabel('Abs. Temp');
    ylabel('Pressure (hPa)');
    %plot(th_input,pm_input,'--r');
    set(gca,'ydir','rev','xgrid','on','ygrid','on')
    title('Sounding DATA')
end
%% now compute just he dry component of the sounding starting from the top and moving down

p_input(nl)=pm_input(nl); % dry pressure at the top (P) is equal to moist pressure (PM)

for k=(nl-1):-1:1;
    dz = h_input(k+1)-h_input(k);
    p_input(k)= p_input(k+1) + .5.*(rho_input(k)+rho_input(k+1)).*g; %note no moisture correction to the density here
end

%% now assign these calculated variables to output names to pass off to the rest of WRF

zk=h_input; % height array
p=pm_input; % total pressure array
p_dry=p_input; % dry pressure
theta=th_input; % pot. temp;
rho= rho_input; % density
u= u_input;
v=v_input;
qv=qv_input;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% END OF SOUNDING PROCESSESING %
% note: these variables are still on the sounding grid, we now need to
% construct the model eta levels and move the sounding data to the model
% grid
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



%% NOW TO CREATE THE MODEL ETA LEVELS


stretch_grid=0; % set to 0 for an un stretched grid;
z_scale=-.4; % determines the curvature/stretching (not sure if we want positive or negative here....)

kde= 100; % number of model levels for the mass grid
%ide=1500; % number of x points
ide=1000; % for ZHong version
ids=1; % start index
jds=1;
jde=10; % narrow 2d slab for no\w

if stretch_grid==1
    for k=1:kde
        znw(k)=(exp(-(k-1)./(kde-1)./z_scale)-exp(-1./z_scale))./(1-exp(-1./z_scale));
    end
else %uniform grid
    for k=1:kde
        znw(k)=1-((k-1)./(kde-1));
    end
end

% for k=1:20
%     znw(k)=1-((k-1)./(kde-1));
% end
% for k=20:kde
%     kdtop=kde;
% znw(k)=znw(20).*(exp(-(k-20-1)./(kdtop-1)./z_scale)-exp(-1./z_scale))./(1-exp(-1./z_scale));
%     
% end

if fgtg==1
    figure(2);clf;
    plot(znw,'linewidth',2);
    xlabel('level #')
    ylabel('Eta Value');
    grid on;
end

%%%% now compute the vertical steps of the mass grid, its inverse and the
%%%% half levels


for k=1:kde-1
    dnw(k)=znw(k+1)-znw(k); % upper minus lower
    rdnw(k)=1./dnw(k); % inverse for use in subsequent calculations
    znu(k)=.5*(znw(k+1)+znw(k)); % layer mean for staggered u grid
end

%% grid topography

hm=100;
xa= 5;
icm=ide./2;

hgt=zeros(ide,jde);
[phb,php,ph0]=deal(zeros(ide,kde,jde));

%set up the topography and geopotential lowest level grids


% note that in matlab we can just load the matrix, but in fortran we must
% read it element by element.
%load('TERRAIN.mat');
load('ZHONG_TERRAIN.mat');
%hgt=hgt_mat; % comment this line out if not using ZHONG versions
hgt=terrain_mat;

%fortran version





%%


for j=jds:jde
    for i = ids:ide
        %hgt(i,j)=hm./(1+((i-icm)./xa).^2); %mountain defined here
        phb(i,1,j)=g.*hgt(i,j); %base state surface geopotential
        php(i,1,j)=0; %geopotential
        ph0(i,1,j)=phb(i,1,j); %initial geopotential
    end
end

%% now set up the model pressure profile


ztop=3000; % top of model domain from namelist inputs
p_top=interp_wrf(p,zk,ztop,nl); % interpolate from sounding to top model level

%{
%%%% recall %%%%%

zk=h_input; % height array
p=pm_input; % total pressure array
p_dry=p_input; % dry pressure
theta=th_input; % pot. temp;
rho= rho_input; % density
u= u_input;
v=v_input;
qv=qv_input;

%}


%%
figure(99)
plot(theta,p);
set(gca,'ydir','rev')
hold on;
tic
for j=5%jds:jde
    for i=ids:ide
        
        p_surf=interp_wrf(p,zk,squeeze(phb(i,1,j)./g),nl); % interp from sounding to surface
        mub(i,j)=p_surf-p_top; % base state mass of dry air in the column
        for k=1:kde-1
            p_level =(znu(k).*mub(i,j)) + p_top; % distribute pressure to eta levels (u-levels in this case????)
            pb(i,k,j)=p_level; % BASE STATE PRESSURE
            t_init(i,k,j)= interp_pcord_wrf(theta,p,p_level,nl)-300;% WRF subtracts 300 K from all temps
            alb(i,k,j)=(rd./p1000mb).*(t_init(i,k,j)+300).*(pb(i,k,j)./p1000mb)^cvpm; % specific density          
            
        end
        
        
        for k=2:kde
            phb(i,k,j)=phb(i,k-1,j)-dnw(k-1)*(mub(i,j).*alb(i,k-1,j)); %base state geopotential
        end
        
    end
end
toc

%%
if fgtg==1
    f=figure(50);clf
    subplot(1,2,1)
    zhgt=squeeze(phb(length(phb)./2,:,5)./g);
    dzhgt=gradient(zhgt);
    plot(zhgt,'-*');
    ylabel('height (m)');
    xlabel('grid #');
    grid on;
    subplot(1,2,2);
    plot(zhgt,dzhgt,'linewidth',2);
    xlabel('height (m)');
    ylabel('vertical spacing (m)');
    grid on;
end

%%
if fgtg==1
    f=figure(100);clf;
    %set(f,'color','w','pos',.4*scrz);
    %movegui('northeast');
    dx=50;
    xdim=ide;
    xmat=(1:xdim).*dx; xmat=repmat(xmat,[99 1]); xmat=permute(xmat,[2 1]);
    pcolor(((1:xdim).*dx),squeeze(phb(:,1:end-1,5)')./g,300+squeeze(t_init(:,:,5)'));shading flat;
    %caxis([min(t_init(:))-1 max(t_init(:))+1]);
    caxis([290 300])
    %contourf(((1:xdim).*dx),squeeze(phb(:,1:end-1,10)'),squeeze(t_init(:,:,10)'));
    %contourf(xmat',squeeze(phb(:,1:end-1,5)'),squeeze(t_init(:,:,5)'),-15:1:15);
    colorbar;
    hold on;
    plot(((1:xdim).*dx),squeeze(phb(:,1:2:end-1,5)')./g,'-w');
    ha=area((1:xdim).*dx,hgt(:,10));
    set(ha,'facecolor',[.7 .7 .7],'linewidth',2)
    set(gca,'xtick',[0:5000:(xdim.*dx)]);
    xtick=get(gca,'xtick');
    xlab=xtick./10^3;
    set(gca,'xticklabel',{xlab});
    xlabel('X Dist [km]');
    ylabel('Altitude [m]');
    %set(gca,'xgrid','on','ygrid','on','layer','top')
    %xlim([10 35]*10^3);
    %daspect([2 1 1])

    
    % now plot the aspect ratio of grid cells
    [~,dzmat]=gradient(squeeze(phb(:,1:end-1,5)')./g);
    f=figure(101);
    subplot(2,1,1)
    hold on;
    pcolor(((1:xdim).*dx),squeeze(phb(:,1:end-1,5)')./g,dzmat);shading flat;
    caxis([min(dzmat(:)) max(dzmat(:))]);
    plot(((1:xdim).*dx),squeeze(phb(:,1:2:end-1,5)')./g,'-w');
    ha=area((1:xdim).*dx,hgt(:,10));
    set(ha,'facecolor',[.7 .7 .7],'linewidth',2)
    set(gca,'xtick',[0:5000:(xdim.*dx)]);
    xtick=get(gca,'xtick');
    xlab=xtick./10^3;
    set(gca,'xticklabel',{xlab});
    xlabel('X Dist [km]');
    ylabel('Altitude [m]');
    title('Vertical Grid Spacing (dz)');
    colorbar;
    
    
    subplot(2,1,2)
    hold on;
    gaspect=dx./dzmat;
    pcolor(((1:xdim).*dx),squeeze(phb(:,1:end-1,5)')./g,gaspect);shading flat;
    title('Grid Apsect Ratio (dx/dz)')
    ha=area((1:xdim).*dx,hgt(:,10));
    set(ha,'facecolor',[.7 .7 .7],'linewidth',2)
    set(gca,'xtick',[0:5000:(xdim.*dx)]);
    xtick=get(gca,'xtick');
    xlab=xtick./10^3;
    set(gca,'xticklabel',{xlab});
    xlabel('X Dist [km]');
    ylabel('Altitude [m]');
    caxis([min(gaspect(:)) max(gaspect(:))]);
    colorbar;
end


    %%% START HERE ON THURSDAY MORNING, PICKING UP FROM module_initialze_les.F
    %%% line 360


return

