loadRAMPdata.m

function ctmData = loadRAMPdata(modelName, years, dataDir, forceReload)
% loadRAMPdata - Load RAMP-corrected CTM data from parquet files
%
% Reads lambda1 (mean) and lambda2 (variance) parquet files for specified
% years and creates a cached .mat file for fast subsequent loading.
%
% SYNTAX:
%   ctmData = loadRAMPdata(modelName, years, dataDir, forceReload)
%
% INPUTS:
%   modelName   - Model identifier (e.g., 'UKML')
%   years       - Years to load (e.g., [2015 2016 2017])
%   dataDir     - Directory containing parquet files (default: '1data/CTM')
%   forceReload - Force reload from parquet (1) or use cache (0, default)
%
% OUTPUTS:
%   ctmData - Structure with fields:
%             .modelName  - Model identifier
%             .years      - Years included
%             .lon        - Longitude vector [nGrid × 1]
%             .lat        - Latitude vector [nGrid × 1]
%             .sMS        - Spatial coordinates [nGrid × 2] as [lon, lat]
%             .tME        - Time vector in decimal years [1 × nMonths]
%             .Z          - Mean field (lambda1) [nGrid × nMonths]
%             .Zv         - Variance field (lambda2) [nGrid × nMonths]
%             .Zunit      - Units string
%             .version    - RAMP version
%             .loadedAt   - Timestamp
%             .gridInfo   - Grid metadata (nGridPoints, yearsChecked, isConsistent)
%
% EXAMPLE:
%   % Load 2015-2020 UKML data
%   ctm = loadRAMPdata('UKML', [2015:2020]);
%
%   % Force reload from parquet files
%   ctm = loadRAMPdata('UKML', [2015:2020], '1data/CTM', 1);
%
% FILE NAMING CONVENTION:
%   Spatial grid .mat files (generated by extractModelSpatialInfo.m):
%     1data/CTM/model_output_data/spatial_grids/{modelName}_spatial_grid.mat
%     Contains: lon, lat, nGridPoints, yearsChecked, isConsistent
%
%   RAMP-corrected parquet files:
%     lambda1_{modelName}_{year}_v{version}-parallel.parquet
%     lambda2_{modelName}_{year}_v{version}-parallel.parquet
%
%   Example:
%     1data/CTM/model_output_data/spatial_grids/M3fusion_spatial_grid.mat
%     1data/CTM/lambda1_M3fusion_2017_v3-parallel.parquet
%     1data/CTM/lambda2_M3fusion_2017_v3-parallel.parquet
%
% DATA STRUCTURE:
%   - Spatial grid (lon, lat) is read from .mat files (see extractModelSpatialInfo.m)
%   - Parquet files contain RAMP-corrected mean (lambda1) and variance (lambda2)
%   - Parquet files have 12 columns (one per month), NO spatial information
%   - Row order in parquet must match grid order from spatial .mat file
%
% PREREQUISITES:
%   Before using this function, you must run extractModelSpatialInfo.m to
%   generate the spatial grid .mat files from the original CSV model outputs.

%% Input Validation
if nargin < 1 || isempty(modelName), modelName = 'UKML'; end
if nargin < 2 || isempty(years), years = 2015:2020; end
if nargin < 3 || isempty(dataDir), dataDir = fullfile('1data', 'CTM'); end
if nargin < 4, forceReload = 0; end

% Ensure years is a vector
if isscalar(years)
    years = years:years;
end

fprintf('\n========================================\n');
fprintf('  LOAD RAMP-CORRECTED CTM DATA\n');
fprintf('========================================\n');
fprintf('Model: %s\n', modelName);
fprintf('Years: %s\n', mat2str(years));
fprintf('Data directory: %s\n', dataDir);

%% Check Cache
cacheDir = fullfile('1data', 'CTM');
if ~exist(cacheDir, 'dir')
    mkdir(cacheDir);
end

% Determine RAMP version from first file
sampleFile = dir(fullfile(dataDir, sprintf('lambda1_%s_%d_v3-parallel.parquet', modelName, years(1))));
if isempty(sampleFile)
    warning('No parquet files found for %s year %d in %s', modelName, years(1), dataDir);
    tokens = [];
