// (c) Rebecca Dell, University of Cambridge and Ruthie Halberstadt, University of Massachusetts (September 2020)
// Example workflow for applying k-means and extracting training data.

// This script groups the pixels within the training images into spectrally similar clusters.
//As the user you interpret these clusters, and form your desired training classes. 
//At the end of this script these new training classes are sampled, and exported as Earth Engine assets 
//(saving the reflectance values and user-interpreted class as a table), for use in subsequent scripts. 

// In your imports you should import your training images from your assets, as pre-processed in the Landsat 8
// pre-processing script. You should also import shapefiles for your ice shelf regions of interest. 

// Set image visualisation parameters (these are set to deal with L8 TOA images)
var vizParamsL8 = {bands: ['B4', 'B3', 'B2'],
  min: 0,
  max: 1,
  gamma: [0.95, 1.1, 1]
};

// Write masking function, which can later be used to manually mask areas of error that you don't want to 
// carry forward into the training data 
var maskInside = function(image, geometry) {
  var mask = ee.Image.constant(1).clip(geometry).mask().not();
  return image.updateMask(mask);
};
 
// Call external module to filter pixels by an NDWI threshold of 0.1
var extra_filter = require('users/beckyld/Chapter_2:Modules/L8_Filter');

// Select the bands we want to use in the K-Means Clusterer
var bands = ['B1','B2','B3','B4','B5','B6','B7'];

// (De)select lines as appropriate to zoom to each ice shelf 
// Focus Viewer on Ice Shelf Nivlisen
Map.setCenter(11.3266, -70.3630,8);
//Focus Viewer on Ice Shelf Roi B 
//Map.setCenter(30.985, -69.727,8);
// Focus Viewer on GVI and Bach
//Map.setCenter(-67.554262, -71.425088,6);
// Focus Viewer on Amery
//Map.setCenter(70.35693446515722, -70.380905933342, 5);
// Focus Viewer on Shackleton
//Map.setCenter(100, -66, 6);
// Focus Viewer on Nansen
//Map.setCenter(163, -74, 6);

// Extract the shapefiles for each ice shelf from the SCAR Antarctic Digital Database, Version 7.0, in Quantarctica
var GVI = (ice_shelves.filter(ee.Filter.eq("system:index", '0000000000000000045e')));
var Amery = (ice_shelves.filter(ee.Filter.eq("system:index", '00000000000000000460')));
var Shackleton = (ice_shelves.filter(ee.Filter.eq("system:index", '000000000000000003f5')));
var Bach = (ice_shelves.filter(ee.Filter.eq("system:index", '0000000000000000044b')));
var Nansen = (ice_shelves.filter(ee.Filter.eq("system:index", '00000000000000000367'))); 
var Nivlisen = (ice_shelves.filter(ee.Filter.eq("system:index", '0000000000000000052d'))); 
var RoiB = (ice_shelves.filter(ee.Filter.eq("system:index", '00000000000000000532'))); 

// Make an image collection from the pre-processed training images, which you have loaded in from assets 
// Images are pre-processed in 'Landsat_8_preprocessing'
var image_collection = ee.ImageCollection([
image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9, image_10, image_11, image_12
]);
print(image_collection); 
//Map.addLayer(image_collection, vizParamsL8);

// Make a feature collection of areas over each image that you want to post-process/ clip out. These polygons are drawn manually in the map viewer based 
// on areas of mis-classification or potential error in an attempt to prevent carrying any error forward into the Random Forest Classifier. 
var cloudpoly = ee.FeatureCollection([IM0_GVI,IM1_GVI,IM2_Amery,IM3_Amery,IM4_Shackleton,IM5_Shackleton,IM6_Nansen,IM7_Nansen,IM8_Nivlisen,
IM9_Nivlisen,IM10_RoiB,IM11_RoiB]);
var polyList = cloudpoly.toList(cloudpoly.size());
//print(polyList);

// Add the full base image to the map viewer for context. I do this here as once I have applied the NDWI filter the image is heavily filtered and harder 
// to contextualise. 
var image2export = ee.Image([image_collection.toList(12).get(8)]);
// print(image2export);
Map.addLayer(image2export,vizParamsL8, 'base image');

// Apply the NDWI filter of 0.1 to the collection - to try and find only water. This gets rid of a lot of confusion in the k-means clusterer
var image_collection_ndwi = image_collection.map(extra_filter.extra_filter);
//Map.addLayer(image_collection_ndwi, vizParamsL8);

//MAKE THE K-MEANS CLUSTERER
//Sample pixels from all images to train one kmeans clusterer

// Make an empty feature collection
var sampleKmeans = ee.FeatureCollection([]); 

