// ****************************************************************************************************************** //
//This script applies the GEE-SedimentBalance Algorithms:
//https://github.com/hydrosolutions/GEE-SedimentBalance

// if export of sediment balance maps fail, it is possible to split the exports in two parts:
// 1. Export of delta h maps by year. See this example: https://code.earthengine.google.com/86c6ce10c98d656f1b0f8f42cee6eeee?noload=true
// 2. Export of sediment balance maps. See this example: https://code.earthengine.google.com/a5a45a95588e8f4916bd0d64b669ccd3?noload=true

// ****************************************************************************************************************** //
// 1. USER INPUTS
// ****************************************************************************************************************** //
//geometries of water bodies
var Lakes = ee.FeatureCollection("projects/ee-adriankreinerak/assets/Lakes_combination_v5");

// define the study period
var start_date = ee.Date.fromYMD(1999,10,1);
var end_date = ee.Date.fromYMD(2022, 10, 1);

//asset path for frequency maps:
var path_Freq = 'projects/ee-adriankreinerak/assets/SurfFreqMaps_Lakes_combination_v2';
//asset path for DEMs:
var path_dem = 'projects/ee-adriankreinerak/assets/DEMs_v2';

var freqMapName='SurfWaterFreq_S2_17to22_';//frequency map export base name
var SedZonesMapFolder='SedZones_1st_Iteration';//sediment balance export image collection name
var SedZonesMapName= 'SedZone_00to22_';//sediment balance map export base name

//sediment balance maps asset path:
var path= 'projects/ee-adriankreinerak/assets/SedZones_Zarafshan/SedZones_1st_Iteration/'; 
var thisname='173003';//ID of the reservoir that will be used as an example for map display

print('Lakes, selected ID',Lakes.filter(ee.Filter.eq('ID_new',ee.Number.parse(thisname))).first());

// parameter cloud cover
var nodata_thresh = 30;
var cc_thresh = 30;
var nodata_thresh_l7 = 50;
//
var cannyThreshold = 0.7;
var cannySigma = 0.7;
var minValue = -0.3;
var maxValue = 0.3;
var scale = 30;
var radius = 20;
var th_veg = 0.5;//NDVI within reservoirs which will be masked from sediment balance map (due to floating vegetation or inundated vegetation that create false shorelines)

// ****************************************************************************************************************** //
// GET LIST OF IDs FOR SPECIFIC DATA TYPES
// ****************************************************************************************************************** //

//get the list of reservoirs for which DEMs are available
var Table = ee.data.listAssets(path_dem);
var l;//number of available DEMs
if (Table.assets===null){
  l=0;
} else {
  l =Table.assets.length;
}var allids=[];
for (var i=0; i<l; i++) {//l
  allids.push(Table.assets[i].id.split('DEMfrom_' + freqMapName)[1]);
}
print('IDs with DEMs',allids);

// New reservoirs_with_DEM, iterate over available IDs (allids): filters Water Bodies, selects the first feature of Water Bodies and adds to new list, converted to new Feature Collection 
// Filters Water Bodies for which DEM is available 
var reservoirs_with_DEM=ee.FeatureCollection(ee.List(ee.List(allids).iterate(function(x,previous){
  return ee.List(previous).add(Lakes.filter(ee.Filter.eq('ID_new',ee.Number.parse(x))));
},ee.List([])))).flatten()

print('Reservoirs with DEMs',reservoirs_with_DEM);
print('Area of Reservoirs',reservoirs_with_DEM.aggregate_array('Lake_area'));

// ****************************************************************************************************************** //
// 2. IMPORT LAYERS AND CODES
// ****************************************************************************************************************** //
// Import external dependencies
var gee_codes = require('users/hydrosolutions/public_functions:SedimentBalance_T1L2');

//Import DEM (of example reservoir)
var dem_thislake = ee.Image(path_dem + "/DEMfrom_"  + freqMapName  + thisname);

//Import a shapefile of the area of interest
var aoi = Lakes.filter(ee.Filter.eq('ID_new', ee.Number.parse(thisname)));
Map.centerObject(aoi);
//keep only the geometry, add a buffer
aoi=aoi.first().geometry()// .buffer(500) //500m buffer around polygon
Map.addLayer(ee.Image().paint(aoi, 0, 2),null, "area_of_interest");
Map.setOptions("HYBRID");

