//////#################################################################################################################################################################
//////########### MONTHLY CROP WATER AND IRRIGATION WATER USE ACROSS CROPLAND #########################################################################################
//////#################################################################################################################################################################
////// See paper: Can remotely-sensed surface water flux be used to estimate net groundwater storage change in an aquifer heavily used for irrigation? 
////// by Cindy Viviers and Michael van der Laan
 
//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
///////////////////////################################################################################################################################################
///////////////////////###################### DATA PREPERATION ########################################################################################################
///////////////////////################################################################################################################################################
//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

// /// Dataset variables names
// // Description - collection ref - band ref
// 1: Monthly cultivation identified per img - classifiedcultivation - 'classification'
// 2: Monthly cultivation resampled (ratio) to WaPOR - cultivation250mMean - 'cultivation250Mean'
// 3: Monthly ETa WaPOR - et250 - 'et250'
// 4: Monthly CHIRPS precipitation resampled (and projected) to WaPOR - rainfall250 - 'precipitation'
// 5: Monthly downscaled GWS resampled (and projected) to WaPOR - downscaledGWS250 - 'downscaledCorre_M2'
 

///////////////////////##################### CROPLAND ##################################
//// See code: XXX
//// Monthly cultivated areas identified from Sentinel-2 Images - exported from script above
var collection_Cropland1 = ee.ImageCollection(Cropland)
    .filterBounds(roi)
    .map(function(image){return image.clip(roi)});
    
var collection_Cropland = collection_Cropland1.sort('system:time_start'); // Sort image collection chronologically
var firstImage1 = collection_Cropland.first();
var projection1 = firstImage1.projection();
print('Nominal scale of the image - Cultivation Classified:', projection1.nominalScale());

/////////////////////##################### CULTIVATION RATIO ACROSS WaPOR RESOLUTION PIXEL 
var waporET = ee.ImageCollection('FAO/WAPOR/2/L1_AETI_D').first();
// Function to reproject and resample cultivated images to match WaPOR ETa using mean.
// Product is a decimal (% if multiplied by 100) that cultivation makes up of WaPOR pixel area
var resampleToET = function(image) {
  var resampledImage = image
    .select('classification').rename('cultivation250Mean')
    .reduceResolution({
      reducer: ee.Reducer.mean(),
      maxPixels: 1024
    })
    .reproject({
      crs: waporET.projection(),
      scale: waporET.projection().nominalScale()
    });
  return resampledImage;
};

// Apply the function to the landcover collection.
var resampledLandcoverCollectionWAPOR = collection_Cropland.map(resampleToET);
Map.addLayer(resampledLandcoverCollectionWAPOR.first(),{min:0, max:1, palette:['white', 'lightgrey', 'grey', 'darkgrey']},'Resampled Cultivation Collection to WAPOR Resolution');

var firstImage2 = resampledLandcoverCollectionWAPOR.first();
var projectionWaPOR = firstImage2.projection();
var scaleWaPOR = projectionWaPOR.nominalScale();
print('Nominal scale of the image - resampled Cultivation:', scaleWaPOR);


///////////////////////##################### WaPOR - ACTUAL EVAPOTRANSPIRATION #########
var ETIa = ee.ImageCollection("FAO/WAPOR/2/L1_AETI_D")
    .filterBounds(roi)
    .filter(ee.Filter.date('2018-12-01','2022-02-01')) 
    .map(function(image){return image.clip(roi)})
    .select('L1_AETI_D');

Map.addLayer(ETIa.first(),{min:0, max:80, palette:['white', 'lightgrey', 'grey', 'darkgrey']},'ETa - Dekad Temporal Resolution');