else
    % Extract version from filename (e.g., "v3" from "..._v3-parallel.parquet")
    tokens = regexp(sampleFile(1).name, '_v(\d+)-parallel', 'tokens');
end

if ~isempty(tokens)
    rampVersion = str2double(tokens{1}{1});
else
    rampVersion = 3;  % Default
end

% Try to find existing cache file first
cachePattern = sprintf('CTM_RAMP_%s_*_v%d.mat', modelName, rampVersion);
existingCaches = dir(fullfile(cacheDir, cachePattern));

cacheFound = false;
if ~forceReload && ~isempty(existingCaches)
    % Check if any existing cache has all requested years
    for i = 1:length(existingCaches)
        cachePath = fullfile(existingCaches(i).folder, existingCaches(i).name);
        try
            load(cachePath, 'ctmData');
            % Check if cached data contains all requested years
            if all(ismember(years, ctmData.years))
                fprintf('\nFound compatible cache: %s\n', existingCaches(i).name);
                fprintf('  Cached years: %s\n', mat2str(ctmData.years));
                fprintf('  Requested years: %s\n', mat2str(years));

                % Extract only requested years
                yearIdx = ismember(ctmData.years, years);
                monthIdx = [];
                for iY = 1:length(ctmData.years)
                    if yearIdx(iY)
                        monthIdx = [monthIdx, (iY-1)*12 + (1:12)];
                    end
                end

                % Subset the data
                ctmData.Z = ctmData.Z(:, monthIdx);
                ctmData.Zv = ctmData.Zv(:, monthIdx);
                ctmData.tME = ctmData.tME(monthIdx);
                ctmData.years = years;
                ctmData.nMonths = length(monthIdx);

                fprintf('  Grid points: %d\n', length(ctmData.lon));
                fprintf('  Time periods: %d\n', length(ctmData.tME));
                fprintf('  Data completeness: %.1f%%\n', 100*sum(~isnan(ctmData.Z(:)))/numel(ctmData.Z));
                fprintf('========================================\n\n');
                cacheFound = true;
                return;
            end
        catch
            continue;
        end
    end
end

if ~cacheFound
    fprintf('\nNo compatible cache found or force reload requested.\n');
end

%% Read Spatial Grid from .mat File
fprintf('\nLoading spatial grid from .mat file...\n');

% Path to spatial grid .mat file
spatialGridDir = fullfile('1data', 'CTM', 'model_output_data', 'spatial_grids');
spatialGridFile = sprintf('%s_spatial_grid.mat', modelName);
spatialGridPath = fullfile(spatialGridDir, spatialGridFile);

% Check if .mat file exists
if ~exist(spatialGridPath, 'file')
    error(['Spatial grid file not found: %s\n' ...
           'Please run extractModelSpatialInfo.m first to generate spatial grid files.'], ...
           spatialGridPath);
end

% Load spatial grid
fprintf('  Loading: %s\n', spatialGridFile);
tic;
gridData = load(spatialGridPath);
tLoad = toc;

% Extract coordinates
lon_grid = gridData.lon(:);  % Ensure column vector
lat_grid = gridData.lat(:);  % Ensure column vector
nGrid = gridData.nGridPoints;

fprintf('  Loaded in %.2f seconds\n', tLoad);
fprintf('  Grid points: %d\n', nGrid);
fprintf('  Lon range: [%.2f, %.2f]\n', min(lon_grid), max(lon_grid));
fprintf('  Lat range: [%.2f, %.2f]\n', min(lat_grid), max(lat_grid));
fprintf('  Spatial consistency: %s\n', iif(gridData.isConsistent, 'PASS', 'FAIL'));

if ~gridData.isConsistent
    warning('Spatial grid is not consistent across all years. Proceed with caution.');
end

%% Load Parquet Files
fprintf('\nLoading RAMP-corrected data from parquet files...\n');

nYears = length(years);

% Track which years were successfully loaded
yearsLoaded = [];
lambda1_list = {};
lambda2_list = {};
tME_list = {};