var FreqMap=ee.Image(path_Freq+'/'+ freqMapName + thisname);
var minDEM=ee.Number(dem_thislake.get('min_elev_perc_fit')).getInfo();
var maxDEM=ee.Number(dem_thislake.get('max_elev_perc_fit')).getInfo();
Map.addLayer(FreqMap,{min:0,max:1,palette:'white,blue'}, 'water surface probability S2',false);
Map.addLayer(dem_thislake.updateMask(dem_thislake.lte(maxDEM)).updateMask(dem_thislake.gte(minDEM)), {min:minDEM,max:maxDEM,palette:'white,black'},'DEM',false);

// ****************************************************************************************************************** //
// 3. ENSEMBLE OF LANDSAT SCENES
// ****************************************************************************************************************** //

// number of years
var years = end_date.difference(start_date,'year').round();
print('number of years considered for analysis',years);

var merged_landsat = ee.ImageCollection(gee_codes.get_landsat_images(aoi,start_date,end_date,
  {nodata_thresh: nodata_thresh,
  cc_thresh: cc_thresh,
  nodata_thresh_l7: nodata_thresh_l7,
  }));
// print('Number of scenes in collection:',merged_landsat.size());
// print('merged_landsat',merged_landsat);

// Map.addLayer(ee.Image(merged_landsat.toList(9999).reverse().get(0)), {bands: (['R', 'G', 'B']), min: 0, max: 0.3}, 'RGB Landsat',false);

// ****************************************************************************************************************** //
// 4. SHORELINES OF ALL LANDSAT SCENES
// ****************************************************************************************************************** //

var edge_imgs_landsat = ee.ImageCollection(merged_landsat
    //returns the DEM clipped to the shorelines for each scene
      .map(function(img){return gee_codes.edgeOtsu_optical(img,dem_thislake,
        {cannyThreshold: cannyThreshold,
          cannySigma: cannySigma,
          minimumValue: minValue,
          maximumValue: maxValue,
          scale: scale,
          radius: radius,
          th_veg: th_veg,
          withDEM: true,
        })}));
// print('edge_imgs_landsat',edge_imgs_landsat);
// ****************************************************************************************************************** //
// 5. MEDIAN ELEVATION OF SHORELINES
// ****************************************************************************************************************** //

//map over the images to obtain the median elevation and also the standard deviation of the elevation ('dh_std')
var edge_img4DEM = ee.ImageCollection(gee_codes.get_corrected_elevation_DEM(edge_imgs_landsat, null, {
  filter_outliers: true, //optional: filter out the noisiest shorelines (dh_std >= two standard deviations of the entire ensemble)
  }));
// print('Number of scenes with data available:',edge_img4DEM.size());
// print('edge_img4DEM',edge_img4DEM);

// Map.addLayer(ee.Image(edge_img4DEM.toList(9999).reverse().get(0)), {bands: (['dem_edge']), min: 0, max: 1}, 'dem_edge',false);

// ****************************************************************************************************************** //
// 6. SHORELINE ANOMALIES OF EACH SCENE
// ****************************************************************************************************************** //

////Extract Shoreline Anomalies of individual contours and get the time information required for linear regression
var edge_sed_landsat=ee.ImageCollection(edge_img4DEM.filter(ee.Filter.neq('otsu_th',99)).map(function(img){return gee_codes.get_sedimentation_zone(img,start_date,{radius: 20})}));
//print('edge_sed_landsat',edge_sed_landsat.first());

// ****************************************************************************************************************** //
// 7. SHORELINE ANOMALIES PER YEAR (1)
// ****************************************************************************************************************** //
//returns a stacked image (one image with one band per year)
var annual_deltah=ee.Image(gee_codes.getannual_deltah(edge_sed_landsat,start_date,end_date));


var func4edgeOtsu = function(img){return gee_codes.edgeOtsu_optical(img,dem_thislake,
          {cannyThreshold: cannyThreshold,
            cannySigma: cannySigma,
            minimumValue: minValue,
            maximumValue: maxValue,
            scale: scale,
            radius: radius,
            th_veg: th_veg,
            withDEM: true,
          })};

