%% Process model data at observations

% Define high level paths
define_paths;

%% Load TEAMER DATA
D = dir(fullfile(dirSimData, 'Agate_*'));

T_SIM = [];
I_SIM = [];
for i = 1:length(D)
   filename = fullfile(dirSimData,D(i).name);
   % Load time-evolving data
   times = ncread(filename,'Times');
   times = times(1:19,:)'; 
   % convert to datetime
   ts = datetime(times,'InputFormat','uuuu-MM-dd''T''HH:mm:ss','TimeZone','UTC');
   T_SIM = [T_SIM; ts];
   I_SIM = [I_SIM; i*ones(size(ts))];
end   

%% LOAD OBSERVATIONAL DATA
% Load meta data
if ~exist('meta','var')   
    libFolder = fullfile(dirObservationalData, 'lib');
    load(fullfile(libFolder, 'AP2_metadata.mat'));  
    load(fullfile(libFolder, 'AgatePass_SiteBoundary.mat'));
    load(fullfile(libFolder, 'AgatePassBathy.mat'));
    tTests = meta.tTests;
    nTest = length(tTests);
    nLocSK = size(meta.locSK,1);
end

dirSensor = fullfile(dirObservationalData, 'sensors');

% Load station keeping data
if ~exist('SK','var')
    SK = importStationKeepingData(fullfile(dirSensor,'ADCP_StationKeeping','compiled'), meta);
end

% Load SWIFT data
if ~exist('SWIFT','var')
    SWIFT = importSWIFTdata(fullfile(dirSensor,'ADCP_SWIFT','compiled'), meta);
    SWIFT.TS = repmat(SWIFT.time, 1, size(SWIFT.X,2));
    SWIFT = rmfield(SWIFT, 'time');
    SWIFT.testID = repmat(SWIFT.testID, 1, size(SWIFT.X,2));
    % For better processing, round all swift data to nearest 30 seconds
    tSec = round(seconds(SWIFT.TS-meta.tTests(1)));
    SWIFT.TS = seconds(tSec)+meta.tTests(1);
end

