%make an index for the location of the surface westerly maximum

clear all; close all;

load('F:\trace\PS\PS_dec.mat','PS', 'lat', 'lon')
load('F:\trace\U\U_dec.mat', 'U');
load('F:\trace\LH_dec.mat', 'YEARS')
[NT, NN,N]=size(U);
P_THRESH=920*100;
LW=2;


  wid=25;
  mid=42;
  mn_lat=mid-wid;
  mx_lat=mid+wid;
  flat2=mn_lat:.1:mx_lat;
  ymn=0; ymx=10;

   for nt=1:1:NT;
       nt
        UU =double(squeeze(U(nt, :,:)));
        PSPS = double(squeeze(PS(nt, :,:)));
        MASK=ones(NN,N);
        MASK(find(PSPS<P_THRESH)) = NaN;
        ZU = nanmean((MASK.*UU)');
        
%         figure(1);clf;
%         hold on; plot(lat, ZU, 'r', 'LineWidth', LW)
   
   
       
       U2= interp1(lat, ZU, flat2);
       
       %search for zone of westerlies
       don = find ( U2 >= 0 & flat2 >=25 & flat2 <=65 ); %indexes of NH mid-latitude Westerlies
       N_mx_lat = max(flat2(don));
       N_mn_lat = min(flat2(don));
       

%        hold on; grid on;
%        hold on; plot(N_mx_lat*[1 1], [ ymn ymx], 'm--', 'LineWidth', LW)
%        hold on; plot(N_mn_lat*[1 1], [ ymn ymx], 'm--', 'LineWidth', LW)
       %interpolate to finer grid in the NH and then find the first mode
       U_max(nt,1) = mean(U2(don).*flat2(don))/mean(U2(don));
       
%       hold on; plot(U_max(nt,1)*[1 1], [ ymn ymx], 'm-', 'LineWidth', LW)
        
       %repeat in SH
       U2= interp1(lat, ZU, -flat2);
       dos = find ( U2 >= 0 & -flat2 <= -25 & -flat2 >=-62 ); %indexes of NH mid-latitude Westerlies
       S_mx_lat = max(-flat2(dos));
       S_mn_lat = min(-flat2(dos));
       
%         figure(1); 
%        hold on; plot(S_mx_lat*[1 1], [ ymn ymx], 'g--', 'LineWidth', LW)
%        hold on; plot(S_mn_lat*[1 1], [ ymn ymx], 'g--', 'LineWidth', LW)
       %interpolate to finer grid in the NH and then find the first mode
       U_max(nt,2) = mean(U2(dos).*-flat2(dos))/mean(U2(dos));
%        hold on; plot(U_max(nt,2)*[1 1], [ ymn ymx], 'g-', 'LineWidth', LW)
       warning last off
   end

figure(2);clf;
hold on; plot(YEARS, U_max(:,1), 'r', 'LineWidth', LW)
hold on; plot(YEARS, -U_max(:,2), 'b', 'LineWidth', LW)

save('F:\trace\scripts\jet_location.mat','U_max', 'YEARS')