var func4sed = function(img){return gee_codes.get_sedimentation_zone(img,start_date,{radius: 20})};
function onlyUnique(value, index, array) {
  return array.indexOf(value) === index;
}

var Table = ee.data.listAssets(path);
if (Table.assets===null){
  l=0;
} else {
  l =Table.assets.length;
}
//print('Table.assets',Table.assets);
var tableids=[];
for (var i=0; i<l; i++) {//l
  tableids.push(Table.assets[i].id.split(path + SedZonesMapName)[1].split('_1st_iteration')[0]);
}
tableids=tableids.filter(onlyUnique);
print('IDs with Exported Sediment Balances',tableids);

var missingIDs=[];
for (var i=0; i<allids.length; i++) {//l
  if( tableids.indexOf(allids[i])<0){
    missingIDs.push(allids[i]);
    //print('Missing Lake ID :'+allids[i]);
  }
}

var reservoirs_missingIDs=ee.FeatureCollection(ee.List(ee.List(missingIDs).iterate(function(x,previous){
  return ee.List(previous).add(Lakes.filter(ee.Filter.eq('ID_new',ee.Number.parse(x))));
},ee.List([])))).flatten();

var missingIDs=reservoirs_missingIDs.aggregate_array('ID_new').getInfo();
print('N° of Reservoirs with missing Sediment Balances:',missingIDs.length);
print('missing IDs', missingIDs);
var thisname_example=ee.Number.parse(thisname);

// print('Description:', SedZonesMapName + thisname + '_1st_iteration');
// print('assetID:',path + SedZonesMapName + thisname + '_1st_iteration')

for (var i=0; i<missingIDs.length; i++) {
  var thisname=missingIDs[i].toString();
  var dem_thislake = ee.Image(path_dem + "/DEMfrom_"  + freqMapName  + thisname);
  //Import a shapefile of the area of interest
  var aoi = Lakes.filter(ee.Filter.eq('ID_new',ee.Number.parse(thisname)));
  aoi=aoi.first().geometry();//.buffer(500) 500m buffer around polygon
  var merged_landsat = ee.ImageCollection(gee_codes.get_landsat_images(aoi,start_date,end_date,
    {nodata_thresh: nodata_thresh,
    cc_thresh: cc_thresh,
    nodata_thresh_l7: nodata_thresh_l7,
    }));
  var edge_imgs_landsat = ee.ImageCollection(merged_landsat
      //returns the DEM clipped to the shorelines for each scene
       .map(func4edgeOtsu));
  var edge_img4DEM = ee.ImageCollection(gee_codes.get_corrected_elevation_DEM(edge_imgs_landsat, null, {
    filter_outliers: true, //optional: filter out the noisiest shorelines (dh_std >= two standard deviations of the entire ensemble)
    }));
  var edge_sed_landsat=ee.ImageCollection(edge_img4DEM.filter(ee.Filter.neq('otsu_th',99)).map(func4sed));
  
  // ****************************************************************************************************************** //
  // 8. SHORELINE ANOMALIES PER YEAR (2)
  // ****************************************************************************************************************** //
  //returns a stacked image (one image with one band per year)
  var annual_deltah=ee.Image(gee_codes.getannual_deltah(edge_sed_landsat,start_date,end_date));
//export the stack as an asset
  Export.image.toAsset({
    image:annual_deltah.set('iterations',1).set('ID',ee.Number.parse(thisname)),
    description: SedZonesMapName + thisname,
    assetId: path + SedZonesMapName + thisname,
    scale: 10,
    maxPixels: 1e13,
    region:aoi
  });
}
// print('annual_deltah',annual_deltah);
// Map.addLayer(annual_deltah,{min:0,max:1,palette:'green,red'}, 'annual_deltah',false);

// ****************************************************************************************************************** //
// 9. OPTIONAL: PERFORM ADDITIONAL ITERATIONS
// ****************************************************************************************************************** //