// Function to calculate total ETa for each dekadal period, including scaling by 0.1
var dekadalTotals = function(image) {
  var startDate = ee.Date(image.get('system:time_start'));
  var endDate = ee.Date(image.get('system:time_end'));
    var dekadDays = endDate.difference(startDate, 'day').add(1); // Calculate the number of days in the dekad (difference plus one)
  // First apply the scale factor to convert to actual ETa values, 
  // then multiply by the number of days to get the total ETa for the dekad.
  var scaledET = image.multiply(0.1); // Apply scale factor
  var totalET = scaledET.multiply(dekadDays).toFloat(); // Multiply by the number of days in the dekadal
      totalET = totalET.copyProperties(image, image.propertyNames())
                  .set({
                    'year': startDate.get('year'), 
                    'month': startDate.get('month'),
                    'dekadDays': dekadDays
                  });

        return totalET;
};

// Apply the dekadalTotals function the collection to get scaled ETa values
var AETI_dekadalTotals = ETIa.map(dekadalTotals);

//// Calculate the monthly ETa (Sum of dekadalTotals per month)
var monthly_sum_WAPOR = ee.ImageCollection.fromImages(
  ee.List.sequence(0, 1*36).map(function(n){
var startMonth = ee.Date('2018-12-01').advance(n,'month');
var endMonth = startMonth.advance(1,'month');
return AETI_dekadalTotals
  .filterDate(startMonth, endMonth)
  .reduce(ee.Reducer.sum())
  .select('L1_AETI_D_sum').rename('et250')
  .reproject({
      crs: waporET.projection(),
      scale: waporET.projection().nominalScale()
    })
  .set({
        'system:time_start': startMonth.millis(),
        'system:time_end': endMonth.millis()
      });
  })
);
Map.addLayer(monthly_sum_WAPOR.first(),{min:20, max:200, palette:['white', 'lightgrey', 'grey', 'darkgrey']},'ETa - Monthly Temporal Resolution');

var firstImage3 = monthly_sum_WAPOR.first();
var projection3 = firstImage3.projection();
print('Nominal scale of the image - Monthly ETa - WAPOR:', projection3.nominalScale());


///////////////////////##################### RAINFAL - PRECIPITATION - CHIRPS ##########
//// Import and calculate monthly rainfall from daily rainfall - Function for monthly sum 
var monthly_sum_CHIRPS = ee.List.sequence(0, 1*36).map(function(n){
var start2 = ee.Date('2018-12-01').advance(n,'month');
var end2 = start2.advance(1,'month');
var monthlyPrecip = ee.ImageCollection('UCSB-CHG/CHIRPS/DAILY')
  .filterBounds(roi)
  .filterDate(start2,end2)
  .map(function(image){return image.clip(roi)})
  .sum()
  .select('precipitation')
  .set('system:time_start',start2.millis());
    return monthlyPrecip;
});

var collection_CHIRPS1 = ee.ImageCollection(monthly_sum_CHIRPS)
    .select('precipitation');

// Map.addLayer(collection_CHIRPS1.first().select('precipitation'),{min:80, max:150, palette: ['white', 'lightgreen', 'green', 'darkgreen']},'CHIRPS Monthly Collection');  

// Function to reproject and resample monthly CHIRPS precipitation images to match WaPOR ETa
var collection_CHIRPS = collection_CHIRPS1.map(function(image) {
  var image4 = image.select('precipitation');
  return image4.resample().reproject({
    crs: projectionWaPOR,
    scale: scaleWaPOR
  });
});    

// Map.addLayer(collection_CHIRPS.first().select('precipitation'),{min:80, max:150, palette:['white', 'lightgreen', 'green', 'darkgreen']},'CHIRPS Monthly Collection_Resampled');

var firstImage3 = collection_CHIRPS.first();
var projection3 = firstImage3.projection();
print('Nominal scale of the image - chirpsResampled250:', projection3.nominalScale());

///////////////////////##################### GROUNDWATER STORAGE - DOWNSCALED GLDAS ####
// Function to reproject and resample monthly downscaled GWS to match WaPOR ETa
var sortedCollection = downscaledGWS.sort('system:time_start');
var collection_GWS = sortedCollection.map(function(image) {
  var image5 = image.select('downscaledCorre_M2');
  return image5.resample().reproject({
    crs: projectionWaPOR,
    scale: scaleWaPOR
  });
});    

// Map.addLayer(collection_GWS.first().select('downscaledCorre_M2'),{min:-50, max:50, palette:['white', 'lightblue', 'blue', 'darkblue']},'GWS Collection');