//Loop through each image in the training collection
for (var i=0; i<12; i=i+1){ 
  var image2cluster = ee.Image([image_collection_ndwi.toList(12).get(i)]);   // Select an image to cluster
  var kimg = image2cluster.select(bands); // select bands 1-7
  var ktraining = kimg.sample({region:kimg.geometry(),scale:30,numPixels:100000}); // sample pixels to feed to the clustering algorithm
sampleKmeans = sampleKmeans.merge(ee.FeatureCollection([ktraining])); // Add the sampled pixels to the feature collection created above this for loop. Merge interatively addsthe pixels as the 
// for-loop progresses .
} //LOOP ENDS HERE
// Here we flatten the merged feature collections into a single feature collection
sampleKmeans = sampleKmeans.flatten(); // print(sampleKmeans.size())
var clusterer = ee.Clusterer.wekaXMeans({minClusters:5, maxClusters:70}).train({features:sampleKmeans,inputProperties:bands}); // train the clustering algorithm

// Subdivide slush and blue ice (classes 4-8)
var sampleKmeanslush = ee.FeatureCollection([]); // Make an empty feature collection
for (var i=0; i<12; i=i+1){ //Loop through each image in the training collection
  var kimg = ee.Image([image_collection_ndwi.toList(12).get(i)]).select(bands); // Select image to cluster and bands
  var kresult = kimg.cluster(clusterer); // Apply the initial K-Means clusterer to the selected image
  var kimg2 = kimg.mask(kresult.gte(4).and(kresult.lte(8))); // Select only clusters 4-8, mask out everything else 
  var ktraining2 = kimg2.sample({region:kimg2.geometry(),scale:30,numPixels:10000}); // sample pixels to feed to the sub-sampling clustering algorithm
  // Note that I sample fewer pixels here ^^^
  // Add the sampled pixels to the feature collection created above this for loop. Merge interatively addsthe pixels as the 
  // for-loop progresses
sampleKmeanslush = sampleKmeanslush.merge(ee.FeatureCollection([ktraining2]));
} // LOOP ENDS HERE
// Here we flatten the merged feature collections into a single feature collection
sampleKmeanslush = sampleKmeanslush.flatten(); 
var clusterer2 = ee.Clusterer.wekaXMeans({minClusters:8, maxClusters:12}).train({features:sampleKmeanslush,inputProperties:bands}); // train the clustering algorithm

// Create empty image collections, which we will store the final training classes in. 
var lakes = ee.ImageCollection([]);
var slushs = ee.ImageCollection([]);
var blues = ee.ImageCollection([]);

// LOOP THROUGH ALL IMAGES AND ASSESS RESULTS 
for (var i=8; i<9; i=i+1){ 
var image2cluster = ee.Image([image_collection_ndwi.toList(12).get(i)]); // select a single image at a time 

// Apply the trained kmeans clusterer to all images, one-by-one, to identify the different clusters (see NOTES below)
//Map.addLayer(image2cluster,vizParamsL8,'img'+i); // adds the image to the map and names it. 
var kimg = ee.Image([image2cluster]).select(bands); // Selects just bands 1-7 from the image
var kresult = kimg.cluster(clusterer); // Applies the trained clusterer to the image. Output = image sorted by kmeans clusters.
Map.addLayer(kresult.randomVisualizer(),{},'kmeans',false); // false command means layer is initially switched off in map-viewer
// I think the below line prints the number of clusters found by k-means
//print('kmeans:' +i,kresult.reduceRegion({reducer:ee.Reducer.max(),geometry:kresult.geometry(),scale:30,maxPixels:1e9}));

var kimg2 = kimg.mask(kresult.gte(4).and(kresult.lte(8))); // Selects just clusters 4-8 for the subdivided section
var kresult2 = kimg2.cluster(clusterer2); // Applies clusterer 2 (the subdivided clusterer) to the image
Map.addLayer(kresult2.randomVisualizer(),{},'kmeans2',false); // Displays as layer in map viewer 

// I think the below line prints the number of clusters found by k-means2
//print('kmeans2:' +i,kresult2.reduceRegion({reducer:ee.Reducer.max(),geometry:kresult2.geometry(),scale:30,maxPixels:1e9}));

// Select clusters that best represent (i) Open Water, (ii) Slush, (iii) Blue Ice. Outputs from this = images.
 var lake = kimg.mask(kresult.gte(0).and(kresult.lte(3)));
 var lake_2 = kimg.mask(kresult.gte(9).and(kresult.lte(14)));
 var slush = kimg.mask(kresult.gte(15).and(kresult.lte(20)));
 var slush_2 = kimg2.mask(kresult2.gte(4).and(kresult2.lte(7))); // N.B these come from kimg2 (the subdivided clusterer)
 var blue = kimg2.mask(kresult2.gte(0).and(kresult2.lte(3))); // N.B these come from kimg2 (the subdivided clusterer)

// Add the newly formed classes to the map
//Map.addLayer(lake.select('B1'), {palette:'800bb1'}, 'lake'+i,false);
//Map.addLayer(lake_2.select('B1'), {palette:'800bb1'}, 'lake_2'+i,false);
//Map.addLayer(slush.select('B1'), {palette:'756bb1'}, 'slush'+i,false);
//Map.addLayer(slush_2.select('B1'), {palette:'756bb1'}, 'slush_2'+i,false);
//Map.addLayer(blue.select('B1'), {palette:'00ffff'}, 'blue'+i,false);

//Clip to manual post-processing polygons (which were arranged into a feature collection above.)
var poly = ee.Feature(polyList.get(i)); // selects feature collection i in poly-list. They are arranged in the correct order so to match the 
// images that they are intended for (see lines 56-57).
var lake_masked = maskInside(lake, poly); // use maskInside(lines8-12) function to mask each class for each image 
var lake_2_masked = maskInside(lake_2, poly);
var slush_masked = maskInside(slush, poly);
var slush_2_masked = maskInside(slush_2, poly);
var blue_masked = maskInside(blue, poly);

// This mosaics together all lake classes, slush classes etc and clips them to the geometry. It then merges these classes image by image.
// Merge is 'better' than mosaic for this purpose, as it retains all information. Mosaic is ok when combining lake classes from the same image as
// they are not overlapping, so no data will be lost.
lakes = lakes.merge(ee.ImageCollection([lake_masked, lake_2_masked]).mosaic().clip(kresult.geometry()));
slushs = slushs.merge(ee.ImageCollection([slush_masked, slush_2_masked]).mosaic().clip(kresult.geometry()));
blues = blues.merge(ee.ImageCollection([blue_masked]).mosaic().clip(kresult.geometry()));

Map.addLayer(slushs.select('B1'), {palette:'756bb1'}, 'slushs'+i,false);
Map.addLayer(lakes.select('B1'), {palette:'800bb1'}, 'lakes'+i,false);
//Map.addLayer(blues.select('B1'), {palette:'a4eaee'}, 'blues'+i,false);
} // LOOP ENDS HERE