/*

//if already exported, load the annual shoreline anomalies from the asset
//var annual_deltah=ee.Image(path +  SedZonesMapName + thisname);

//NUMBER OF ADDITIONAL ITERATIONS
var it= 1;

var annual_deltah2=ee.Image(ee.List.sequence(1,it).iterate(function(x,previous){
  previous=ee.Image(previous);
  //SEDIMENT BALANCE
  var sed_thislake=ee.Image(gee_codes.stacktoimage(ee.Image(previous)));
  //MEDIAN ELEVATION OF SHORELINES
  var edge_img4DEM = ee.ImageCollection(gee_codes.get_corrected_elevation_DEM(edge_imgs_landsat, sed_thislake, {
    filter_outliers: true, //optional: filter out the noisiest shorelines (dh_std >= two standard deviations of the entire ensemble)
    pvalues:true,
  }));
  //SHORELINE ANOMALIES OF EACH SCENE
  var edge_sed_landsat=ee.ImageCollection(edge_img4DEM.filter(ee.Filter.neq('otsu_th',99)).map(function(img){return gee_codes.get_sedimentation_zone(img,start_date,{radius: 20})}));
  //SHORELINE ANOMALIES PER YEAR
  var annual_deltah=ee.Image(gee_codes.getannual_deltah(edge_sed_landsat,start_date,end_date));

  return annual_deltah;
},annual_deltah));

*/

//Unfortunately, the execution of several iterations at once will fail, due to out of memory. (Error code: 8)
//It is necessary to export after every 1-2 iterations
/* var path2='projects/ee-india-reservoirs/assets/SedZones_Gujarat/SedZones_2nd_Iteration/';
print('annual_deltah2',annual_deltah2);
Export.image.toAsset({
  image:annual_deltah2.set('iterations', 2),
  description: SedZonesMapName + thisname + '_2nd_iteration',
  assetId: path2 +  SedZonesMapName + thisname + '_2nd_iteration',
  scale: 10,
  maxPixels: 1e13,
  region:aoi
});*/

// ****************************************************************************************************************** //
// 10. CALCULATE SEDIMENT BALANCE PER PIXEL WITHIN AOI
// ****************************************************************************************************************** //
//comment these next two lines in case you want to visualize your own results
var annual_deltah=ee.Image('projects/ee-adriankreinerak/assets/SedZones_Zarafshan/SedZones_1st_Iteration/SedZone_00to22_173003_1st_iteration');
aoi=Lakes.filter(ee.Filter.eq('ID_new',thisname_example)).first().geometry();

//Sen's slope of shoreline anomalies per pixel --> sediment balance
var sed_thislake=ee.Image(gee_codes.stacktoimage(annual_deltah));
// print('sed_thislake',sed_thislake)

// ****************************************************************************************************************** //
// 11. OVERALL LAKE SEDIMENT BALANCE
// ****************************************************************************************************************** //  

//convert the stack of shoreline anomalies back to an image collection
var sed_zone_collection=ee.ImageCollection(annual_deltah.bandNames().map(function(b){
  var nr=ee.Number.parse(ee.String(b).slice(1));
  return annual_deltah.select(ee.String(b)).rename('b1').set('t',nr);
}));
// print('sed_zone_collection',sed_zone_collection);
//Count the number of available annual SEAs per pixel (2000 to 2021)
var count=sed_thislake.select('count');
count=count.focal_mode({radius: 2, kernelType :'circle',units: 'pixels',iterations: 1}).blend(count)
.reproject({crs: 'EPSG:32642', scale:10});

var pvalue=sed_thislake.select('p-value');
pvalue=pvalue.focal_mean({radius: 2, kernelType :'circle',units: 'pixels',iterations: 1}).blend(pvalue)
.reproject({crs: 'EPSG:32642', scale:10});

Map.addLayer(count,{bands: (['count']),palette:['white','blue'],min: 0, max: 20}, "count filled",false);
Map.addLayer(pvalue,{bands: (['p-value']),palette:['white','blue'],min: 0, max: 0.1}, "pvalue filled",false);

var slope_final=
    sed_thislake.select('slope').updateMask(count.gte(6)).focal_mean({radius: 2, kernelType :'circle',units: 'pixels',iterations: 1})
  .blend(
    //focal_mean with radius 1 to smooth the image
    sed_thislake.select('slope').updateMask(count.gte(6)).focal_mean({radius: 1, kernelType :'circle',units: 'pixels',iterations: 1}).updateMask(count.gte(6))
  )
  .reproject({crs: 'EPSG:32642', scale:10});
