%% Script to assimilate the Argo probe velocity data onto a regular grid
% Uses a zeroth order Kalman filter to match the interpolation back
% onto the probe data (Molcard et al 2003)

% Before running this script the Argo data must be downloaded from 
% http://www.coriolis.eu.org/cdc/data_selection.htm
% (this site can be very slow), and then imported using the script
% "readArgoData.m" which will create a file called probeData.mat.
% The probeData.mat file used in the paper can be found in the
% Supplementary Files
%
% These data are then sub-sampled, and this sample is used to make one 
% reconstruction of the ocean currents. This sub-sampling is repeated many
% times (typically a few hundred) to obtain many estimates of the deep-sea
% ocean current velocity field.

% These velocity fields are then used in the script "trackParticles.m" to 
% drive the particle tracking model

% Jon Yearsley & Julia Sigwart (Aug 2010)
% jon.yearsley@ucd.ie

% Molcard et al, 2003,
% "Assimilation of drifter observations for the reconstruction of the Eulerian circulation field",
% Journal of Geophysical Research, Vol 108, No C3, 3056
% doi:10.1029/2001JC001240
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear all

nFiles = 1; % Number of current vector fields to generate (each saved in a separate file)
outputDir = '.'; % Directory where output files will be saved
filePref = 'Shallow_800m_1400m_Data'; % File prefix for output files

% Define the long lat grid of interest
gridStep = 0.1; % Size of each grid square
latVec = -26:gridStep:-2;
longVec = 143:gridStep:175;
% Define month range
monthRange = [1,12];
% Define depth range (in dBar~m)
% 800dBar is a cut off for "shallow", little data less than this
% 1300dBar is a sensible cut off between deep, and "shallow"
% 2599dBar is deeper than vast majority of Argo data, and chiton data
depthRange = [800, 1400];
yearRange = [2000, 2009];

% Error on positions (in degrees) of argo data and model
% (only the ratio matters)
errorSD_argo = 1;
errorSD_model = 10;

nIter = 3; % Number of iterations for the interpolation
subSampStep = 4; % Take data from one probe cycle every subSampStep grid squares (for a 0.1 grid we sample every 4th grid square)
distWeight = 1; % The smaller this is the more 'local' the corrections to the velocity field
options = odeset('AbsTol', 1e-4); % ODE solver soptions
displayOn = false;  % If true plot figures

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% All parameters are specified above this line.
% Do not edit anything below this line without first understanding the code
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