% Load FLOAT SWARM data
if ~exist('FS','var')
    FILES = dir(fullfile(dirSensor,'FloatSwarm','compiled','Float*'));
    for i = 1:length(FILES)
    FS(i) = load(fullfile(FILES(i).folder, FILES(i).name)); 
    end
    clear FILES i;
    
     % Rename variables in FS to match SWIFT and SK name convention.
    for i = 1:length(FS)
        FS(i).X = FS(i).xSS;
        FS(i).Y = FS(i).ySS;
        FS(i).Z = FS(i).zSS;
        FS(i).UX = FS(i).UXS;
        FS(i).UY = FS(i).UYS;
        FS(i).TS = repmat((FS(i).tRef+seconds(FS(i).tTestSet))',1,size(FS(i).X,2));
    end
    FS = rmfield(FS,{'UXS','UYS','xSS','ySS','zSS'});
end

% Load CTD data
if ~exist('CTD','var')
    FILES = dir(fullfile(dirSensor,'CTD','compiled','CC*'));
    for i = 1:length(FILES)
        CTD(i) = load(fullfile(FILES(i).folder, FILES(i).name));
    end
    CTD = CTD(~[CTD.IsInvalid]);
    CTD = CTD([3:7,9:12]); % Grab relevant profiles (skips doubles, and 1st)   
    castTimes = datetime([CTD.CastTimeUtc]', 'TimeZone','UTC');
    for i = 1:length(CTD)
        meanSOS(i) = mean(CTD(i).Sound_velocity);
    end
    %sos =  interp1(castTimes, meanSOS, meta.tTests(testID)+minutes(10));
end 

%% Get unique times of all data points, rounded to 30 seconds,
for i = 1:length(FS)
    FS_TS(:,i) = seconds(FS(i).tTestSet)'+FS(i).tRef;
end
TQ = [SWIFT.time; SK.TS(:); FS_TS(:)];
secOff = mod(TQ.Second,30);
TQ.Second = TQ.Second-secOff;
TQ = unique(TQ);

[~,IA,IB] = intersect(T_SIM, TQ);
% Fill in short gaps and get times on front and back of surveys. 
bAll = false(size(T_SIM));
bAll([IA; IA-1; IA-2; IA+1; IA+2]) = 1;
dbAll = [0; diff(bAll)];
tStart = T_SIM(dbAll==1);
tEnd = T_SIM(dbAll==-1);

% Assign survey number [for each timestep T_SIM]
survey_num = cumsum([0; diff(bAll)]>0);
survey_num(~bAll) = 0;

% Index of time steps within each hour
iRel =  repmat(1:length(ts),1,length(D))';

% For SWIFT and SK, assign survey number
SWIFT.survey_num = zeros(size(SWIFT.X));
SK.survey_num = zeros(size(SK.X));
for i = 1:9
    SWIFT.survey_num(SWIFT.time>=tStart(i) & SWIFT.time <= tEnd(i),:) = i;
    SK.survey_num(SK.TS >= tStart(i) & SK.TS <= tEnd(i)) = i;
    % TBS: Confirmed. These are correctly identified. 
end

%% Convert FIELD DATA to local coordinate frame
latRef = meta.locSK(3,1);
lonRef = meta.locSK(3,2);

% Define station keeping locations (local coordinates)
[locSK(:,2), locSK(:,1)] = geodetic2ned(meta.locSK(:,1),meta.locSK(:,2),0,latRef,lonRef,0,referenceEllipsoid('wgs84'));

% Need coastlines for plotting
[yCLwest, xCLwest] =  geodetic2ned(CLwest(:,1), CLwest(:,2), 0,...
    latRef, lonRef, 0, referenceEllipsoid('wgs84'));
[yCLeast, xCLeast] =  geodetic2ned(CLeast(:,1), CLeast(:,2), 0,...
    latRef, lonRef, 0, referenceEllipsoid('wgs84'));

% Bathymetry (natively in UTM, convert to lat lon, then to local 
xB = bathyData.xB;
yB = bathyData.yB;
zB = -bathyData.zB; % Our bathymetry 
[latBs, lonBs] = utm2ll(xB, yB, 10);
[yBs, xBs] = geodetic2ned(latBs, lonBs, 0, latRef, lonRef, 0, referenceEllipsoid('wgs84'));
% Crop section we care about
xCrop = [-1000, 1500];
yCrop = [-1000, 2500];
bInside = xBs>xCrop(1) & xBs<xCrop(2) & yBs>yCrop(1) & yBs<yCrop(2);
xBs = xBs(bInside);
yBs = yBs(bInside);
zB = zB(bInside);
G_B = scatteredInterpolant(xBs, yBs, zB,'linear','none');
kBounds = boundary(xBs, yBs,1);

% Compiled Observational data is all in [m] reference to SK3. Do not
% require any adjustment here. 

%% INITIALIZE SIMULATION PROPERTIES - TIME-INVARIENT VARIABLES
filename = fullfile(dirSimData,D(1).name);

% Load constant grid parameters
x_n = ncread(filename,'x');
y_n = ncread(filename,'y');
%lat_n = ncread(filename,'lat');
%lon_n = ncread(filename,'lon');
h_n = ncread(filename,'h');
siglay_n = ncread(filename,'siglay'); % centers [1x20]
siglev_n = ncread(filename,'siglev'); % edges   [1x21]

x_c = ncread(filename,'xc');
y_c = ncread(filename,'yc');
%lat_c = ncread(filename,'latc');
%lon_c = ncread(filename,'lonc');
h_c = ncread(filename,'h_center'); % depth below geoid (MLLW?)
siglay_c = ncread(filename,'siglay_center'); % range = [-0.975, -0.025]
siglev_c = ncread(filename,'siglev_center'); % range = [-1, 0] 
nLayer = size(siglay_c,2);

% [x, y] = geodetic2ned(lat_n,lon_n,0,latRef,lonRef,0,referenceEllipsoid('wgs84'));
% Convert reference location to zonal offset. 
[xRef, yRef] = ll2utm(latRef, lonRef);

% Sim node coordinates: local
xn_l = x_n-xRef;
yn_l = y_n-yRef;
% Sim center coordinates: local
xc_l = x_c-xRef;
yc_l = y_c-yRef;


%% FOR EACH SURVEY, GET SIM DATA AND INTERPOLATE

% ASSEMBLED INDECES - SUMMARY
%      All with dimensions [1320 1]
% T_SIM - start time-step of all time-steps simulated (across all files)
% I_SIM - index of SIMULATION FILE (by hour)
% I_REL - index (out of n=120 timesteps) within the simulation hour
% survey_num - survey number (0 if timestep not used in any survey) 

% FOR EACH SURVEY, ASSEMBLE TIME-VARYING DATA AND INTERPOLATE

% Initialize array for storing interpolation from model
SWIFT.UXQ = nan(size(SWIFT.UX));
SWIFT.UYQ = nan(size(SWIFT.UY));
SK.UXQ = nan(size(SK.UX));
SK.UYQ = nan(size(SK.UY));
for i = 1:length(FS)
    FS(i).UXQ = nan(size(FS(i).UXS));
    FS(i).UYQ = nan(size(FS(i).UYS));
end

% Initialize Flags for Identifying points to interpolate
T_SWIFT = repmat(SWIFT.time,1,size(SWIFT.X,2));
bQ_SWIFT = false(size(T_SWIFT));
bV_SWIFT = ~isnan(SWIFT.X) & ~isnan(SWIFT.Y) & ~isnan(SWIFT.Z); % all valid
T_SK = SK.TS;
bQ_SK = false(size(SK.TS));
bV_SK = ~isnan(SK.X) & ~isnan(SK.Y) & ~isnan(SK.Z); % all valid

for iSurvey = 1:9
    fprintf('>>>    Processing Survey %i   <<<\n ', iSurvey);
    cSurvey = tic;
    % Initialize the data structure
    bSurvey = survey_num == iSurvey;
    ts = T_SIM(bSurvey);
    TS = repmat(permute(ts,[3,2,1]),size(xc_l,1),nLayer,1);
    XS = repmat(xc_l,1,nLayer,length(ts));
    YS = repmat(yc_l,1,nLayer,length(ts));
    ZS = 0*YS; % initialize, changes with each time step. 
    UXS = 0*YS; % initialize
    UYS = 0*YS; % initialize
    
    % Load time-varying data
    fprintf(' >> Loading simulation data ...');
    timer_start = tic;
    sim = I_SIM(bSurvey);
    iSim = unique(sim); % For stitching surveys over multiple sim_hour
    for i = 1:length(iSim)
       filename = fullfile(dirSimData,D(iSim(i)).name); % sim file
       bOut = sim == iSim(i); % logical index for placing data subset
       iSub = iRel(bSurvey & I_SIM==iSim(i)); % index for pulling data subset
       % Load time-evolving data
       u = ncread(filename,'u'); % @ centers
       v = ncread(filename,'v'); % @ centers
       UXS(:,:,bOut) = u(:,:,iSub);
       UYS(:,:,bOut) = v(:,:,iSub); 
       % ww = ncread(filename,'ww'); % upward water velocity. 
       
       % Interpolate zeta to center points, only for this time step
         % zeta is water level above geoid (MLLW) at nodes.
           % dimensions [node, time]
       zeta = ncread(filename,'zeta');   
       t_sim = T_SIM(I_SIM==iSim(i));
       X_L = repmat(xn_l,1,size(zeta,2));
       Y_L = repmat(yn_l,1,size(zeta,2));
       T_L = repmat(t_sim',length(xn_l),1);
       tRef = T_L(1);
       T_L = seconds(T_L-tRef);
       G_zeta = scatteredInterpolant(double(X_L(:)),double(Y_L(:)),T_L(:),double(zeta(:)), 'linear','none');
       XC_L = repmat(xc_l,1,length(t_sim)); % Create query set at centers
       YC_L = repmat(yc_l,1,length(t_sim));
       TC_L = repmat(t_sim',length(xc_l),1);
       TC_L = seconds(TC_L-tRef);
       
       % Compute zeta at centers
       zeta_c = G_zeta(double(XC_L(:)), double(YC_L(:)), double(TC_L(:)));
       zeta_c = reshape(zeta_c,size(XC_L));

       %d_n = permute(zeta+h_n,[1,3,2]); % time-varying water depth
       %z_n = -d_n.*repmat(siglay_n,1,1,size(d_n,3)); % time-varying nodal depth (positive down, surface relative)
       
       % water depth at centers.
       d_c = permute(zeta_c+h_c,[1,3,2]); % dimensions [centers, 1, time]
       Z_C = -d_c.*repmat(siglay_c,1,1,size(d_c,3));
       
       % Grab subset (in time) associated with this survey
       ZS(:,:,bOut) = Z_C(:,:,iSub);
    end
    
       
    % Clear some memory before proceeding
    clear zeta_c Z_C XC_L YC_L TC_L X_L Y_L G_zeta u v zeta 
    c1 = toc(timer_start);
    fprintf(' completed. Elapsed time: %0.1f sec. (once per survey) \n', c1);
    
    % Copy top layer to zero depth
    XS = cat(2,XS(:,1,:), XS);
    YS = cat(2,YS(:,1,:), YS);
    ZS = cat(2,-0.01+0*ZS(:,1,:), ZS);
    UXS = cat(2, UXS(:,1,:), UXS);
    UYS = cat(2, UYS(:,1,:), UYS);
    
    % Convert data to double: this is data covering entire survey. 
    XS = double(XS);
    YS = double(YS);
    ZS = double(ZS);
    UXS = double(UXS);
    UYS = double(UYS);
    
    fprintf(' >> Initializing interpolation functions... ');
    timer_start = tic;
    % Initialize interpolation functions 
    X = XS(:,:,1);
    Y = YS(:,:,1);
    Z = ZS(:,:,1);
    U = UXS(:,:,1);
     % Step 1 - UX interp function
    GUX1 = scatteredInterpolant(X(:),Y(:),Z(:),U(:),'linear','none'); 
     % Step 1 - UY interp function
    GUY1 = GUX1;
    GUY1.Values = reshape(UYS(:,:,1),[],1);
     % Step 2 - UX interp function
    GUX2 = GUX1; 
    GUX2.Points(:,3) = reshape(ZS(:,:,2),[],1); % Update Z coord
    GUX2.Values = reshape(UXS(:,:,2),[],1);
     % Step 2 - UY interp function
    GUY2 = GUX2; 
    GUY2.Values = reshape(UYS(:,:,2),[],1);
   
    % Internal test
    if 1 == 1
        xq = locSK(3,1); 
        yq = locSK(3,2);
        zq = 0:1:20;
        [XQ, YQ, ZQ] = ndgrid(xq, yq, zq);
        ux = GUX1(XQ(:), YQ(:), ZQ(:));
        uy = GUY1(XQ(:), YQ(:), ZQ(:));
        
        % Will interpolate from deeper areas, remove those data
        G_D = scatteredInterpolant(XS(:,1,1), YS(:,1,1), double(d_c(:,1,1)));
        dq = G_D(xq,yq);
        bValid = zq < dq; 
        ux(~bValid) = nan;
        uy(~bValid) = nan;
        
        uh = sqrt(ux.^2 + uy.^2);
        plot(uh, zq, ux, zq, uy, zq);
        legend('|U|','Ux', 'Uy');
        
        ux = GUX2(XQ(:), YQ(:), ZQ(:));
        uy = GUY2(XQ(:), YQ(:), ZQ(:));
        % technically, the second time step of depth, but doesn't change
        % enough to matter
        ux(~bValid) = nan;
        uy(~bValid) = nan;
        
        uh = sqrt(ux.^2 + uy.^2); hold on;
        plot(uh, zq, ux, zq, uy, zq);
        legend('|U|','Ux', 'Uy');
        % layers are correctly oriented. 
    end
    
    % Initiate the float times [same size as all observations in survey]
    T_FS = repmat(seconds(FS(iSurvey).tTestSet)'+FS(iSurvey).tRef,...
                    1, size(FS(iSurvey).xSS,2));
                
    % Initiate query flag            
    bQ_FS = false(size(T_FS));
    % Initiate valid float positions
    bV_FS = ~isnan(FS(iSurvey).xSS) & ~isnan(FS(iSurvey).ySS) & ...
            ~isnan(FS(iSurvey).ySS);
        
    c1 = toc(timer_start);
    fprintf(' completed. Elapsed time: %0.1f sec. \n', c1);
    
    %% Assemble query points
    for i = 1:length(ts)-1
        timer_start = tic;
        %if mod(i,5)==0
        fprintf('  > Interpolating data at timestep %i of %i... \n',i, length(ts)-1);
        % end
        
        % Find time points within range data
        t1 = ts(i);
        t2 = ts(i+1);
        % Reset query mask
        bQ_FS(:) = 0; bQ_SWIFT(:) = 0; bQ_SK(:) = 0;
        % Find points within time and valid
        bQ_SWIFT(T_SWIFT >= t1 & T_SWIFT < t2 & bV_SWIFT) = 1;
        bQ_SK(T_SK >= t1 & T_SK < t2 & bV_SK) = 1;
        bQ_FS(T_FS >= t1 & T_FS < t2 & bV_FS) = 1;
        
        if any(bQ_SWIFT, 'all') || any(bQ_SK,'all') || any(bQ_FS,'all')
            % Get corresponding query points

            % Get swift points
            xq_s = SWIFT.X(bQ_SWIFT);
            yq_s = SWIFT.Y(bQ_SWIFT);
            zq_s = SWIFT.Z(bQ_SWIFT);
            tq_s = T_SWIFT(bQ_SWIFT);
            n_s = 1+zeros(size(xq_s));
            
            % Get station keeping points
            xq_sk = SK.X(bQ_SK);
            yq_sk = SK.Y(bQ_SK);
            zq_sk = SK.Z(bQ_SK);
            tq_sk = SK.TS(bQ_SK);
            n_sk = 2+zeros(size(xq_sk));
            
            % Get float data points
            xq_f = FS(iSurvey).xSS(bQ_FS);
            yq_f = FS(iSurvey).ySS(bQ_FS);
            zq_f = FS(iSurvey).zSS(bQ_FS);
            tq_f = T_FS(bQ_FS);            
            n_f = 3+zeros(size(xq_f));
            
            % Concatenate all query points
            xq = [xq_s(:); xq_sk(:); xq_f(:)];
            yq = [yq_s(:); yq_sk(:); yq_f(:)];
            zq = [zq_s(:); zq_sk(:); zq_f(:)]; 
            tq = [tq_s(:); tq_sk(:); tq_f(:)];
            nq = [n_s(:); n_sk(:); n_f(:)]; % for reassigning points at end
            
            % run interpolations
            % Takes about 0.1 second to evaluate. Most of the time is spent
            % updating the interpolation function (~5 sec),
            % so not breaking out timing. 
            % fprintf('  > Running intepolations ...'); 
            uxq1 = GUX1(xq, yq, zq);
            uxq2 = GUX2(xq, yq, zq);
            uyq1 = GUY1(xq, yq, zq);
            uyq2 = GUY2(xq, yq, zq);
            
            % compute time weighting
            t_rel = seconds(tq-t1)./seconds(t2-t1);
            uxq = uxq1+t_rel.*(uxq2-uxq1);
            uyq = uyq1+t_rel.*(uyq2-uyq1);
            
            % assign to outputs
            if any(bQ_SWIFT,'all')
                SWIFT.UXQ(bQ_SWIFT) = uxq(nq==1);
                SWIFT.UYQ(bQ_SWIFT) = uyq(nq==1);
            end
            
            if any(bQ_SK,'all')
                SK.UXQ(bQ_SK) = uxq(nq==2);
                SK.UYQ(bQ_SK) = uyq(nq==2);
            end
            
            if any(bQ_FS,'all')
                FS(iSurvey).UXQ(bQ_FS) = uxq(nq==3);
                FS(iSurvey).UYQ(bQ_FS) = uyq(nq==3);
            end
        end
        
        
        % Update interpolation functions
        if i<length(ts)-1
            % fprintf('  > Updating interpolation functions...');
           
            % Step 2 now assigned to Step 1
            GUX1 = GUX2;
            GUY1 = GUY2;
            % Update Step 2
            % Note that since we've started on index i = 1, the next+1
            % index is forward 2 time steps from the current index. 
            GUX2.Points(:,3) = reshape(ZS(:,:,i+2),[],1); % Update Z coord
            GUX2.Values = reshape(UXS(:,:,i+2),[],1);
            % Step 2 - UY interp function
            GUY2 = GUX2;
            GUY2.Values = reshape(UYS(:,:,i+2),[],1);    
        end
        c1 = toc(timer_start);
        fprintf('    completed (%0.1f sec) \n',c1);
        
    end
    c2 = toc(cSurvey);
    fprintf(' >> SURVEY %i COMPLETE. Total elapses time %0.1f min. \n', iSurvey, c2/60);
        
end % End loop of surveys    


%% Save the data
swift.UXQ = SWIFT.UXQ;
swift.UYQ = SWIFT.UYQ;
sk.UXQ = SK.UXQ;
sk.UYQ = SK.UYQ;

for i=1:length(FS)
    fs(i).UXQ = FS(i).UXQ;
    fs(i).UYQ = FS(i).UYQ;
end

if 
save('sim_estimates_2','swift','sk','fs');        
        

%% Check simulation in same coordinates as field data
f1 = figure(1); clf(f1); f1.Color = 'w';
ax1 = gca;
set(ax1,'NextPlot','add');
f1.Units = 'inches'; 
f1.Position = [1,1,6.5,7];

xOffset = 0; 
yOffset = 0;

ax = ax1;
pLandE = patch(ax, xCLeast-xOffset, yCLeast-yOffset, 0.9*[1,1,1],'EdgeColor','none');
pLandW = patch(ax, xCLwest-xOffset, yCLwest-yOffset, 0.9*[1,1,1],'EdgeColor','none');
  

set(ax1, 'XLim', [-1000, 1500], 'YLim', [-1250, 2000]);
%set(ax1(3:4), 'XLim', [-500, 1250], 'YLim', [-1250, 1000]);
set(ax1,'XTick', -1000:500:2000, 'YTick', -1000:500:2000,...
    'XGrid','on', 'YGrid', 'on',...
    'DataAspectRatio',[1,1,1]);

p1 = plot(ax1, xc_l, yc_l, 'r.'); hold on;
p2 = plot(ax1, xn_l, yn_l, 'g.');
title(ax1, 'Model coordinates shifted to observation reference frame (SK3)');
legend(ax1,[p1,p2], {'Center Points','Node Points'},'Location','NorthWest');
xlabel(ax1, 'Easting [m]');
ylabel(ax1, 'Northing [m]');
%print(f1, fullfile(dirFig,'model_coord_2_local'),'-djpeg','-r600');