for iYear = 1:nYears
    year = years(iYear);

    fprintf('\n--- Processing Year %d ---\n', year);

    % Construct filenames
    lambda1File = sprintf('lambda1_%s_%d_v%d-parallel.parquet', modelName, year, rampVersion);
    lambda2File = sprintf('lambda2_%s_%d_v%d-parallel.parquet', modelName, year, rampVersion);

    lambda1Path = fullfile(dataDir, lambda1File);
    lambda2Path = fullfile(dataDir, lambda2File);

    % Try to load this year's data
    try
        % Check files exist
        if ~exist(lambda1Path, 'file')
            error('Lambda1 file not found: %s', lambda1Path);
        end
        if ~exist(lambda2Path, 'file')
            error('Lambda2 file not found: %s', lambda2Path);
        end

        % Read lambda1 (mean)
        fprintf('  Reading %s...\n', lambda1File);
        tic;
        lambda1_table = parquetread(lambda1Path);
        tRead1 = toc;
        fprintf('    Loaded in %.2f seconds (%d rows)\n', tRead1, height(lambda1_table));

        % Read lambda2 (variance)
        fprintf('  Reading %s...\n', lambda2File);
        tic;
        lambda2_table = parquetread(lambda2Path);
        tRead2 = toc;
        fprintf('    Loaded in %.2f seconds (%d rows)\n', tRead2, height(lambda2_table));

        % Verify parquet file has correct number of rows (should match grid)
        if height(lambda1_table) ~= nGrid
            error('Lambda1 file has %d rows, expected %d (grid size)', ...
                height(lambda1_table), nGrid);
        end
        if height(lambda2_table) ~= nGrid
            error('Lambda2 file has %d rows, expected %d (grid size)', ...
                height(lambda2_table), nGrid);
        end

        % Extract monthly data (columns 1-12, no spatial info in parquet)
        lambda1_year = single(lambda1_table{:, 1:12});
        lambda2_year = single(lambda2_table{:, 1:12});

        % Create time vector for this year
        tME_year = year + (0:11)/12;

        fprintf('  Data range - Mean: [%.2f, %.2f] ppb\n', ...
            min(lambda1_year(:), [], 'omitnan'), max(lambda1_year(:), [], 'omitnan'));
        fprintf('  Data range - Variance: [%.4f, %.2f] ppb²\n', ...
            min(lambda2_year(:), [], 'omitnan'), max(lambda2_year(:), [], 'omitnan'));

        % Store successfully loaded data
        lambda1_list{end+1} = lambda1_year;
        lambda2_list{end+1} = lambda2_year;
        tME_list{end+1} = tME_year;
        yearsLoaded(end+1) = year;

        fprintf('  ✓ Year %d loaded successfully\n', year);

    catch ME
        fprintf('\n');
        fprintf('========================================\n');
        fprintf('  WARNING: Failed to load year %d\n', year);
        fprintf('========================================\n');
        fprintf('Model: %s\n', modelName);
        fprintf('Year: %d\n', year);
        fprintf('Error: %s\n', ME.message);
        fprintf('Expected files:\n');
        fprintf('  - %s\n', lambda1Path);
        fprintf('  - %s\n', lambda2Path);
        fprintf('\nThis year will be SKIPPED. Continuing with remaining years...\n');
        fprintf('========================================\n\n');
    end
end

% Check if any years were successfully loaded
if isempty(yearsLoaded)
    error('No data could be loaded for any requested year. Requested: %s', mat2str(years));
end

% Report loading summary
fprintf('\n--- Loading Summary ---\n');
fprintf('  Requested years: %s\n', mat2str(years));
fprintf('  Successfully loaded: %s\n', mat2str(yearsLoaded));
if length(yearsLoaded) < length(years)
    missingYears = setdiff(years, yearsLoaded);
    fprintf('  ⚠ Missing years: %s\n', mat2str(missingYears));
end

% Concatenate all loaded data
lambda1_all = cat(2, lambda1_list{:});
lambda2_all = cat(2, lambda2_list{:});
tME_all = cat(2, tME_list{:});

% Update years to reflect what was actually loaded
years = yearsLoaded;
nMonths = length(tME_all);

%% Quality Control
fprintf('\n--- Quality Control ---\n');

% Count NaN values
nNaN_mean = sum(isnan(lambda1_all(:)));
nNaN_var = sum(isnan(lambda2_all(:)));
nTotal = numel(lambda1_all);