//Because surface water may not be correctly identified under canopies, mask all boundary pixels that are adjacent to lush vegetation
//var ndvi_wetseason=ee.Image(1).where(merged_landsat.filter(ee.Filter.calendarRange(7,9, 'month')).mean().select('NDVI').gte(th_veg),0);
slope_final=slope_final
  //.updateMask(ndvi_wetseason.focal_min({radius: 10,kernelType:'square', units: 'meters'}))
  .updateMask(count.gte(1))//mask out all pixels that have never represented the shoreline
  .reproject({crs: 'EPSG:32642', scale:10});

Map.addLayer(slope_final,{bands: (['slope']),palette:['red', 'white','blue'],min: -0.05, max: 0.05}, "slope filled, raw",false);
//use this one for additional iterations:
Map.addLayer(slope_final.updateMask(sed_thislake.select('sedzone_filled').neq(0)),{bands: (['slope']),palette:['red', 'white','blue'],min: -0.05, max: 0.05}, "slope filled");
//Map.addLayer(slope_final.updateMask(pvalue.lte(0.1)),{bands: (['slope']),palette:['red', 'white','blue'],min: -0.05, max: 0.05}, "slope filled, pvalue <= 0.1",false);

var dem_filled = ee.Image(-99).rename('b1').where(dem_thislake.gt(0),dem_thislake).clip(aoi);

print('average sediment balance in mm (negative: deposition, positive: erosion',
  ee.Number(slope_final
  .reduceRegion(ee.Reducer.mean(),aoi).get('slope')).multiply(1000).round());

print('how much storage volume is lost each year, mio. m3',
  ee.Number(slope_final
  .reduceRegion(ee.Reducer.mean(),aoi).get('slope')).multiply(1000).round().multiply(-1)
  .divide(1000)
  .multiply(ee.Number(aoi.area())));

print('average sediment balance in mm where there is deposition',
  ee.Number(slope_final.updateMask(slope_final.lte(-0.01))
.reduceRegion(ee.Reducer.mean(),aoi).get('slope')).multiply(1000).round());

print('average sediment balance in mm where there is erosion',
  ee.Number(slope_final.updateMask(slope_final.gte(0.01))
.reduceRegion(ee.Reducer.mean(),aoi).get('slope')).multiply(1000).round());


print('neutral area %',
  ee.Number(slope_final.updateMask(slope_final.abs().lt(0.01))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope')) 
  .divide(ee.Number(slope_final.reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));

//count=sed_thislake.select('count');

print('erosion area %',
  ee.Number(slope_final.updateMask(slope_final.gte(0.01))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope'))
  .divide(ee.Number(slope_final.reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));
print('deposition area %',
  ee.Number(slope_final.updateMask(slope_final.lte(-0.01))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope'))
  .divide(ee.Number(slope_final.reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));
  
print('masked area %',
  ee.Number(slope_final.updateMask(slope_final.abs().gte(0.01)).updateMask(count.gte(3))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope'))
  .divide(ee.Number(slope_final.updateMask(count.gte(3)).reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));  

print('erosion area %, only significant slopes',
  ee.Number(slope_final.updateMask(slope_final.gte(0.01)).updateMask(sed_thislake.select('sedzone_filled').neq(0)).updateMask(count.gte(3))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope'))
  .divide(ee.Number(slope_final.updateMask(count.gte(3)).reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));
print('deposition area %, only significant slopes',
  ee.Number(slope_final.updateMask(slope_final.lte(-0.01)).updateMask(sed_thislake.select('sedzone_filled').neq(0)).updateMask(count.gte(3))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope'))
  .divide(ee.Number(slope_final.updateMask(count.gte(3)).reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));

print('masked area %, only significant slopes',
  ee.Number(slope_final.updateMask(slope_final.abs().gte(0.01)).updateMask(sed_thislake.select('sedzone_filled').neq(0)).updateMask(count.gte(3))
  .reduceRegion(ee.Reducer.count(),aoi).get('slope'))
  .divide(ee.Number(slope_final.updateMask(count.gte(3)).reduceRegion(ee.Reducer.count(),aoi).get('slope')))
  .multiply(100));  