// Here we repeat the KMeans clustering process to get some additional classes that were not included in the first training images

var dirtyimage_collection = ee.ImageCollection([
dirty_image_1, dirty_image_2 // these are images of Nansen Ice Shelf (see manuscript)
]);
//print(dirtyimage_collection); 

var dirtypolyList = dirtycloudpolys.toList(dirtycloudpolys.size()); // Dirty cloud polys can be drawn in a different GEE map viewer and imported to this script, or drawn in this script - makes no difference
//print(dirtypolyList);

// Add the full base image to the map viewer for context
var dirtyimage2export = ee.Image([dirtyimage_collection.toList(2).get(1)]);
// print(image2export);
//Map.addLayer(dirtyimage2export,vizParamsL8, 'dirty base image');

// Apply the NDWI filter of 0.1 to the collection - to try and find only water. This gets rid of a lot of confusion
var dirtyimage_collection_ndwi = dirtyimage_collection.map(extra_filter.extra_filter);
//Map.addLayer(image_collection_ndwi, vizParamsL8);

//MAKE THE EXTRA DIRTY ICE K-MEANS CLUSTERER
//Sample pixels from all images to train one kmeans clusterer

var dirtysampleKmeans = ee.FeatureCollection([]); 
for (var i=0; i<2; i=i+1){ 
  
  var dirtyimage2cluster = ee.Image([dirtyimage_collection_ndwi.toList(2).get(i)]);
  var dirtykimg = dirtyimage2cluster.select(bands); 
  var dirtyktraining = dirtykimg.sample({region:dirtykimg.geometry(),scale:30,numPixels:100000}); // sample pixels to feed to the clustering algorithm
dirtysampleKmeans = dirtysampleKmeans.merge(ee.FeatureCollection([dirtyktraining]));   
}

dirtysampleKmeans = dirtysampleKmeans.flatten(); // print(sampleKmeans.size())
var dirtyclusterer = ee.Clusterer.wekaXMeans({minClusters:5, maxClusters:70}).train({features:dirtysampleKmeans,inputProperties:bands}); // train the clustering algorithm

var dirtys = ee.ImageCollection([]);
var d_blues = ee.ImageCollection([]);
var snows = ee.ImageCollection([]);
var extras = ee.ImageCollection([]); // This captures cluster 18 which looks like a blue ice snow/ mix. Doesn't really matter as is just 'not_wet'.