var firstImage5 = collection_GWS.first();
var projection5 = firstImage5.projection();
print('Nominal scale of the image - GWSResampled250:', projection5.nominalScale());






//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
///////////////////////################################################################################################################################################
///////////////////////###################### JOIN IMAGE COLLECTIONS ##################################################################################################
///////////////////////################################################################################################################################################
//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

var classifiedcultivation = collection_Cropland.filterDate('2018-12-01','2021-12-01').filterBounds(roi).select('classification');
var cultivation250mMean = resampledLandcoverCollectionWAPOR.filterDate('2018-12-01','2021-12-01').filterBounds(roi).select('cultivation250Mean');
var et250 = monthly_sum_WAPOR.filterDate('2018-12-01','2021-12-01').filterBounds(roi).select('et250');
var rainfall250 = collection_CHIRPS.filterDate('2018-12-01','2021-12-01').filterBounds(roi).select('precipitation');
var downscaledGWS250 = collection_GWS.filterDate('2018-12-01','2021-12-01').filterBounds(roi).select('downscaledCorre_M2');


// Create a date filter to identify images with matching dates
var dateFilter = ee.Filter.equals({
    leftField: 'system:time_start',
    rightField: 'system:time_start'
  });

var joinType = ee.Join.inner();


var innerJoin1 = ee.ImageCollection(joinType.apply(classifiedcultivation, cultivation250mMean, dateFilter));
var joined1 = innerJoin1.map(function(feature) {
  return ee.Image.cat(feature.get('primary'),feature.get('secondary'));
});

var innerJoin2 = ee.ImageCollection(joinType.apply(joined1, et250, dateFilter));
var joined2 = innerJoin2.map(function(feature) {
  return ee.Image.cat(feature.get('primary'),feature.get('secondary'));
});

var innerJoin3 = ee.ImageCollection(joinType.apply(joined2, rainfall250, dateFilter));
var joined3 = innerJoin3.map(function(feature) {
  return ee.Image.cat(feature.get('primary'),feature.get('secondary'));
});

var innerJoin4 = ee.ImageCollection(joinType.apply(joined3, downscaledGWS250, dateFilter));
var forBalance250m = innerJoin4.map(function(feature) {
  return ee.Image.cat(feature.get('primary'),feature.get('secondary'));
});

function addBALANCE_250(image){
  var selected_bands1 = image.select('precipitation','et250');
  // Calculate 'irrigation' band: ETa - Precipitation and set negative values to 0
  var irrigation = selected_bands1.expression('(b(1) - b(0))').where(selected_bands1.expression('b(1) - b(0)').lt(0), 0); // b(0) = precipitation and b(1) = ETa
  // Calculate 'surface water flux' band: Precipitation - ETa
  var flux = selected_bands1.expression('b(0) - b(1)');
  // Add the new bands to the image and rename them
  return image.addBands(irrigation).addBands(flux).rename('classification','cultivation250Mean','et250','precipitation','downscaledCorre_M2','irrigation','waterflux');
}

//Map function over entire time-series collection
var balance250 = forBalance250m.map(addBALANCE_250);
// Map.addLayer(balance250,{},'Image Collection of Joined Datasets and Balances');  


//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
///////////////////////################################################################################################################################################
///////////////////////########## COMPUTE MONTHLY CROP - AND IRRIGATION WATER USE + SURFACE WATER FLUX ################################################################
///////////////////////################################################################################################################################################
//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
var pixelArea = ee.Image.pixelArea();

var volume_CropWaterUse = balance250.map(function(image) {
  var classMask = image.select('classification').eq(1);  // Select where classification is cultivated
  var cultThreshold = image.select('cultivation250Mean').gt(0.25); // Select where classification covers at least 25% of WaPOR pixel
  var combinedMask = classMask.and(cultThreshold); 
  var masked = image.updateMask(combinedMask); // Only keep pixels which meet both criteria
  var aoi = masked.clip(SDC_Irrigated); 
  var inMeter = aoi.select('et250').multiply(0.001); // mm to m conversion .multiply(0.001);
  var volume = inMeter.multiply(pixelArea); // multiply rate (in meter) by area (in m2) for volume (in m3)
  return volume.set('system:time_start', image.get('system:time_start'));
});