fprintf('  NaN in mean: %d / %d (%.2f%%)\n', nNaN_mean, nTotal, 100*nNaN_mean/nTotal);
fprintf('  NaN in variance: %d / %d (%.2f%%)\n', nNaN_var, nTotal, 100*nNaN_var/nTotal);

% Check for negative variance
nNegVar = sum(lambda2_all(:) < 0, 'omitnan');
if nNegVar > 0
    warning('Found %d negative variance values (%.2f%%). Setting to small positive.', ...
        nNegVar, 100*nNegVar/nTotal);
    lambda2_all(lambda2_all < 0) = 0.01;  % Small positive variance
end

% Check for infinite values
nInf_mean = sum(isinf(lambda1_all(:)));
nInf_var = sum(isinf(lambda2_all(:)));
if nInf_mean > 0 || nInf_var > 0
    warning('Found %d infinite mean and %d infinite variance values. Setting to NaN.', ...
        nInf_mean, nInf_var);
    lambda1_all(isinf(lambda1_all)) = NaN;
    lambda2_all(isinf(lambda2_all)) = NaN;
end

%% Package Output Structure
fprintf('\n--- Creating CTM Data Structure ---\n');

% Create spatial coordinates matrix [lon, lat]
sMS = [lon_grid, lat_grid];

ctmData.modelName = modelName;
ctmData.years = years;
ctmData.lon = lon_grid;
ctmData.lat = lat_grid;
ctmData.sMS = sMS;  % Spatial coordinates [nGrid × 2] as [lon, lat]
ctmData.tME = tME_all;
ctmData.Z = lambda1_all;      % Mean field (lambda1)
ctmData.Zv = lambda2_all;     % Variance field (lambda2)
ctmData.Zname = sprintf('%s-RAMP', modelName);
ctmData.Zunit = 'ppb';
ctmData.Zlabel = sprintf('%s RAMP-corrected MDA8 Ozone', modelName);
ctmData.version = rampVersion;
ctmData.loadedAt = datestr(now);
ctmData.nGrid = length(lon_grid);
ctmData.nMonths = nMonths;
ctmData.gridInfo = struct('nGridPoints', gridData.nGridPoints, ...
                         'yearsChecked', gridData.yearsChecked, ...
                         'isConsistent', gridData.isConsistent, ...
                         'source', 'mat_file', ...
                         'matFile', spatialGridFile);  % Include grid metadata

fprintf('  Model: %s (RAMP v%d)\n', modelName, rampVersion);
fprintf('  Grid points: %d\n', ctmData.nGrid);
fprintf('  Time periods: %d months (%d years)\n', nMonths, nYears);
fprintf('  Mean range: [%.2f, %.2f] %s\n', ...
    min(ctmData.Z(:), [], 'omitnan'), max(ctmData.Z(:), [], 'omitnan'), ctmData.Zunit);
fprintf('  Variance range: [%.4f, %.2f] %s²\n', ...
    min(ctmData.Zv(:), [], 'omitnan'), max(ctmData.Zv(:), [], 'omitnan'), ctmData.Zunit);

%% Save Cache
fprintf('\n--- Saving Cache ---\n');

% Update cache filename to reflect actually loaded years
cacheFile = sprintf('CTM_RAMP_%s_%d-%d_v%d.mat', ...
    modelName, min(years), max(years), rampVersion);
cachePath = fullfile(cacheDir, cacheFile);

fprintf('  File: %s\n', cacheFile);
fprintf('  Years: %s\n', mat2str(years));

tic;
save(cachePath, 'ctmData', '-v7.3');
tSave = toc;

fileInfo = dir(cachePath);
fprintf('  Saved in %.2f seconds\n', tSave);
fprintf('  File size: %.1f MB\n', fileInfo.bytes / 1024^2);

fprintf('\n========================================\n');
fprintf('  CTM DATA LOADING COMPLETE\n');
fprintf('========================================\n\n');

end

%% Helper Functions

function result = iif(condition, trueVal, falseVal)
    % Inline if function
    if condition
        result = trueVal;
    else
        result = falseVal;
    end
end