// Loop thru all dirty images to assess results all together
for (var i=0; i<2; i=i+1){ 
var dirtyimage2cluster = ee.Image([dirtyimage_collection_ndwi.toList(2).get(i)]);

// Apply the trained kmeans clusterer to all images, one-by-one, to identify the slush clusters (see NOTES below)
//Map.addLayer(dirtyimage2cluster,vizParamsL8,'img'+i);
var dirtykimg = ee.Image([dirtyimage2cluster]).select(bands); 
var dirtykresult = dirtykimg.cluster(dirtyclusterer);
//Map.addLayer(dirtykresult.randomVisualizer(),{},'kmeans',false);

//print('dirtykmeans:' +i,dirtykresult.reduceRegion({reducer:ee.Reducer.max(),geometry:dirtykresult.geometry(),scale:30,maxPixels:1e9}));

// Plot Becky's Categories for (i) Slush and Shallow Water, (ii) Deep Water, (iii) Blue Ice
var dirty_ice = dirtykimg.mask(dirtykresult.gte(2).and(dirtykresult.lte(3)));
//var lake = kimg.mask(kresult.gte(5).and(kresult.lte(9))); NOT NEEDED
var d_blue_ice = dirtykimg.mask(dirtykresult.gte(12).and(dirtykresult.lte(15)));
var extra = dirtykimg.mask(dirtykresult.eq(18));
var snow = dirtykimg.mask(dirtykresult.gte(19).and(dirtykresult.lte(21)));

//Map.addLayer(dirty_ice.select('B1'), {palette:'bc8f8f'}, 'dirty_ice'+i,false);
//Map.addLayer(lake.select('B1'), {palette:'800bb1'}, 'lake'+i,false); NOT NEEDED
//Map.addLayer(d_blue_ice.select('B1'), {palette:'40e0d0'}, 'd_blue_ice'+i,false);
//Map.addLayer(extra.select('B1'), {palette:'40e0d0'}, 'extra'+i,false);
//Map.addLayer(snow.select('B1'), {palette:'ff1493'}, 'snow'+i,false);

//Clip to manual post-processing polygons
var dirtypoly = ee.Feature(dirtypolyList.get(i));
var d_blue_ice_masked = maskInside(d_blue_ice, dirtypoly);
var extra_masked = maskInside(extra, dirtypoly);
var dirty_ice_masked = maskInside(dirty_ice, dirtypoly);
var snow_masked = maskInside(snow, dirtypoly);

d_blues = d_blues.merge(ee.ImageCollection([d_blue_ice_masked]).mosaic().clip(dirtykresult.geometry()));
dirtys = dirtys.merge(ee.ImageCollection([dirty_ice_masked]).mosaic().clip(dirtykresult.geometry()));
snows = snows.merge(ee.ImageCollection([snow_masked]).mosaic().clip(dirtykresult.geometry()));
extras = extras.merge(ee.ImageCollection([extra_masked]).mosaic().clip(dirtykresult.geometry()));

/*Map.addLayer(d_blues.select('B1'), {palette:'a4eaee'}, 'blues'+i,false);
Map.addLayer(dirtys.select('B1'), {palette:'bc8f8f'}, 'dirtys'+i,false);
Map.addLayer(snows.select('B1'), {palette:'ff1493'}, 'snow'+i,false);*/
 
}

// extract 1000 points for each class from the main clusterer (3 classes)
var lake1sampled = lakes.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:27700})});//1001 pixels
  var trainlake1 = lake1sampled.flatten().map(function(f){return f.set('class',1)}); 
  //print('trainlake1',trainlake1.size());
var slushsampled = slushs.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:16500})});//1015 pixels
  var trainslush1 = slushsampled.flatten().map(function(f){return f.set('class',2)});
 // print('trainslush1',trainslush1.size());
var bluesampled = blues.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:32800})});//1008 pixels
  var trainblue1 = bluesampled.flatten().map(function(f){return f.set('class',3)});
  //print('trainblue1',trainblue1.size());
  
// extract 1000 points for each class from the dirty clusterer (4 classes)
var d_bluesampled = d_blues.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:27000})});//1095 pixels
  var d_trainblue1 = d_bluesampled.flatten().map(function(f){return f.set('class',4)});
  //print('d_trainblue1',d_trainblue1.size());
var dirtysampled = dirtys.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:466600})});//1000 pixels 
  var traindirty1 = dirtysampled.flatten().map(function(f){return f.set('class',5)});
  //print('traindirty1',traindirty1.size());
var snowsampled = snows.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:102000})});// 1009 pixels
  var trainsnow1 = snowsampled.flatten().map(function(f){return f.set('class',6)});
  //print('trainsnow1',trainsnow1.size());
var extrasampled = extras.map(function(f){return f.sample({region:f.geometry(),scale:30,numPixels:107000})});//1000 pixels
  var trainextra1 = extrasampled.flatten().map(function(f){return f.set('class',7)});
  //print('trainextra1',trainextra1.size());

var training_all = trainlake1.merge(trainslush1).merge(trainblue1).merge(d_trainblue1).merge(traindirty1).merge(trainsnow1).merge(trainextra1);  

Export.table.toAsset({
  collection: training_all,
  description:'training_data',
});