var volume_IrrigationWaterUse = balance250.map(function(image) {
  var classMask = image.select('classification').eq(1);  // Select where classification is cultivated
  var cultThreshold = image.select('cultivation250Mean').gt(0.25); // Select where classification covers at least 25% of WaPOR pixel
  var combinedMask = classMask.and(cultThreshold); 
  var masked = image.updateMask(combinedMask); // Only keep pixels which meet both criteria
  var aoi = masked.clip(SDC_Irrigated); 
  var inMeter = aoi.select('irrigation').multiply(0.001); // mm to m conversion .multiply(0.001);
  var volume = inMeter.multiply(pixelArea); // multiply rate (in meter) by area (in m2) for volume (in m3)
  return volume.set('system:time_start', image.get('system:time_start'));
});

var volume_SurfaceWaterFlux = balance250.map(function(image) {
  var classMask = image.select('classification').eq(1);  // Select where classification is cultivated
  var cultThreshold = image.select('cultivation250Mean').gt(0.25); // Select where classification covers at least 25% of WaPOR pixel
  var combinedMask = classMask.and(cultThreshold); 
  var masked = image.updateMask(combinedMask); // Only keep pixels which meet both criteria
  var aoi = masked.clip(SDC_Irrigated); 
  var inMeter = aoi.select('waterflux').multiply(0.001); // mm to m conversion .multiply(0.001);
  var volume = inMeter.multiply(pixelArea); // multiply rate (in meter) by area (in m2) for volume (in m3)
  return volume.set('system:time_start', image.get('system:time_start'));
});

var volume_GWSchange = balance250.map(function(image) {
  var classMask = image.select('classification').eq(1);  // Select where classification is cultivated
  var cultThreshold = image.select('cultivation250Mean').gt(0.25); // Select where classification covers at least 25% of WaPOR pixel
  var combinedMask = classMask.and(cultThreshold); 
  var masked = image.updateMask(combinedMask); // Only keep pixels which meet both criteria
  var aoi = masked.clip(SDC_Irrigated); 
  var inMeter = aoi.select('downscaledCorre_M2').multiply(0.001); // mm to m conversion .multiply(0.001);
  var volume = inMeter.multiply(pixelArea); // multiply rate (in meter) by area (in m2) for volume (in m3)
  return volume.set('system:time_start', image.get('system:time_start'));
});

// Create a time series chart
var chart_CropWaterUse = ui.Chart.image.series({
  imageCollection: volume_CropWaterUse,
  region: SDC_Irrigated,
  reducer: ee.Reducer.sum(),
  xProperty: 'system:time_start'
}).setOptions({
  title: 'Cumulative Monthly Crop Water Use for Irrigated Cultivation across SDC',
  vAxis: {title: 'Volume (m³)'},
  hAxis: {title: 'Date'},
  lineWidth: 1,
  pointSize: 6,
  series: {0: {color: 'purple'}}, // Adjust color as needed
});
print(chart_CropWaterUse);

// Create a time series chart
var chart_IrrigationWaterUse = ui.Chart.image.series({
  imageCollection: volume_IrrigationWaterUse,
  region: SDC_Irrigated,
  reducer: ee.Reducer.sum(),
  xProperty: 'system:time_start'
}).setOptions({
  title: 'Cumulative Monthly Irrigation Water Use for Irrigated Cultivation (Blue Water) across SDC',
  vAxis: {title: 'Volume (m³)'},
  hAxis: {title: 'Date'},
  lineWidth: 1,
  pointSize: 6,
  series: {0: {color: 'blue'}}, // Adjust color as needed
});
print(chart_IrrigationWaterUse);