for f = 1:nFiles    % Each vector field estimate for the deep sea ocean currents is saved in a separate file
    tic % Start clock to time execution

    %% Load data
    load('probeData.mat') % Load Argo probe data (created with readArgoData.m)
    load('coastline.mat') % Load coastline data
    
    % Select coastline within our specified region
    coastSelect = (longlatCoast(:,1)>=min(longVec)-gridStep & longlatCoast(:,1)<=max(longVec)+gridStep & ...
        longlatCoast(:,2)>=min(latVec)-gridStep & longlatCoast(:,2)<=max(latVec)+gridStep) | isnan(longlatCoast(:,1)) ;
    
    % Select probe data within specified region, depth and date range
    select = depth >=depthRange(1) & depth<=depthRange(2) & ...
        month>=monthRange(1) & month<=monthRange(2) & ...
        year>=yearRange(1) & year<=yearRange(2);
    select = select & longlatStart(:,1)>=min(longVec)-gridStep & longlatStart(:,1)<=max(longVec)+gridStep & ...
        longlatStart(:,2)>=min(latVec)-gridStep & longlatStart(:,2)<=max(latVec)+gridStep;
    select = select & longlatEnd(:,1)>=min(longVec)-gridStep & longlatEnd(:,1)<=max(longVec)+gridStep & ...
        longlatEnd(:,2)>=min(latVec)-gridStep & longlatEnd(:,2)<=max(latVec)+gridStep;
    
    nProbe = nnz(select); % Number of selected probe cycles
    % Selected probe cycle start and finish points, and the time taken to travel between these points
    xStart = longlatStart(select,:); 
    xEnd = longlatEnd(select,:);
    dt = timeStep(select);
    
    xData = (longlatStart(select,:) + longlatEnd(select,:))/2;  % Mid-point of each probe cycle
    v1Data = (longlatEnd(select,1) - longlatStart(select,1))./dt; % Rate of change of longitude for each probe cycle
    v2Data = (longlatEnd(select,2) - longlatStart(select,2))./dt;% Rate of change of latitude for each probe cycle
    
    % Month, year and depth of each selected probe cycle
    month = month(select);
    year = year(select);
    depth = depth(select);
    
    %% Randomly subsample the probe data to thin it out
    % Select one cycle every subSampleStep grid square
    keep = false(nProbe,1);
    offseti = randi(subSampStep);  % Randomise starting grid square
    for i=offseti:subSampStep:length(latVec)
        offsetj = randi(subSampStep);  % Randomise starting grid square
        for j=offsetj:subSampStep:length(longVec)
            % Find probes in speicifed grid
            probeselect = abs(xData(:,1)-longVec(j))<gridStep & abs(xData(:,2)-latVec(i))<gridStep;
            if any(probeselect)
                ind = find(probeselect);
                keep(ind(randi(nnz(probeselect)))) = true; % Pick one probe at random
            end
        end
    end
    xData = xData(keep,:);
    v1Data = v1Data(keep);
    v2Data = v2Data(keep);
    xStart = xStart(keep,:);
    xEnd = xEnd(keep,:);
    dt = dt(keep);
    nProbe = nnz(keep);
    month = month(keep);
    year = year(keep);
    depth = depth(keep);
    
    % Generate grids for latitude, longitude
    [longGrid,latGrid] = meshgrid(longVec,latVec);
    
    %% Find areas of sea and land on the grid
    % If land is in the central quadrant of a square then the square is land
    sea = true(size(latGrid));
    ind = find(isnan(longlatCoast(:,1)));
    ind = ind(ind(1:(length(ind)-1))+1 ~= ind(2:end));
    for i=1:(length(ind)-1),
        polyx = longlatCoast((ind(i)+1):(ind(i+1)-1),1);
        polyy = longlatCoast((ind(i)+1):(ind(i+1)-1),2);
        land = inpolygon(longGrid,latGrid,polyx,polyy);
        sea(land) = false;
    end
    
    %% Interpolate probe data to estimate the velocity field
    % Apply zeroth order approximation to the Kalman filter
    
    % Generate initial interpolation of the velocity components v1Data and
    % v2Data  (Note, the griddata functions may be replaced by the TriScatteredInterp function)    
    v1Grid = griddata(xData(:,1),xData(:,2),v1Data,longGrid,latGrid,'cubic');
    v2Grid = griddata(xData(:,1),xData(:,2),v2Data,longGrid,latGrid,'cubic');
    % Now touch up any NaN with another interpolation routine just to avoid
    % NaN appearing
    ind = (isnan(v1Grid) | isnan(v2Grid));
    if any(ind(:))
        v1Grid(ind) = griddata(xData(:,1),xData(:,2),v1Data,longGrid(ind),latGrid(ind),'v4');
        v2Grid(ind) = griddata(xData(:,1),xData(:,2),v2Data,longGrid(ind),latGrid(ind),'v4');
    end
    
    % Save the original guess at vector field
    v1Save = v1Grid;
    v2Save = v2Grid;
    
    if displayOn
        figure(1)
    end
    
    disp('90% upper quantile of the error in longitude and latitude')
    for iter = 1:nIter,  % Improve the interpolation nIter times
        if displayOn
            if iter==1
                quiver(longGrid,latGrid,v1Save,v2Save,0,'Color',[0.5,0.5,1])
                hold on
            end
        end
        
        % Loop around all the probes and integrate their position, based upon the
        % model velocity data
        odefun = @(t,y) dlonglatdt(y, longGrid, latGrid, v1Grid, v2Grid);
        xEndModel = zeros(size(xEnd));
        for p = 1:nProbe % Simulate tracks for each probe using the latest guess at the vector field
            tspan = [0 dt(p)];
            y0 = xStart(p,:); % Initial position
            [T,Y] = ode23(odefun,tspan,y0,options);  % Could use ode45 if want more accuracy
            xEndModel(p,:) = Y(end,:); % Record end position of simulated track
            
            if displayOn
                if iter==1
                    plot([xStart(p,1);xEnd(p,1)],[xStart(p,2); xEnd(p,2)],'rx-')
                    text(xData(p,1),xData(p,2),num2str(month(p)))
                    plot(Y(:,1),Y(:,2),'.-','Color',[0.5,0.5,1])
                elseif iter==nIter
                    plot(Y(:,1),Y(:,2),'.-b')
                end
            end
        end
        
        % Calculate simulated track mid-point, and simulated average velocity
        xDataModel = (xEndModel + xStart) / 2;
        v1Model = (xEndModel(:,1) - xStart(:,1)) ./ dt;
        v2Model = (xEndModel(:,2) - xStart(:,2)) ./ dt;
        
        % Update model velocity vectors based upon deviaton between model and
        % observations
        alpha = 1 + (errorSD_argo/median(dt) / errorSD_model)^2;  % Equation 7 Molcard et al
        ind = ~isnan(v1Model) & ~isnan(v2Model);
        
        correction1 = zeros(size(v1Grid));
        correction2 = zeros(size(v2Grid));
        
        for i=1:length(latVec)
            for j=1:length(longVec)
                % Correct the vector field at the start, middle and end of a
                % track
                
                % Equation 5 Molcard et al (Correct at path's midpoint)
                d2 = ((latVec(i)-xDataModel(:,2)).^2 + (longVec(j)-xDataModel(:,1)).^2)/ 2 / gridStep^2;
                gamma_ijm = exp(- d2 / distWeight^2); % Equation 6 Molcard et al
                correction1(i,j) = sum(gamma_ijm(ind).*(v1Data(ind)-v1Model(ind)));
                correction2(i,j) = sum(gamma_ijm(ind).*(v2Data(ind)-v2Model(ind)));
                
                % Equation 5 Molcard et al (correct at path's endpoint)
                d2 = ((latVec(i)-xEndModel(:,2)).^2 + (longVec(j)-xEndModel(:,1)).^2)/ 2 / gridStep^2;
                gamma_ijm = exp(- d2 / distWeight^2); % Equation 6 Molcard et al
                correction1(i,j) = correction1(i,j) + sum(gamma_ijm(ind).*(v1Data(ind)-v1Model(ind)));
                correction2(i,j) = correction2(i,j) + sum(gamma_ijm(ind).*(v2Data(ind)-v2Model(ind)));
                                
                % Equation 5 Molcard et al (correct at paths start point)
                d2 = ((latVec(i)-xStart(:,2)).^2 + (longVec(j)-xStart(:,1)).^2)/ 2 / gridStep^2;
                gamma_ijm = exp(- d2 / distWeight^2); % Equation 6 Molcard et al
                correction1(i,j) = correction1(i,j) + sum(gamma_ijm(ind).*(v1Data(ind)-v1Model(ind)));
                correction2(i,j) = correction2(i,j) + sum(gamma_ijm(ind).*(v2Data(ind)-v2Model(ind)));
            end
        end
        v1Grid = v1Grid + 1/alpha * correction1/3;
        v2Grid = v2Grid + 1/alpha * correction2/3;
        
        % Print 90% quantile of longitude and latitude error on screen
        disp(['Iteration ' num2str(iter) ': ' num2str(quantile(abs(xEndModel(ind,:) - xEnd(ind,:)),0.90))])
    end
    
    if displayOn
        plot(longlatCoast(coastSelect,1)+0*gridStep,longlatCoast(coastSelect,2)+0*gridStep,'k');
        xlabel('Longitude')
        ylabel('Latitude')
        %hold off
        axis tight
    end
    disp('===============================')
    
    %% Correct vector field around the coastline
    % to prevent particles from being swept across land
    nn4 = nan(numel(longGrid),4);
    nn8 = nan(numel(longGrid),8);
    
    coast = false(numel(longGrid),1);
    for i=1:numel(longGrid)
        % Nearest 4 neighbours
        [xInd,yInd] = find((abs(abs(longGrid-longGrid(i))-gridStep)<0.5*gridStep & abs(latGrid-latGrid(i))<0.5*gridStep) | ...
            (abs(longGrid-longGrid(i))<0.5*gridStep & abs(abs(latGrid-latGrid(i))-gridStep)<0.5*gridStep));
        nnTmp = sub2ind(size(longGrid),xInd,yInd);
        nn4(i,1:length(nnTmp)) = nnTmp;
        
        % Nearest 8 neighbours
        [xInd,yInd] = find(abs(abs(longGrid-longGrid(i))-gridStep)<0.5*gridStep & ...
            abs(abs(latGrid-latGrid(i))-gridStep)<0.5*gridStep);
        nnTmp = [nnTmp; sub2ind(size(longGrid),xInd,yInd)];
        nn8(i,1:length(nnTmp)) = nnTmp;
        
        if (any(sea(nn4(i,~isnan(nn4(i,:))))==0) && sea(i)==1) || ...
            (any(sea(nn8(i,~isnan(nn8(i,:))))==1) && sea(i)==0)
            coast(i) = true;
        end
    end
    
    ind = find(coast & ~isnan(v1Grid(:)) & ~isnan(v2Grid(:)));
    for a=1:length(ind),
        i = ind(a);
        nnInd = nn4(i,~isnan(nn4(i,:)));
        nn4Vecs = [(longGrid(nnInd)-longGrid(i))', (latGrid(nnInd)-latGrid(i))'];
        dotProd4 = (nn4Vecs * [v1Grid(i), v2Grid(i)]')./sqrt(diag(nn4Vecs*nn4Vecs'));
        [dum4,indSort4] = sort(dotProd4,'descend');
        
        nnInd = nn8(i,~isnan(nn8(i,:)));
        nn8Vecs = [(longGrid(nnInd)-longGrid(i))', (latGrid(nnInd)-latGrid(i))'];
        dotProd8 = (nn8Vecs * [v1Grid(i), v2Grid(i)]')./sqrt(diag(nn8Vecs*nn8Vecs'));
        [dum8,indSort8] = sort(dotProd8,'descend');
        
        if sea(nn4(i,indSort4(1)))==0  % If v1Grid, v2Grid pointing to land
            newInd = 1;
            while sea(nn4(i,indSort4(newInd)))==0 && newInd<nnz(isfinite(nn4(i,:)))
                % Find direction pointing toward sea
                newInd = newInd + 1;
            end
            vel = sqrt(v1Grid(i)^2 + v2Grid(i)^2);
            normFac = sqrt(nn4Vecs(indSort4(newInd),:)*nn4Vecs(indSort4(newInd),:)');
            
            %Correct vector field so that it points away from land and
            % towards sea
            v1Grid(i) = vel * nn4Vecs(indSort4(newInd),1) / normFac;
            v2Grid(i) = vel * nn4Vecs(indSort4(newInd),2) / normFac;
        end
    end
    
    if displayOn
        quiver(longGrid,latGrid,v1Grid,v2Grid,0,'b')
        hold on
        plot(longlatCoast(coastSelect,1)+0*gridStep,longlatCoast(coastSelect,2)+0*gridStep,'k');
        xlabel('Longitude')
        ylabel('Latitude')
        hold off
        axis tight
    end
    toc
    
    % Save the data to a file with a unique filename
    if ~exist(outputDir,'dir') % Check output folder exists, and create it if it doesn't
        mkdir(outputDir)
    end

    x = f;
    filename = [outputDir '\' filePref, '_', num2str(x) '.mat'];
    while exist(filename,'file')==2
        x = x+1;
         filename = [outputDir '\' filePref, '_', num2str(x)  '.mat'];
    end
    
    save(filename,'longGrid','latGrid','v1Grid','v2Grid','xData','v1Data',...
        'v2Data','xStart','xEnd','month','year','depth', 'dt','nProbe'); 
end