function [] = machineLearnDrop(plot_importance,predictions,statistics,flight_stats)
%The following code extracts and calculates predictors for the dropsonde ML algorithm 

%Inputs:    campaign: name of the field campaign for machine learning
%           split: switch for randomizing files for testing (0 or 1)
%           plot_importance: switch for generating importance plot (0 or 1)
%           predictions: switch for outputting flight predictions as .txt
%           file
%           statistics: switch for outputting stats as .txt file


%Author: Julianna Christopoulos (juliechristo@g.ucla.edu)

%Last Updated: October 14, 2024

%Output:
%           Produces predictor importance plot for regression tree 
%           Writes MLH predictions to .txt files for each flight
%           Writes statistics to .txt files for each flight 

%------------------------------ Code Begin------------------------------
    % -- Define Field Campaigns to Inlcude --
    %campaigns = {'2020_ACTIVATE', '2021_ACTIVATE', '2022_ACTIVATE', 'CPEX_CV', 'CAMP2Ex'};
    %plot_importance = 1;
    %predictions = 0;
    %statistics = 0;
    %flight_stats = 0;


   % -- Gather Files and Split into Training and Testing
    numCampaigns = length(campaigns);

     MLH_filenames = {};
     MLH_pathnames = {};

    for i = 1:numCampaigns
        campaignName = campaigns{i};
        mlh_directory = sprintf('/data2/juliechristo/%s/PBLH/', campaignName);
        MLH_directories.(['direc' num2str(i)]) = mlh_directory;
        % Get list of files in the directory
        mlh_files = dir(fullfile(mlh_directory));
        % Extract filenames and store in the array
        filenames = {mlh_files.name};
        % Filter files based on extension (.ict or .h5)
        valid_extensions = {'.csv'};
        valid_files = filenames(endsWith(filenames, valid_extensions));
        % Store filtered filenames
        MLH_filenames{i} = valid_files;
        % Construct full pathnames and store in the array
        pathnames = cellfun(@(filename) fullfile(mlh_directory, filename), valid_files, 'UniformOutput', false);
        MLH_pathnames{i} = pathnames;
    end
    all_files = horzcat(MLH_filenames{:});
    [trainFiles, testFiles] = split_files(all_files, 0.8);


     % -- Make a Strucutre of All the Data Directories -- 
    names = struct();
    for i = 1:numCampaigns
        campaignName = campaigns{i};
        %-- Initialize Data Directories --
        names.(['Direc' num2str(i)]) = sprintf('/data2/juliechristo/%s/', campaignName);
    end

    %--Extract full directories for training --
    drop_train = gatherFiles(trainFiles,names,'PBLH')';

    %--Extract full directories for testing --
    drop_test = gatherFiles(testFiles,names,'PBLH')';




    %-------------------------------------Training-----------------------------

    %-- Get Other Variables --
    [time_tr,drop_tr] = gather_Drop(drop_train);
    drop_tr = cellfun(@(x) x(:)', drop_tr, 'UniformOutput', false);
    indices_tr = find_index(drop_train,time_tr);
    alt_tr = gather_Altitude(drop_train,indices_tr,'/Altitude');
    [bsc_tr] = gather_DataProducts(drop_train,indices_tr, '/532_bsc');
    [aer_depol_tr] = gather_DataProducts(drop_train,indices_tr ,'/1064_aer_dep');
    [lat_tr] = gather_NavData(drop_train,indices_tr, '/gps_lat');
    [lon_tr] = gather_NavData(drop_train,indices_tr, '/gps_lon');
    [gps_alt_tr] = gather_NavData(drop_train,indices_tr, '/gps_alt'); %altitude of aircraft 
    lidar_tr = matchVars_MLH(drop_train,indices_tr,'/Archive/');
    [depol_tr] = matchVars_Depol(drop_train,indices_tr,'/depol/neg/p1/h5_files/');
    [rh_tr] = gather_Water(drop_train,indices_tr,'/water/RH/p1/','/MixedLayerHeight');
    rh_tr = cellfun(@(x) x(:)', rh_tr, 'UniformOutput', false);
    [mr_tr] = gather_Water(drop_train,indices_tr,'/water/MR/p1/','/MixedLayerHeight');
    mr_tr = cellfun(@(x) x(:)', mr_tr, 'UniformOutput', false);
    [p1_tr] = matchVars_Depol(drop_train,indices_tr,'/p1_drop/');
    [p2_tr] = matchVars_Depol(drop_train,indices_tr,'/p2_drop/');
    [p3_tr] = matchVars_Depol(drop_train,indices_tr,'/p3_drop/');
    [p4_tr] = matchVars_Depol(drop_train,indices_tr,'/p4_drop/');
    h_predic_tr = {p1_tr, p2_tr, p3_tr, p4_tr};     

    
    %--Compute Zenith and Solar Hour Angles--
    [zenith_tr,hangle_tr] = calcAngle(trainFiles, time_tr, lat_tr, lon_tr);
    hangle_tr = cellfun(@(x) x(:)', hangle_tr, 'UniformOutput', false);
 
    %-- Initialize Variances --
    %-- Calculate for 532nm backscatter --
    [horz_variances, mlh_means] = cellfun(@(x) calcVar_MLH(x, time_tr, 10), h_predic_tr, 'UniformOutput', false);
    [hvar_p1_tr, hvar_p2_tr, hvar_p3_tr, hvar_p4_tr] = horz_variances{:};
   
    %-- Calculate for 1064nm Aerosol Depolarization Ratio --
    [hvar_depol_tr,mean_depol_tr] = calcVar_MLH(depol_tr, time_tr, 10);
     
    %-- Get land/water flags from .txt files --
    flag_tr = checkCoords(lat_tr,lon_tr);

    % -- Add on Variance Predictors to Original Height Predictors
    predic_all_tr = {p1_tr,p2_tr,p3_tr,p4_tr,hvar_p1_tr,hvar_p2_tr,hvar_p3_tr,hvar_p4_tr,hangle_tr,depol_tr,rh_tr,mr_tr};
    

        %-- Find Minimums and Maximums -- 
    [min_tr, max_tr] = cellfun(@find_range, predic_all_tr, 'UniformOutput', false);
    [min_p1,min_p2,min_p3,min_p4,min_hvar_p1,min_hvar_p2,min_hvar_p3,min_hvar_p4,min_hangle,min_depol,min_rh,min_mr] = min_tr{:};

    [max_p1,max_p2,max_p3,max_p4,max_hvar_p1,max_hvar_p2,max_hvar_p3,max_hvar_p4,max_hangle,max_depol,max_rh,max_mr] = max_tr{:};

    
    %Archived Mins and Maxs
    [min_drop, max_drop] = find_range(drop_tr);

    %-- Standardize Values for Machine Learning --
    std_values_tr = cellfun(@(x, min_val, max_val) standardizeData(x, min_val, max_val), predic_all_tr, min_tr, max_tr, 'UniformOutput', false);
    [sp1_tr,sp2_tr,sp3_tr,sp4_tr,shvar_p1_tr,shvar_p2_tr,shvar_p3_tr,shvar_p4_tr,shangle_tr,sdepol_tr,srh_tr,smr_tr] = std_values_tr{:};
    [sdrop_tr] = standardizeData(drop_tr,min_drop,max_drop);

    %-------------------------------------Testing-----------------------------
 
    %-- Get Other Variables --
    [time_tst,drop_tst] = gather_Drop(drop_test);
    drop_tst = cellfun(@(x) x(:)', drop_tst, 'UniformOutput', false);
    indices_tst = find_index(drop_test,time_tst);
    alt_tst = gather_Altitude(drop_test,indices_tst,'/Altitude');
    [bsc_tst] = gather_DataProducts(drop_test,indices_tst, '/532_bsc');
    [aer_depol_tst] = gather_DataProducts(drop_test,indices_tst ,'/1064_aer_dep');
    gps_time_tst = gather_NavData(drop_test,indices_tst, '/gps_time');
    [lat_tst] = gather_NavData(drop_test,indices_tst, '/gps_lat');
    [lon_tst] = gather_NavData(drop_test,indices_tst, '/gps_lon');
    [gps_alt_tst] = gather_NavData(drop_test,indices_tst, '/gps_alt'); %altitude of aircraft 
    lidar_tst = matchVars_MLH(drop_test,indices_tst,'/Archive/');
    [depol_tst] = matchVars_Depol(drop_test,indices_tst,'/depol/neg/p1/h5_files/');
    [rh_tst] = gather_Water(drop_test,indices_tst,'/water/RH/p1/','/MixedLayerHeight');
    rh_tst = cellfun(@(x) x(:)', rh_tst, 'UniformOutput', false);
    [mr_tst] = gather_Water(drop_test,indices_tst,'/water/MR/p1/','/MixedLayerHeight');
    mr_tst = cellfun(@(x) x(:)', mr_tst, 'UniformOutput', false);
    [p1_tst] = matchVars_Depol(drop_test,indices_tst,'/p1_drop/');
    [p2_tst] = matchVars_Depol(drop_test,indices_tst,'/p2_drop/');
    [p3_tst] = matchVars_Depol(drop_test,indices_tst,'/p3_drop/');
    [p4_tst] = matchVars_Depol(drop_test,indices_tst,'/p4_drop/');
    h_predic_tst = {p1_tst, p2_tst, p3_tst, p4_tst};



    %--Compute Zenith and Solar Hour Angles--
    [zenith_tst,hangle_tst] = calcAngle(testFiles, time_tst, lat_tst, lon_tst);
    hangle_tst = cellfun(@(x) x(:)', hangle_tst, 'UniformOutput', false);

    %-- Initialize Variances --
    %-- Calculate for 532nm backscatter --
    [horz_variances_tst,mlh_means_tst] = cellfun(@(x) calcVar_MLH(x, time_tst, 10), h_predic_tst, 'UniformOutput', false);
    [hvar_p1_tst, hvar_p2_tst, hvar_p3_tst, hvar_p4_tst] = horz_variances_tst{:};

    %-- Calculate for 1064nm Aerosol Depolarization Ratio --
    [hvar_depol_tst,mean_depol_tst] = calcVar_MLH(depol_tst, time_tst, 10);


    %-- Get land/water flags from .txt files --
    flag_tst = checkCoords(lat_tst,lon_tst);

    % -- Add on Variance Predictors to Original Height Predictors
    predic_all_tst = {p1_tst,p2_tst,p3_tst,p4_tst,hvar_p1_tst,hvar_p2_tst,hvar_p3_tst,hvar_p4_tst,hangle_tst,depol_tst,rh_tst,mr_tst}; 
 
    %-- Standardize Values for Machine Learning --
    std_values_tst = cellfun(@(x, min_val, max_val) standardizeData(x, min_val, max_val), predic_all_tst, min_tr, max_tr, 'UniformOutput', false);
    [sp1_tst,sp2_tst,sp3_tst,sp4_tst,shvar_p1_tst,shvar_p2_tst,shvar_p3_tst,shvar_p4_tst,shangle_tst,sdepol_tst,srh_tst,smr_tst] = std_values_tst{:};
    [sdrop_tst] = standardizeData(drop_tst,min_drop,max_drop);

 
    % -- Put the training and testing predictros into a table and output as
    % a .mat -- 
    %Flatten the Training Arrays 
    train_drop = cat(2, sdrop_tr{:})';
    % Initialize flattened array
    train_data = [];
    % Flatten vertically along columns
    for i = 1:numel(std_values_tr)
        train_data = [train_data; cat(2, std_values_tr{i}{:})];
    end
    %Write into Columns
    train_data = train_data';

    %Remove the last two columns for linear and neural
    train_data_removed = train_data(:, 1:end-2);

   
%Edit the Columns for training linear regression and neural (no nans)   
     medians_train = zeros(1, size(train_data, 2));
     train_data_clean = train_data;
     for col = 1:size(train_data_clean,2)
         columnData = train_data_clean(:,col);
         nanIndices = isnan(columnData);
         median_value = nanmedian(columnData);
         medians_train(col) = median_value;
         columnData(nanIndices) = median_value;
         train_data_clean(:,col) = columnData;
     end

    
     train_data_clean_removed = train_data_removed;
     for col_removed = 1:size(train_data_clean_removed,2)
         columnData_removed = train_data_clean_removed(:,col_removed);
         nanIndices_removed = isnan(columnData_removed);
         medianValue_removed = nanmedian(columnData_removed);
         columnData_removed(nanIndices_removed) = medianValue_removed;
         train_data_clean_removed(:,col_removed) = columnData_removed;
     end  

     %Do the same for the responses
     train_drop_clean = train_drop;
     medianDrop = nanmedian(train_drop_clean);
     nanDrop = isnan(train_drop_clean);
     train_drop_clean(nanDrop) = medianDrop;



    %------------------------Do the Machine Learning---------------------------

     %--Linear Regression Model--
     linear_water = fitrlinear(train_data_clean,train_drop_clean,'Regularization','ridge');

     linear_no_water = fitrlinear(train_data_clean_removed,train_drop_clean,'Regularization','ridge');

     %--Tree Ensembles Model--
     treeTemplate = templateTree('Surrogate','on');
     ensemble = fitrensemble(train_data,train_drop,'Method','Bag','Learners',treeTemplate);
     
     %--Neural Networks model--
     neural_water = fitrnet(train_data_clean,train_drop_clean);

     neural_no_water = fitrnet(train_data_clean_removed,train_drop_clean);


    %--Predictor Importance--

    if plot_importance == 1
        % Calculate the predictor importance
        importance = predictorImportance(ensemble);
        predictor_labels = {'P1','P2','P3','P4','P5','P6','P7','P8','P9','P10','P11','P12'};

        % Plot the bar chart
        bar_handle = bar(importance);

        % Set the color for each bar
        bar_handle.EdgeColor = 'none';

        % Adjust x-axis tick positions and labels
        xlabel('Predictor','FontWeight','bold','FontSize',12);
        ylabel('Importance','FontWeight','bold','FontSize',12);
        xticks(1:numel(importance));
        xticklabels(predictor_labels);
        xtickangle(90); % Optional: Rotate the x-axis labels if needed
        title('Marine Ensemble Model','FontWeight','bold','FontSize',14);

        ax = gca;
        set(gca, 'XTickLabel', get(gca, 'XTickLabel'), 'FontWeight', 'bold');

        % Save the plot as an image file
        print('importance_water_mdl_bold.png', '-dpng', '-r300');
    end    


 
     mlh_linear = cell(length(sp1_tst), 1);
     mlh_ens = cell(length(sp1_tst), 1);
     mlh_neural = cell(length(sp1_tst), 1);
     %Make Predictions for each flight
     for i = 1:length(testFiles)
         test_MLH = sp1_tst{i};
         nan_index = find(isnan(test_MLH));
         test_data = [sp1_tst{i}', sp2_tst{i}', sp3_tst{i}', sp4_tst{i}',shvar_p1_tst{i}', shvar_p2_tst{i}', shvar_p3_tst{i}', shvar_p4_tst{i}',shangle_tst{i}',sdepol_tst{i}',srh_tst{i}',smr_tst{i}'];
         test_data_clean = test_data;
         %Fix nans for linear and neural problem
         for column = 1:size(test_data_clean,2)
             colNans_test = isnan(test_data_clean(:, column));
             test_data_clean(colNans_test, column) = medians_train(column);
         end

         %Linear Prediction + recover units
         if  all(isnan(test_data_clean(:, end-1))) && all(isnan(test_data_clean(:, end))) 
             test_data_clean_removed = test_data_clean(:, 1:end-2);
             linear_dat = predict(linear_no_water,test_data_clean_removed).* (max_drop - min_drop) + min_drop;
             %linear_dat = exp(linear_dat) .* (max_drop - min_drop) + min_drop;
             linear_dat(nan_index) = nan ;
             mlh_linear{i} = linear_dat';
         else
             linear_dat = predict(linear_water,test_data_clean).* (max_drop - min_drop) + min_drop;
             %linear_dat = exp(linear_dat) .* (max_drop - min_drop) + min_drop;
             linear_dat(nan_index) = nan ;
             mlh_linear{i} = linear_dat';
         end 
         %Ensemble Prediction + recover units
         ens_dat = predict(ensemble,test_data).* (max_drop - min_drop) + min_drop;
         %ens_dat = exp(ens_dat) .* (max_drop - min_drop) + min_drop;
         ens_dat(nan_index) = nan;
         mlh_ens{i} = ens_dat';
         %Neural Prediction + recover units
         if all(isnan(test_data_clean(:, end-1))) && all(isnan(test_data_clean(:, end)))    
             test_data_clean_removed = test_data_clean(:, 1:end-2);
             neural_dat = predict(neural_no_water,test_data_clean_removed).* (max_drop - min_drop) + min_drop;
             %neural_dat = exp(neural_dat) .* (max_drop - min_drop) + min_drop;
             neural_dat(nan_index) = nan;
             mlh_neural{i} = neural_dat';
         else
             neural_dat = predict(neural_water,test_data_clean).* (max_drop - min_drop) + min_drop;
             %neural_dat = exp(neural_dat) .* (max_drop - min_drop) + min_drop;
             neural_dat(nan_index) = nan;
             mlh_neural{i} = neural_dat';
         end 
     end


    %---------Ouput the Predictions for each flight-------------------------- 
    if predictions == 1
        for o = 1:length(mlh_linear)
            time_col = gps_time_tst{o}';
            linear_col = mlh_linear{o}';
            ens_col = mlh_ens{o}';
            neural_col = mlh_neural{o}';
            drop_col = drop_tst{o}';
            lidar_col = lidar_tst{o}';
 
            data = table(time_col,linear_col, ens_col, neural_col,drop_col,lidar_col,'VariableNames', {'Time','Linear', 'Tree', 'Ensemble', 'Neural','Drop','Lidar'});
            currentfile = testFiles{o};
            date = currentfile(1:8);
            filename = sprintf('%s.txt', date);
            writetable(data, filename, 'Delimiter', '\t');
        end
    end 
    %-----------------------Compute Statistics of ML Predictions----------------------
 
    %---------------------------------------Statistics for the entire campaign--------------------------------
    if statistics == 1 
        dropsonde_all = [];
        lidar_all = [];
        def_all = [];
        ensemble_all = [];
        linear_all = [];
        neural_all = [];
 
 
        % Flatten vertically along columns
        for i = 1:numel(drop_tst)
            dropsonde_all = [dropsonde_all; cat(2, drop_tst{i}(:))];
            lidar_all = [lidar_all; cat(2, lidar_tst{i}(:))];
            def_all = [def_all; cat(2, p2_tst{i}(:))];
            ensemble_all = [ensemble_all; cat(2, mlh_ens{i}(:))];
            linear_all = [linear_all; cat(2, mlh_linear{i}(:))];
            neural_all = [neural_all; cat(2, mlh_neural{i}(:))];
        end
 
        % Nans
        drop_nans = ~isnan(dropsonde_all);
        lidar_nans = ~isnan(lidar_all);
        def_nans = ~isnan(def_all);
        ensemble_nans = ~isnan(ensemble_all);
        linear_nans = ~isnan(linear_all);
        neural_nans = ~isnan(neural_all);
        common_nans = drop_nans & lidar_nans & def_nans & ensemble_nans & linear_nans & neural_nans;

        %Lidar
        [nmb_lidar,nme_lidar,mb_lidar,me_lidar,rmse_lidar,corr_lidar] = calcStats_all_updated(common_nans,dropsonde_all,lidar_all); %default
 
        %Default
        [nmb_def,nme_def,mb_def,me_def,rmse_def,corr_def] = calcStats_all_updated(common_nans,dropsonde_all,def_all); %default

        %Ensemble Tree Model
        [nmb_ens,nme_ens,mb_ens,me_ens,rmse_ens,corr_ens] = calcStats_all_updated(common_nans,dropsonde_all,ensemble_all);
 
        %Linear Model
        [nmb_linear,nme_linear,mb_linear,me_linear,rmse_linear,corr_linear] = calcStats_all_updated(common_nans,dropsonde_all,linear_all);
 
        %Neural Model 
        [nmb_neural,nme_neural,mb_neural,me_neural,rmse_neural,corr_neural] = calcStats_all_updated(common_nans,dropsonde_all,neural_all);
 
        %--Output to .txt file--
        NMB = [nmb_lidar, nmb_def,nmb_ens, nmb_linear, nmb_neural];
        NME = [nme_lidar, nme_def,nme_ens, nme_linear, nme_neural];
        MB = [mb_lidar, mb_def,mb_ens, mb_linear, mb_neural];
        ME = [me_lidar, me_def,me_ens, me_linear, me_neural];
        RMSE = [rmse_lidar, rmse_def, rmse_ens, rmse_linear, rmse_neural];
        CORR = [corr_lidar, corr_def, corr_ens, corr_linear, corr_neural];
 
 
        % Round the values to three decimal places
        NMB = round(NMB, 2);
        NME = round(NME, 2);
        MB = round(MB, 2);
        ME = round(ME, 2);
        RMSE = round(RMSE, 2);
        CORR = round(CORR, 2);
 
        columnNames = {'NMB', 'NME', 'MB', 'ME', 'RMSE', 'CORR'};
        rowNames = {'Lidar', 'Default', 'Ensemble', 'Linear', 'Neural'};
        tbl = table(NMB', NME', MB', ME', RMSE', CORR', 'VariableNames', columnNames, 'RowNames', rowNames);
 
 
 
        filename = 'DropStats_check.txt';
        writetable(tbl, filename, 'Delimiter', '\t');
 
    end 
 
 
     %----------------------------------------- Statistics for each flight---------------------------------------------------- 
     
    if flight_stats == 1
        %Default
        [nmb_def,nme_def,mb_def,me_def,rmse_def,corr_def] = calcStats(drop_tst,lidar_tst); %default
  
        %Ensemble Tree Model
        [nmb_ens,nme_ens,mb_ens,me_ens,rmse_ens,corr_ens] = calcStats(drop_tst,mlh_ens);
 
        %Linear Model
        [nmb_linear,nme_linear,mb_linear,me_linear,rmse_linear,corr_linear] = calcStats(drop_tst,mlh_linear);
 
        %Neural Model 
        [nmb_neural,nme_neural,mb_neural,me_neural,rmse_neural,corr_neural] = calcStats(drop_tst,mlh_neural);
 
 
        for p = 1:length(all_files)
            currentfile = all_files{p};
            date = currentfile(1:8); % Retrieve the date for the current iteration
 
            NMB = [nmb_def{p}, nmb_ens{p}, nmb_linear{p}, nmb_neural{p}];
            NME = [nme_def{p}, nme_ens{p}, nme_linear{p}, nme_neural{p}];
            MB = [mb_def{p}, mb_ens{p}, mb_linear{p}, mb_neural{p}];
            ME = [me_def{p}, me_ens{p}, me_linear{p}, me_neural{p}];
            RMSE = [rmse_def{p}, rmse_ens{p}, rmse_linear{p}, rmse_neural{p}];
            CORR = [corr_def{p}, corr_ens{p}, corr_linear{p}, corr_neural{p}];
 
            %Round the values to three decimal places
            NMB = round(NMB, 2);
            NME = round(NME, 2);
            MB = round(MB, 2);
            ME = round(ME, 2);
            RMSE = round(RMSE, 2);
            CORR = round(CORR, 2);
 
            columnNames = {'NMB', 'NME', 'MB', 'ME', 'RMSE', 'CORR'};
            rowNames = {'Lidar', 'Ensemble', 'Linear', 'Neural'};
            tbl = table(NMB', NME', MB', ME', RMSE', CORR', 'VariableNames', columnNames, 'RowNames', rowNames);
 
            %Save the table as a variable with a unique name
            varName = sprintf('tbl_%s', date);
            assignin('base', varName, tbl);
 
            filename = sprintf('flight_%s.txt', date);
            writetable(tbl, filename, 'Delimiter', '\t');
        end
    end 
 
 end 
 












 
     