// Create a time series chart
var chart_SurfaceWaterFlux = ui.Chart.image.series({
  imageCollection: volume_SurfaceWaterFlux,
  region: SDC_Irrigated,
  reducer: ee.Reducer.sum(),
  xProperty: 'system:time_start'
}).setOptions({
  title: 'Cumulative Monthly Water Flux for Irrigated Cultivation across SDC',
  vAxis: {title: 'Volume (m³)'},
  hAxis: {title: 'Date'},
  lineWidth: 1,
  pointSize: 6,
  series: {0: {color: 'green'}}, // Adjust color as needed
});
print(chart_SurfaceWaterFlux);

// Create a time series chart
var chart_GWSchange = ui.Chart.image.series({
  imageCollection: volume_GWSchange,
  region: SDC_Irrigated,
  reducer: ee.Reducer.sum(),
  xProperty: 'system:time_start'
}).setOptions({
  title: 'Cumulative Monthly change in GWS',
  vAxis: {title: 'Volume (m³)'},
  hAxis: {title: 'Date'},
  lineWidth: 1,
  pointSize: 6,
  series: {0: {color: 'darkblue'}}, // Adjust color as needed
});
print(chart_GWSchange);


//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
///////////////////////################################################################################################################################################
///////////////////////########## COMPUTE ANNUAL IRRIGATION WATER SUMMARY #############################################################################################
///////////////////////################################################################################################################################################
//////%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
var rate_IrrigationWaterUse = balance250.map(function(image) {
  var classMask = image.select('classification').eq(1);  // Select where classification is cultivated
  var cultThreshold = image.select('cultivation250Mean').gt(0.25); // Select where classification covers at least 25% of WaPOR pixel
  var combinedMask = classMask.and(cultThreshold); 
  var masked = image.updateMask(combinedMask); // Only keep pixels which meet both criteria
  var aoi = masked.clip(SDC_Irrigated); 
  var inMeter = aoi.select('irrigation');
  return inMeter.set('system:time_start', image.get('system:time_start'));
});

var IrrigationWaterUse2019 = rate_IrrigationWaterUse.filterDate('2018-12-01','2019-11-30').sum();
var IrrigationWaterUse2020 = rate_IrrigationWaterUse.filterDate('2019-12-01','2020-11-30').sum();
Map.addLayer(IrrigationWaterUse2020,{min:0, max:500, palette:['white', 'lightblue', 'blue', 'darkblue']},'2020- Cumulative Annual Irrigation Use Rate');

// Export.image.toDrive({
//   image: IrrigationWaterUse2020,
//   description: 'date_blueWater_2020', // Set the name for the exported image.
//   folder: 'Google_Drive_Folder_Name', // Specify the folder in your Google Drive where you want to save the image.
//   region: SDC_Irrigated, // Set the region of interest.
//   scale: 20 // Set the scale (in meters per pixel) of the exported image.
// });

////###################################################################################################################################################################
///////////////////////////////////////////////////// LOCALITY ////////////////////////////////////////////////////////////////////////////////////////////////////////
////###################################################################################################################################################################
// Define the styling parameters
var styleBlueParams = {
  color: '0000FF', // blue (0000FF) color for the lines
  fillColor: '00000000', // no fill (transparent)
  width: 1 // width of the lines
};
var styleGreenParams = {
  color: '00FF00', // green (00FF00) color for the lines
  fillColor: '00000000', // no fill (transparent)
  width: 1.5 // width of the lines
};
var styleBlackParams = {
  color: '000000', // black (000000) color for the lines
  fillColor: '00000000', // no fill (transparent)
  width: 1.5 // width of the lines
};
var styledIrrigated = SDC_Irrigated.style(styleBlueParams);// Blue water dependent
//var styledRainfed = SDC_Rainfed.style(styleGreenParams); // Green water dependent
var styledroi = roi.style(styleBlackParams); // Steenkoppies Dolomitic Compartment Study Area
Map.addLayer(styledIrrigated,{}, 'FCB_Irrigated');
//Map.addLayer(styledRainfed,{}, 'FCB_Rainfed');
Map.addLayer(styledroi,{}, 'Steenkoppies Dolomitic Compartment');

Map.centerObject(roi, 11);
Map.setOptions('SATELLITE');