tailingsClassifier

/*********************************************
*                                            *
*        Tailings Classification Model       *
*                                            *
* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ *
*                                            *
*             Author: Dan Jewell             *
*       Saints Mary's University, 2022       *
*                                            *
*********************************************/

/* This model is intended to sample Sentinel-2 surface reflectance imagery and classify it to indicate waste from
historic gold mines. Model is trained on data from the NS Mine Tailings Database.

/*********************************************
*                                            * 
*    Initialize Variables and Load Assets    *
*                                            *
*********************************************/

//////////////////////////////////////////////
//             User Parameters              //
//////////////////////////////////////////////

//These parameters will determine behaviour of model

//How many times sampler/classifier will loop
var iterations = 5;

//Change scale if changing bands. Should only be 10 or 20.
var scale = 20;

//Total sample points for non-tailings. Generally, more = better precision, fewer = better recall
//Tailings samples are collected as proportion of tailings pixels
var totalSamplePoints = 7500;

//Determine whether indices are used in classifier. Indices will still be used for filtering training data
var includeIndices = true;

//Determine whether spectral bands are used
var includeBands = true;

// Use best bands based on testing. 
// Overrides includeIndices and includeBands
var useBestBands = true;

//Use random districts, or the target districts below
//Randomizes for EACH iteration - monte carlo cross-validation (see Lyons et al, 2018)
var randomDistricts = false;
//List of districts that will be used if not random
var targetDistrictNames = ee.List(["Montague"]);

//Use custom-drawn polygon
var userTargetArea = false;

//Adjust the precision/recall weight for F Score.
//1 = equal weight. 2 = recall is twice as important as precision
//Must be in range [0,2]
var b = 1; 

// Thresholds for masking vegetation and water in pixels
var vegetationThreshold = 0.6;
var waterThreshold = 0.0;

// Parameter Dictionary for Records
var parameterDictionary = ee.Dictionary({
  "iterations": iterations,
  "scale": scale,
  "nSamplePoints": totalSamplePoints,
  "randomDistricts": randomDistricts,
  "usingBestBands": useBestBands,
  "vegetationThreshold": vegetationThreshold,
  "waterThreshold": waterThreshold,
  "fScoreBValue": b
});

print(parameterDictionary, "Model Parameters");

//////////////////////////////////////////////
//             Misc. Variables              //
//////////////////////////////////////////////
  
//Image to be used for training (Halifax area)
// var trainingImage = trainingImageAsset;

var trainingImage = s2sr
  .filterBounds(ROI)
  .filterDate("2019-07-01", "2019-09-01")
  .sort("CLOUDY_PIXEL_PERCENTAGE")
  .first();

//Mainland NS polygon to cut out large amounts of ocean
var mainlandNS = ee.Feature(mainlandNS);

//Clip training image to mainland
trainingImage = trainingImage.clip(mainlandNS);

//Radius of buffered circle around districts in meters
var districtBufferDistance = 2500;

//A list of seeds for various server-side uses
// var seedList = ee.List.sequence({start: 0, count: 1000}).shuffle({seed: ee.Number(Math.random() * 1000).floor()});
var seedList = ee.List.sequence({start: 0, count: 1000}).shuffle({seed: 5283});

//////////////////////////////////////////////
//                Band Sets                 //
//////////////////////////////////////////////

//All bands with 10m resolution plus SCL (land classification band)
var bands10 = ["B2", "B3", "B4", "B8"];

//All bands with 20m resolution plus SCL
var bands20 = ["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12"];

// Based on testing, the best-performing set of variables
var bestBands = ["B3", "B6", "B7", "B8A", "B12", "ndvi", "mndwi"];

//////////////////////////////////////////////
//                Cloud Mask                //
//////////////////////////////////////////////

// Make sure that clouds are indicated when switching training images 

//Clouds look similar spectrally to concrete/urban, which can look like tailings. 
//Automatic cloud removal is difficult and risks removing tailing as well

//Display cloudiness of whole image
//Unless clouds can be separated from tailings more reliably, simply displaying cloudiness is best
var cloudChance = ee.Number(trainingImage.get("CLOUDY_PIXEL_PERCENTAGE"));
print("The % of cloudy pixels in this image:", cloudChance);

//////////////////////////////////////////////
//      Projection and Band Selection       //
//////////////////////////////////////////////

//A message on iterations
// print("Number of iterations:", iterations);
//A Message on FScore b value
// print("B Value of F-Score:", b);

//Set automatically based on parameters above:

//Projection for 10m bands
var projection10 = trainingImage.select("B4").projection();

//Projection for 20m bands
var projection20 = trainingImage.select("B8A").projection();

//Bands and projection are set to scale
if(scale == 10){
  var bands = bands10;
  var projection = projection10.atScale(scale);
} else if(scale == 20){
  var bands = bands20;
  var projection = projection20.atScale(scale);
}

//Add SCL band to bands list for sampling later
bands = ee.List(bands).add("SCL");

//Set training image with relevant bands
trainingImage = trainingImage.select(bands);

/*********************************************
*                                            * 
*           Initial Data Processesing        *
*                                            *
*********************************************/

//////////////////////////////////////////////
//             Tailings Areas               //
//////////////////////////////////////////////

/* Original data from the NS Mine Tailings Database (from E. W. Hennick and J. C. Poole, 2020).
 * This data is filtered to show only tailings areas from gold mines in mainland NS. */

// Cast imported asset to a feature collection 
var tailings = ee.FeatureCollection(tailingsNS);

// Clip to only mainland NS
tailings = tailings.map(function(f){
  return f.intersection(mainlandNS);
});

// Filter to only tailings from gold mines
tailings = tailings.filter(ee.Filter.and(
  ee.Filter.eq("Commodity", "Gold"), 
  ee.Filter.eq("GCODE", "TAILINGS_AREA"), 
  ee.Filter.neq("Crusher1", "0"))); //Crushers with no associated tailings are labeled "Crusher1 = 0"
  
// print(tailings, "Tailings Feature Collection");
Map.addLayer(tailings.draw({color: "Grey", strokeWidth: 1}), null, "Tailings Database Polygons");

// Convert tailings features to a raster image
function CreateTailingsImage(tailings, image){
  
  // Create a background image with value of 0 and reduce its resolution to match lowest of input bands
  var background = ee.Image(0)
    .reproject(projection)
    .clip(image.geometry()); //Created background image is boundless. Clip to input (training image) extent
  
  // Give each tailings feature a property, "tailings", with value = 1
  var tailingsWithProperty = tailings.map(function(f){return f.set("tailings", 1);});
  
  // Convert tailings features to an image (raster)
  var tailingsImage = tailingsWithProperty.reduceToImage(["tailings"], ee.Reducer.max()) //Reducer max takes highest value (0 or 1)
    .reproject(projection)
    .reduceResolution(ee.Reducer.mean());
  
  // Where tailings image (excluding No Data) is not 0, replace background pixels with value of 1
  return background.where(tailingsImage.eq(1), tailingsImage)
    .reproject(projection) 
    .rename("tailings")
    .eq(1); //Quickly changes range to [0,1]
}

var tailingsImage = CreateTailingsImage(tailings, trainingImage);

// print(tailingsImage, "Tailings Image")
  // Map.addLayer(tailingsImage)

//Rename "DistrictNa" property of districtPoints
districtPoints = districtPoints.map(function(f){
  
  var name = f.get("DistrictNa");
  
  return ee.Feature(f.geometry(), {"ShortName": name}).copyProperties({source: f, exclude: ["DistrictNa", "Watershed"]});
});

//////////////////////////////////////////////
//        Training/Validation Image         //
//////////////////////////////////////////////

/*  Tailings image is filtered to remove pixels that indicate vegetation or open water. 
 *  This is to simulate the surfaces that are visible to the Sentinel-2 sensors and
 *  may be classified as tailings.
 *  NDVI and NDWI are also used as inputs to the classifier (when includeIndices is true)
 */

//Thresholds for masking
//Set above in model parameters
// var vegetationThreshold = 0.5;
// var waterThreshold = 0.0; //Currently using s2's SCL layer for water

//Create normalized difference images for masking
var ndvi = trainingImage.normalizedDifference(["B8", "B4"]).rename("ndvi"); //Normalized Difference Vegetation Index
var ndwi = trainingImage.normalizedDifference(["B3", "B8"]).rename("ndwi"); //Normalized Difference Water Index
var mndwi = trainingImage.expression("mndwi = (b('B3') - b('B11')) / (b('B3') + b('B11'))"); // Modified Normalized Difference Water Index

// Histogram to show distribution of vegetation and water index values
var ndviSample = ndvi.sample({numPixels: 10000, dropNulls: true});
var ndviHisto= ui.Chart.feature.histogram(ndviSample, "ndvi");
// print(ndviHisto);

var mndwiSample = mndwi.sample({numPixels: 10000, dropNulls: true});
var mndwiHisto= ui.Chart.feature.histogram(mndwiSample, "mndwi");
// print(mndwiHisto);

//Water layer of S2's SCL band
var water = trainingImage.select("SCL").eq(6);

//Create mask images. Any pixel with value under threshold will return with value of 1
var vegetationMask = ndvi.lte(vegetationThreshold).reproject(projection).rename("vegetationMask");
var waterMask = mndwi.lte(waterThreshold).reproject(projection).rename("waterMask");
// var waterMask = water.eq(0).reproject(projection);

function ApplyMasks(imageToMask, vegetationMask, waterMask){
  
  return imageToMask.multiply(vegetationMask).multiply(waterMask);
  
}

var maskedTailings = ApplyMasks(tailingsImage, vegetationMask, waterMask).reproject({crs: projection, scale: scale});

//Masks for visualization:
Map.addLayer(vegetationMask, {min: 0, max: 1, palette: ["green", "brown"]}, "Vegetation Mask", false);
Map.addLayer(waterMask.eq(0).selfMask(), {min: 0, max: 1, palette: ["blue", "brown"]}, "Water Mask", false);

//Images used for vegetation mask:
Map.addLayer(ndvi, {palette: ["brown", "green"], min: -1, max: 1}, "NDVI", false);

//Image used for water mask:
Map.addLayer(mndwi, {palette: ["brown", "blue"]}, "MNDWI", false);
// Map.addLayer(water.selfMask(), {palette: ["white", "blue"]}, "Water Mask from SCL")

/*********************************************
*                                            * 
*               Multi-Use Tools              *
*                                            *
*********************************************/
//General-use functions that are used multiple times

//Counts pixels in an image and divides by a given quotient
function GetSampleCount(image, region, factor){
  
  image = ee.Image(image);
  //Convert to vector
  var imgSamples = image.sample({
    region: region,
    scale: scale,
    projection: projection, 
    factor: factor,
    dropNulls: true});
  
  return imgSamples.size().floor();
}

//Classify points
//This should be used for the sample points that are selected after filtering training image for viable training areas
//Returns one classified image with two bands: 
//1) "classified" with two classes (0 or 1 - tailings or not-tailings),
//2) "tailings" with two classes. This is the masked tailings database layer for reference and accuracy assessment

function ClassifyFromSamples(points, variables){
  
  points = ee.Dictionary(points);
  variables = ee.List(variables);
  var testingVector = ee.FeatureCollection(points.get("districtVector"));
  var classificationArea = trainingImage.clipToCollection(testingVector);
  
  //Get a new seed each time
  //points contains a "seed" property
  var classifierSeed = ee.Number(points.get("seed"));
  
  //Unpack feature collections from dictionary
  var tailingsPoints = ee.FeatureCollection(points.get("tailingsPoints"));
  var nonTailingsPoints = ee.FeatureCollection(points.get("nonTailingsPoints"));
  
  //Combine all points to train classifier
  var trainingPoints = nonTailingsPoints.merge(tailingsPoints);
  
  //Random Forest
  var classifier = ee.Classifier.smileRandomForest({
    numberOfTrees: 250,
    seed: classifierSeed
  })
    .train({
      features: trainingPoints,
      classProperty: "tailings",
      inputProperties: variables
  });
  
  var classifiedImage = classificationArea.classify(classifier)
    .addBands(maskedTailings) //For confusion matrix later
    .reproject({crs: projection, scale: scale})
    .eq(1) //sets bitdepth to 1
    .set("seed", classifierSeed); //Seed is attached for later use.
    
  //Return a classified image. 
  return classifiedImage;
}

// To show what variables were used in classification
function ClassifyFromSamples_SchemaTest(points, variables){
  
  points = ee.Dictionary(points);
  variables = ee.List(variables);
  var testingVector = ee.FeatureCollection(points.get("districtVector"));
  var classificationArea = trainingImage.clipToCollection(testingVector);
  
  //Get a new seed each time
  //points contains a "seed" property
  var classifierSeed = ee.Number(points.get("seed"));
  
  //Unpack feature collections from dictionary
  var tailingsPoints = ee.FeatureCollection(points.get("tailingsPoints"));
  var nonTailingsPoints = ee.FeatureCollection(points.get("nonTailingsPoints"));
  
  //Combine all points to train classifier
  var trainingPoints = nonTailingsPoints.merge(tailingsPoints);
  
  //Random Forest
  var classifier = ee.Classifier.smileRandomForest({
    numberOfTrees: 250,
    seed: classifierSeed
  })
    .train({
      features: trainingPoints,
      classProperty: "tailings",
      inputProperties: variables
  });
  
  var classifierSchema = classifier.schema();
  return classifierSchema;
}

//Get Variable Importance
//Similar to ClassifyFromSamples, but returns a dictionary with information on classifier results
function GetVariableImportance(points, variableList){
  
  points = ee.Dictionary(points);
  variableList = ee.List(variableList); //A list of strings containing the names of variables
  
  //Get a new seed each time
  //points contains a "seed" property
  var classifierSeed = ee.Number(points.get("seed"));
  
  //Unpack feature collections from dictionary
  var tailingsPoints = ee.FeatureCollection(points.get("tailingsPoints"));
  var nonTailingsPoints = ee.FeatureCollection(points.get("nonTailingsPoints"));
  
  //Combine all points to train classifier
  var trainingPoints = nonTailingsPoints.merge(tailingsPoints);
  
  //Random Forest
  var classifier = ee.Classifier.smileRandomForest({
    numberOfTrees: 250,
    seed: classifierSeed //Get a new seed from the list at top for each iteration of classifier
  })
    .train({
      features: trainingPoints,
      classProperty: "tailings",
      inputProperties: variableList
  });
  
  //.explain() returns a dictionary with variable importance, out of bag error, and information on the classifier parameters
  var classifierExplained = classifier.explain();
  
  return classifierExplained;
}

//Pass a list of results from GetVariableImportance here to obtain mean VI of individual variables
function AggregateImportance(listOfVarImportance){
  
  listOfVarImportance = ee.List(listOfVarImportance); //A list of dictionaries
  // listOfVariables = ee.Dictionary(listOfVariables); //This should be a list of strings
  
  //Map through dictionary where each key is a variable and value is the corresponding variable importance
  //Using keys from the first dictionary in the list - they each have the same keys
  //Replace value with a list of features, each containing VI property for that variable
  //e.g.: B2: [0.7, 0.6, 0.7, 0.5, 0.5, ...] this will be aggregated to obtain mean VI
  
  var keyValue = ee.Dictionary(ee.Dictionary(listOfVarImportance.get(0)).get("importance"));
  
  var VIFromList = keyValue.map(function(key, value){
    
    var variable = key; //The variable being assessed, from list of bands
    //Map through list of variable importance dictionaries
    var listOfVI = listOfVarImportance.map(function(currentImportance){
      
      currentImportance = ee.Dictionary(currentImportance);
      var allVI = ee.Dictionary(currentImportance.get("importance")); //Dictionary of every variable's importance for this iteration
      
      var thisImportance = allVI.get(variable); //Get the importance of each variable independantly
      
      var viFeature = ee.Feature(null, {"VariableImportance": thisImportance});
      return viFeature;
    });
    
    var viFC = ee.FeatureCollection(listOfVI);
    var stats = viFC.aggregate_stats("VariableImportance");
      
    return stats;
  });

  return VIFromList;
}

// Mean spectra per district
// Input Feature Collection of districts and bands, return FC of means
// Image must include tailings
function MeanDistrictTailingsSpectra(districtsFC, image, bands){
  
  districtsFC = ee.FeatureCollection(districtsFC);
  image = ee.Image(image);
  // Remove SCL. Could remove tailings as well, but it works for QC
  bands = ee.List(bands).filter(ee.Filter.inList("item", ["SCL"]).not());
  
  // Mask by tailings pixels
  var tailingsImage = image.updateMask(image.select("tailings").eq(1));
  
  // Sample from each district
  var spectralMeans = districtsFC.map(function(district){
    
    district = ee.Feature(district);
    var districtName = district.get("ShortName");
    
    var districtSample = tailingsImage.sample({
      region: district.geometry(), 
      scale: scale,
      projection: projection,
      factor: 0.1,
      dropNulls: true});
    
    // Get mean for each band as features
    var bandMeans = bands.map(function(band){
      band = ee.String(band);
      var bandMean = districtSample.aggregate_mean(band);
      return bandMean;
    });
    
    // Combine bands and band means into a dictionary
    var featureProperties = ee.Dictionary.fromLists(bands, bandMeans)
      .set("District", districtName); //Add district name
    
    // Return feature collection from list of features
    return ee.Feature(null, featureProperties);
    
  });
    
  return spectralMeans;
  
}



/*********************************************
*                                            * 
*                  Indices                   *
*                                            *
*********************************************/
{
var indicesImages = ee.List([]);

//Add indices generated for masking to indices list
var indices = ee.List(["ndvi", "mndwi"]);
indicesImages = ee.List(indicesImages).add(ndvi).add(mndwi);

//////////////////////////////////////////////
//              Ferric Iron                 //
//////////////////////////////////////////////

//Defined in Mielke et. al. (2014) as "SWIR2 / NIR + RED / GREEN"
if(scale == 20){
  
  var fii = trainingImage.expression(
    "fii = b('B12') / b('B8') + b('B4') / b('B3')");

  indicesImages = indicesImages.add(fii);
  indices = ee.List(indices).add("fii");
  // Map.addLayer(fii, {palette: ["white", "brown"]}, "Ferric Iron Index");
}

//////////////////////////////////////////////
//                  IFD                     //
//////////////////////////////////////////////

//Iron feature depth is designed for HS images
if(scale == 20){
  
  var rInt = trainingImage.expression(
    "rInt = b('B4') + (b('B12') - b('B4')) * (835 - 665) / (1640 - 665)");
  //We don't want to keep rInt once IFD is calculated
  var ifdTrainingIntermediate = trainingImage.addBands(rInt); 
  var ifd = ifdTrainingIntermediate.expression(
    "ifd = b('rInt') - b('B8A')");
  //Only values above 0 are meaningful
  // var ifdMax = ee.Dictionary(ifd.reduceRegion({reducer: ee.Reducer.max(), bestEffort: true})).get("ifd");
  // var ifdMaxImage = ee.Image.constant(ifdMax);
  // ifd = ifd.updateMask(ifd.gte(0)).divide(ifdMaxImage).multiply(ee.Image(100*100));
  // var ifdSample = ifd.sample({numPixels: 10000, dropNulls: true});
  // var ifdHisto = ui.Chart.feature.histogram(ifdSample, "ifd");
  
  indicesImages = indicesImages.add(ifd);
  indices = ee.List(indices).add("ifd");
  // Map.addLayer(ifd, {min: 0, max: 100, palette: ["white", "blue"]}, "IFD");
}

//////////////////////////////////////////////
//                Coordinates               //
//////////////////////////////////////////////

// var coordinateImage = ee.Image.pixelLonLat();

// indicesImages = indicesImages.add(coordinateImage);
// indices = indices.add("longitude").add("latitude");


//////////////////////////////////////////////
//         Combine all Used Indices         //
//////////////////////////////////////////////

//Add each index image as a band to training image
if(includeIndices && !useBestBands){
  // print(indices);
  bands = ee.List(bands).cat(indices);
} else if(useBestBands){
  indices = indices.filter(ee.Filter.inList("item", bestBands));
}

trainingImage = ee.Image(indicesImages.iterate(function(img, lastImg){
  img = ee.Image(img);
  lastImg = ee.Image(lastImg);
  return lastImg.addBands(img);
}, trainingImage));

}

//////////////////////////////////////////////
//       Distribution of Band Values        //
//////////////////////////////////////////////

// Values of masked pixels
// var maskedSample = trainingImage.addBands(maskedTailings.rename("maskedTailings")).stratifiedSample({
//     numPoints: 0,
//     classBand: "maskedTailings",
//     classValues: [0,1],
//     classPoints: [0, 5000],
//     dropNulls: true,
//     scale: scale
// });

/*********************************************
*                                            * 
*       Set Training/Testing Districts       *
*                                            *
*********************************************/

///////////////////////////////////////////////
//     Filter Districts Without Tailings     //
///////////////////////////////////////////////

//Create district variable from table asset
var districts = ee.FeatureCollection(districtPoints);
//Remove Moose River, which is significantly different from database record
districts = districts.filter(ee.Filter.neq("ShortName", "Moose River"));
//Filter to only gold districts that contain a certain area of tailings
function filterDistrictsByTailings(region, districts, tailingsMask, scale){
  
  //Filter to tailings within target region
  var regionTailings = tailings.filterBounds(region);
  
  //Buffer all district points by distance set in Misc. Variable section above
  var districtBuffer = ee.FeatureCollection(districts.map(function(f){
      var bufferedDistrict = f.buffer(districtBufferDistance);
      
      //Sample every masked tailings pixel, then get count of sample features
      var tailingsCount = ee.FeatureCollection(maskedTailings.selfMask().sample({
        region: bufferedDistrict.geometry(), 
        scale: scale, 
        projection: projection,
        factor: 1,
        dropNulls: true,
        geometries: true
        })).size();
      
      return bufferedDistrict.set("TailingsPixelCount", tailingsCount);
    })
  );
  
  //Keep districts with at least 5 masked tailings pixels
  var districtsWithTailings = districtBuffer.filter(ee.Filter.gte("TailingsPixelCount", 5));

  return ee.FeatureCollection(districtsWithTailings);
}

var districtsFiltered = filterDistrictsByTailings(trainingImage.geometry(), districts, maskedTailings, scale);

//Build dictionary of districts for storing stats
if(randomDistricts){
  var districtNames = districtsFiltered.toList(districtsFiltered.size())
    .map(function(f){return ee.Feature(f).get("ShortName")});
  //Create dictionary of dictionaries to later be filled with stats. Empty for now.
  var districtStats = ee.Dictionary.fromLists(districtNames, ee.List.repeat(ee.Dictionary({}), districtNames.size()));
} else{
  //Create dictionary of dictionaries to later be filled with stats. Empty for now.
  var districtStats = ee.Dictionary.fromLists(targetDistrictNames, ee.List.repeat(ee.Dictionary({}), targetDistrictNames.size()));
}

// print(districtsFiltered.filterBounds(trainingImage.geometry()));

///////////////////////////////////////////////
//   Split Districts Into Training/Testing   //
///////////////////////////////////////////////

//Sorts districts into training and testing, then creates a buffered area around them
function subsetDistricts(districtsCleared, scale, seed){
  
  //For quickly changing the train/test ratio
  var split = ee.Number(0.8);
  
  //Add random column to district features and buffer
  //Creates a training set containing 80% of districts and testing set with 20%
  var districtsRand = districtsCleared.randomColumn("random", seed)
    .sort("random");
  var districtsCount = ee.Number(districtsRand.size()); //Size of districts collection
  var districtList = districtsRand.toList(districtsCount); //Create list to call by index
  
  //Set an index to each feature, ranked by random value
  var districtsOrdered = ee.List.sequence({start: 0, count: districtsCount}).map(function(n){
    
    var feature = ee.Feature(districtList.get(n));
    
    return feature.set("index", n);
  });
  districtsOrdered = ee.FeatureCollection(districtsOrdered);

  //Split districts based on number of input districts, rather than actual random value, to maintain consistent # of sample areas
  var trainingSplit = ee.Number(districtsCount.multiply(split)).floor();
  
  //Create district buffer, split into training and validation
  return ee.Dictionary({
    training: ee.FeatureCollection(districtsOrdered
      .filter(ee.Filter.lt("index", trainingSplit))),
    testing: ee.FeatureCollection(districtsOrdered
      .filter(ee.Filter.gte("index", trainingSplit)))
      });
}

//Test Subsets
// print(districtSubset.get("training"), "Training Districts");
// print(districtSubset.get("testing"), "Testing Districts");
// Map.addLayer(ee.FeatureCollection(districtSubset.get("training")), {color: "blue"}, "Training Districts");
// Map.addLayer(ee.FeatureCollection(districtSubset.get("testing")), {color: "red"}, "Testing Districts");

/*********************************************
*                                            * 
*           Collect Sample Points            *
*                                            *
*********************************************/

//Add masked tailings and indices
//Tailings band is used for stratified sampling. It is removed before classifier is trained.
trainingImage = trainingImage.addBands(maskedTailings, ["tailings"]);
bands = ee.List(bands).add("tailings");

//If not using spectral bands, they are removed here (set up top)
if(!includeBands && scale == 20){
  bands = bands.removeAll(bands20).add("SCL"); //SCL has to be re-added
} else if(!includeBands && scale == 10){
  bands = bands.removeAll(bands10).add("SCL");
}

//If using set districts, the following is run once. 
//Otherwise, it's run in the subset loop further below
if(!randomDistricts){
  var targetDistricts = districtsFiltered.filter(ee.Filter.inList("ShortName", targetDistrictNames));
  var nonTargetDistricts = districtsFiltered.filter(ee.Filter.inList("ShortName", targetDistrictNames).not());
  var districtSubset = ee.Dictionary({training: nonTargetDistricts, testing: targetDistricts});
  var testingVector = targetDistricts;
}

//Show districts being tested
// Map.addLayer(testingVector.style({color: "grey", fillColor: "00000000"}), null, "Testing Areas");

//Sample Points Quantities
var totalSamplePoints = ee.Number(totalSamplePoints); //Set up top

//Set sample points in each land cover class according to their areal proportion
var classVals = ee.List.sequence(2, 9);
/*Classes of interest in S2's land cover band:
// 2 Dark Area Pixels
// 3 Cloud Shadows
// 4 Vegetation
// 5 Bare Soils
// 6 Water
// 7 Clouds Low Probability / Unclassified
// 8 Clouds Medium Probability
// 9 Clouds High Probability
*/

var pixelCount = classVals.map(function(n){ //Loop through each class value and count number of pixels
    //Return a list of the number of pixels in each class
    return trainingImage.select("SCL").eq(ee.Number(n))
            .selfMask()
            .reduceRegion({reducer: ee.Reducer.count(), scale: scale, maxPixels: 1e12})
            .get("SCL");
});
var totalPixels = ee.Number(pixelCount.reduce(ee.Reducer.sum())); //Sum of non-null pixels in training image

//Produce list of counts for each SCL value, based on proportion of pixels in image per class
var nonTailingsClassPoints = pixelCount.map(function(n){
  n = ee.Number(n); //n = total pixels per SCL label
  
  return ee.Number(n
    .divide(totalPixels) //Get ratio of pixels in each class to total image pixels
    .multiply(totalSamplePoints) //Make proportional samples match the total number of samples
    .floor() //Convert to int, rounding down
    .add(25)); //Add base number to ensure under-represented classes are still sampled. If there are fewer than n pixels, all are sampled.
});

// var classPoints = [100, 25, 500, 50, 100, 50, 25, 25];
//                 //[287,250,3495,298,1897,268,251,250]

// print(classPoints, "Class Points");

//Collect training points. Training image and training area (districts) are set above and don't change
function CollectTrainingPoints(bands, seed, districtSubsetDictionary){
  
  //Set to true for illustrating sample point collection, debugging, or exporting
  var drawGeometries = true;
  
  //Unpack district subset dictionary
  districtSubsetDictionary = ee.Dictionary(districtSubsetDictionary);
  var trainingVector = ee.FeatureCollection(districtSubsetDictionary.get("training"));
  var testingVector = ee.FeatureCollection(districtSubsetDictionary.get("testing"));
  var allDistricts = ee.FeatureCollection(trainingVector.merge(testingVector));
  
  //Sample non-tailings outside of districts
  var nonTailingsTrainingRegion = trainingImage.geometry().difference(allDistricts);

  //Collect Points
  //Non-tailings points are stratified by S2's land classification band
  var nonTailingsSample = trainingImage.stratifiedSample({
    numPoints: 0, //Overridden by classPoints
    region: nonTailingsTrainingRegion,
    scale: scale,
    seed: seed,
    classBand: "SCL",
    classValues: classVals, //Which SCL labels to use. Set above
    classPoints: nonTailingsClassPoints, //How many points per class
    dropNulls: true,
    geometries: drawGeometries
  });
  
  //Sample tailings from everywhere except test areas
  var tailingsTrainingRegion = trainingImage.geometry().difference(testingVector);
  //Adjust sample count by total available pixels
  //Take 2/3 of tailings pixels
  var tailingsPoints = GetSampleCount(maskedTailings.selfMask(), tailingsTrainingRegion, 0.33);
  var tailingsClassPoints = ee.List([0, tailingsPoints]); 
  
  var tailingsSample = trainingImage.stratifiedSample({
    numPoints: 0, //Overridden by classPoints
    region: tailingsTrainingRegion,
    scale: scale,
    seed: seed,
    classBand: "tailings",
    classValues: [0,1],
    classPoints: tailingsClassPoints,
    dropNulls: true,
    geometries: drawGeometries
  });
    
  return ee.Dictionary({nonTailingsPoints: nonTailingsSample, 
  tailingsPoints: tailingsSample, 
  seed: seed,
  districtVector: testingVector});
}

///////////////////////////////////////////////
//      Create List of Sample Points         //
///////////////////////////////////////////////

var sampleBands = ee.List(bands).remove("SCL"); 

//We want to classify points many times and collect accuracy as we go. Create as many points as "iterations" the model is to be run
//Use district subset from list set above
var samplePointsSetFixed = ee.List.sequence({start: 0, count: iterations}).map(function(n){
  
  n = ee.Number(n);
  //Get seed from earlier seedlist using n
  var sampleSeed = seedList.get(n);
  //Add sample selection to list
  return CollectTrainingPoints(sampleBands, sampleSeed, districtSubset);
  
});

//Use new, random districts each time samples are collected
var samplePointsSetRandom = ee.List.sequence({start: 0, count: iterations}).map(function(n){
  
  n = ee.Number(n);
  //Get seed from earlier seedlist using n
  var sampleSeed = seedList.get(n);
  //Get new district subset each iteration
  var districtSubset = subsetDistricts(districtsFiltered, scale, sampleSeed);
  //Add samples to list
  return CollectTrainingPoints(sampleBands, sampleSeed, districtSubset);
  
});

//Based on boolean variable set at top, one of these will be called
if(!randomDistricts){
  var samplePointsSet = samplePointsSetFixed;
} else{
  var samplePointsSet = samplePointsSetRandom;
}

//Some tests to show that each list entry is unique (turn on geometries in sample function)
// Map.addLayer(ee.FeatureCollection(ee.Dictionary(ee.List(samplePointsSet).get(0)).get("nonTailingsPoints")), {color: "red"}, "Non tailings points", false);
// Map.addLayer(ee.FeatureCollection(ee.Dictionary(ee.List(samplePointsSet).get(0)).get("tailingsPoints")), {color: "green"}, "Training points", false);
// Map.addLayer(ee.FeatureCollection(ee.Dictionary(ee.List(samplePointsSet).get(9)).get("nonTailingsPoints")), {color: "blue"}, "Non training points");
// Map.addLayer(ee.FeatureCollection(ee.Dictionary(ee.List(samplePointsSet).get(9)).get("tailingsPoints")), {color: "yellow"}, "Training points");

/*********************************************
*                                            * 
*              Classification                *
*                                            *
*********************************************/

//Section includes accuracy assessment tools, variable importance ranking tools, 
//and final classification

//This is set up top. Classify full image or just test areas.
// if(fullMap){
//   var classificationArea = trainingImage;
// }else{
//   var classificationArea = testingArea;
// }

/////////////////////////////////////////////
//       Accuracy Assessment Function      //
/////////////////////////////////////////////

//Each image has a classification and a tailings band
//Sample all tailings in reference layer and see if they were classified - recall
//Sample all classified tailings and see if they're marked as tailings in reference band - precision

//Calculate FScore of classified image. Set b at top
function GetAccuracy(img, b, testVector){
  
  b = ee.Number(b); //Adjust the precision/recall weight for F Score. 1 = equal weight. 2 = recall is twice as important as precision
  img = ee.Image(img);
  
  //Sample count will adjust based total pixels
    
  //Sample portion of tailings
  var tailingsCount = GetSampleCount(img.select("tailings").selfMask(), testingVector, 0.5);
  //Sample portion of positive classified pixels
  var classifiedCount = GetSampleCount(img.select("classification").selfMask(), testingVector,  0.5);
  
  var tailingsSampleCount = ee.List([tailingsCount.multiply(2), tailingsCount]); 
  var classifiedSampleCount = ee.List([tailingsCount.multiply(2), tailingsCount]);
  
  //Sample from both reference ("tailings") and classified pixels, then combine. Ensures sampling of FP and FN
  var samplePointsTailings = img.stratifiedSample({
    numPoints: 0, //Overridden by classPoints
    classBand: "tailings",
    region: testingVector,
    scale: scale,
    projection: projection,
    seed: seedList.shuffle().get(0),
    classValues: [0,1],
    classPoints: tailingsSampleCount,
    dropNulls: true,
    geometries: false
  });

  var samplePointsClassification = img.stratifiedSample({
    numPoints: 0, //Overridden by classPoints
    classBand: "classification",
    region: testingVector,
    scale: scale,
    projection: projection,
    seed: seedList.shuffle().get(0),
    classValues: [0,1],
    classPoints: classifiedSampleCount,
    dropNulls: true,
    geometries: false
  })
  .filter(ee.Filter.bounds(samplePointsTailings, 1).not()); //Pixels that were sampled twice are removed by filter before merging.
  
  var samplePoints = samplePointsTailings.merge(samplePointsClassification); //Combine sample sets
  
  //Generate error matrix as an array
  var errorMatrix = samplePoints.errorMatrix("tailings", "classification").array();
  //Unpack error matrix to calculate variables for f1-score
  //True positives
  var TP = errorMatrix.get([1,1]);
  //False negatives
  var FN = errorMatrix.get([1,0]);
  //False positives
  var FP = errorMatrix.get([0,1]);
  //Recall is correct classifications out of total positives in reference data (i.e., masked tailings pixels)
  var recall = ee.Number.expression("TP / (TP + FN)", {"TP": TP, "FN": FN});
  //Precision is correct classifications out of total positive classifications
  var precision = ee.Number.expression("TP / (TP + FP)", {"TP" : TP, "FP": FP});
  //F-Score
  var fScore = ee.Number.expression("(1 + b**2) * ((recall * precision) / ((b**2 * precision) + recall))", 
    {"b": b, "recall": recall, "precision": precision});
  
  return ee.Dictionary({fScore: fScore, recall: recall, precision: precision});
}

///////////////////////////////////////////////
//            Variable Importance            //
///////////////////////////////////////////////

//Variables to be tested, one at a time
var variableTestList = bands.removeAll(["SCL", "tailings", "B8"]); //Remove bands that aren't used for training/testing
if(!includeIndices){
  variableTestList = variableTestList.removeAll(indices);
}
if(!includeBands){
  variableTestList = variableTestList.removeAll(bands);
}

// print(variableTestList, "Variables Used in Classifier");

//Importance List
//Generate variable importance for each sample point set
var variableImportanceList = samplePointsSet.map(function(points){
  
  var importance = GetVariableImportance(points, variableTestList);
  
  return importance;
});
//Loop through all variables being used and aggregate stats for each
var importanceStats = AggregateImportance(variableImportanceList);
// print("Importance Stats", importanceStats);
//Create list of lists by combining variables with aggregated stats
var importanceKeys = importanceStats.keys();
var importanceValues = importanceStats.values().map(function(dict){return ee.Dictionary(dict).get("mean")});
var importanceList = importanceKeys.zip(importanceValues);
//Convert to FC by mapping through list of lists
var importanceFC = ee.FeatureCollection(importanceList.map(function(list){
  
  list = ee.List(list);
  var variable = list.get(0);
  var mean = list.get(1);
  
  return ee.Feature(null, {"Variable": variable, "Mean": mean});
})).sort("Mean", false); //Sort by mean in descending order

if(!useBestBands){
  print(ui.Chart.feature.byProperty(importanceFC, "Mean", "Variable"));
}

// var meanImportance = importanceFC.aggregate_mean("Mean");

//////////////////////////////////////////////////
//  Variable Importance Iterative Introduction  //
//////////////////////////////////////////////////

function CreateFScoreFeature(variableSet){
  // Get FScore and convert to a feature. Easier to iterate through
  variableSet = ee.List([variableSet]).flatten();
  // Points should be consistent across each accuracy test
  // If districts change, this could be the cause of of accuracy loss
  var points = ee.Dictionary(samplePointsSet.get(0)); //Get set of points
  var testingVector = ee.FeatureCollection(points.get("districtVector"));
  var classifiedImage = ClassifyFromSamples(points, variableSet);
  
  var fScore = GetAccuracy(classifiedImage, b, testingVector).get("fScore");
  
  return ee.Feature(null, {FScore: fScore, Variables: variableSet});
}

function CreateImportanceFeature(variableSet, points){
  // Get importance and OOB error and convert to a feature. Easier to iterate through
  variableSet = ee.List([variableSet]).flatten();
  // Unlike FScore feature, districts can change each iteration
  // Using OOB error shouldn't be as affected by test district
  var explainedDict = ee.Dictionary(GetVariableImportance(points, variableSet));
  var importance = explainedDict.get("importance");
  var oobError = explainedDict.get("outOfBagErrorEstimate");
  
  return ee.Feature(null, {"Importance": importance, "OobError": oobError, "Variables": variableSet});
}

// Right now the mean FScore uses same number of iterations as set above. This may be expensive
// Could be fixed by either sampling new points within function or ensuring that mean fscore iterations
// is <= model iterations
function CreateMeanFScoreFeature(variableSet, iterations){
  // Get mean FScore by repeating classification and convert to feature
  variableSet = ee.List([variableSet]).flatten();
  // A list to store features. Indexed to get points from sample collection
  var featureList = ee.List.sequence({start: 0, count: iterations})
    .map(function(n){
      
      var points = samplePointsSet.get(n);
      var testVector = ee.Dictionary(points).get("districtVector");
      var classifiedImage = ClassifyFromSamples(points, variableSet);
      var fScore = ee.Number(GetAccuracy(classifiedImage, b, testVector).get("fScore"));
      
      return ee.Feature(null, {FScore: fScore});
    });
  var meanFScore = ee.FeatureCollection(featureList)
    .aggregate_mean("FScore");
  
  return ee.Feature(null, {FScore: meanFScore, Variables: variableSet});
}

function OptimizeVariables(importanceFC, variableList){
  // Subtractive iteration:
  // Start with full list and remove variables. Start with least important
  // If removing variable reduces accuracy, keep variable in set
  var firstVariable = importanceFC.sort("Mean").first().get("Variable");
  var startFeature = CreateFScoreFeature(variableList);//Start with full list
  
  // Whole list minus first variable, which is passed back in as second parameter for .iterate()
  var subtractiveVariables = importanceFC.filter(ee.Filter.eq("Variable", firstVariable).not()) 
    .iterate(function(currentFeature, previousFScoreFeature){
    
      //Create current feature
      var previousVariableSet = ee.List(ee.Feature(previousFScoreFeature).get("Variables"));
      //Results of last iteration, minus current variable
      var currentVariableSet = previousVariableSet.remove(ee.Feature(currentFeature).get("Variable"));
      var currentFScoreFeature = CreateFScoreFeature(currentVariableSet); //Convert to feature
      var currentFScore = currentFScoreFeature.get("FScore");
      
      //Compare to previous feature
      var previousFScore = ee.Feature(previousFScoreFeature).get("FScore");
      var fScoreList = ee.List([currentFScore, previousFScore]); //Put values in list compare
    
      //Return highest FScore variables
      //If current Fscore (set - current variable) is lower than previous FScore, keep it in variable set.
      //ee.Algorithms.If interprets a non-empty list as "true" and empty list as "false".
      //If FScore is a tie (unlikely), it will keep previous set.
      var bestSet = ee.Algorithms.If(fScoreList.filter(ee.Filter.lt('item', previousFScore)), previousFScoreFeature, currentFScoreFeature);
    
      return bestSet;
      
    }, startFeature);
  
  var startFScore = startFeature.get("FScore");
  return ee.Feature(subtractiveVariables).set("AllVariablesFScore", startFScore);
}

// Creates FScore feature 'iterations' times, then takes mean to determine change in FScore,
// rather than just testing it once per variable
function OptimizeVariablesByMean(importanceFC, variableList, iterations){
  // Subtractive iteration:
  // Start with full list and remove variables. Start with least important
  // If removing variable reduces accuracy, keep variable in set
  var firstVariable = importanceFC.sort("Mean").first().get("Variable");
  var startFeature = CreateFScoreFeature(variableList);//Start with full list
  
  // Whole list minus first variable, which is passed back in as second parameter for .iterate()
  var subtractiveVariables = importanceFC.filter(ee.Filter.eq("Variable", firstVariable).not()) 
    .iterate(function(currentFeature, previousFScoreFeature){
    
      //Create current feature
      var previousVariableSet = ee.List(ee.Feature(previousFScoreFeature).get("Variables"));
      //Results of last iteration, minus current variable
      var currentVariableSet = previousVariableSet.remove(ee.Feature(currentFeature).get("Variable"));
      // var currentFScoreFeature = CreateFScoreFeature(currentVariableSet); //Convert to feature
      // var currentFScore = currentFScoreFeature.get("FScore");
      
      var currentFScoreFeature = CreateMeanFScoreFeature(currentVariableSet, iterations);
      var currentFScore = currentFScoreFeature.get("FScore");
      
      //Compare to previous feature
      var previousFScore = ee.Feature(previousFScoreFeature).get("FScore");
      var fScoreList = ee.List([currentFScore, previousFScore]); //Put values in list compare
    
      //Return highest FScore variables
      //If current Fscore (set - current variable) is lower than previous FScore, keep it in variable set.
      //ee.Algorithms.If interprets a non-empty list as "true" and empty list as "false".
      //If FScore is a tie (unlikely), it will keep previous set.
      var bestSet = ee.Algorithms.If(fScoreList.filter(ee.Filter.lt('item', previousFScore)), previousFScoreFeature, currentFScoreFeature);
    
      return bestSet;
      
    }, startFeature);
  
  return ee.Feature(subtractiveVariables);
}

function importanceToFC(importanceFeature){
  // Takes each dictionary entry from the importance feature and converts to an FC for sorting, etc
  importanceFeature = ee.Feature(importanceFeature);
  var importanceDict = ee.Dictionary(importanceFeature.get("Importance"));
  var featureList = importanceDict.keys().map(function(key){
    
    var importance = importanceDict.get(key);
    return ee.Feature(null, {"Variable": key, "Importance": importance});
  });
  
  return ee.FeatureCollection(featureList);
}

function OptimizeVariables_nonRecursive(points, variableList){
  // Subtractive iteration:
  // Start with full list and remove variables. Start with least important
  // If removing variable reduces accuracy, keep variable in set
  // var firstVariable = importanceFC.sort("Mean").first().get("Variable");
  // var startFeature = CreateFScoreFeature(variableList);//Start with full list
  
  var startFeature = CreateImportanceFeature(variableList, points); // Importance converted to feature
  var allVariables = importanceToFC(startFeature);
  var firstVariable = allVariables.sort("Importance").first();

  // Whole list minus first variable, which is passed back in as second parameter for .iterate()
  var subtractiveVariables = allVariables.filter(ee.Filter.eq("Variable", firstVariable).not()) 
    .iterate(function(currentFeature, previousImportanceFeature){
    
      //Create current feature by removing least importance variable
      var previousVariableSet = ee.List(ee.Feature(previousImportanceFeature).get("Variables"));
      //Results of last iteration, minus current variable
      var currentVariableSet = previousVariableSet.remove(ee.Feature(currentFeature).get("Variable"));
      var currentImportanceFeature = CreateImportanceFeature(currentVariableSet, points); //Convert to feature
      var currentOob = currentImportanceFeature.get("OobError");
      
      //Compare to previous feature
      var previousOob = ee.Feature(previousImportanceFeature).get("OobError");
      var oobList = ee.List([currentOob, previousOob]); //Put values in list compare
    
      //Return lowest OOB Error Variable Set
      //If current OOB Error (set - current variable) is higher than previous FScore, keep it in variable set.
      //ee.Algorithms.If interprets a non-empty list as "true" and empty list as "false".
      //If FScore is a tie (unlikely), it will keep previous set.
      var bestSet = ee.Algorithms.If(oobList.filter(ee.Filter.gt('item', previousOob)), previousImportanceFeature, currentImportanceFeature);
    
      return bestSet;
      
    }, startFeature);
  
  var startOob = startFeature.get("OobError");
  return ee.Feature(subtractiveVariables).set("AllVariablesOobError", startOob);
}

if(!useBestBands){
  var subtractiveVariables = OptimizeVariables_nonRecursive(samplePointsSet.get(0), variableTestList);
  print(subtractiveVariables, "Subtractive Variables");

  var bandsString = ee.List(subtractiveVariables.get("Variables")).iterate(function(thisVar, varString){
    
    thisVar = ee.String(thisVar);
    return ee.String(varString).cat(thisVar).cat(", ");
    
  }, ee.String(""));
  
  print(bandsString, "Bands Used");
}

///////////////////////////////////////////////////////
//             Pixel-Wise Classification             //
///////////////////////////////////////////////////////

// Variables can be procedurally selected, or "best bands" can be used (set up top)
if(!useBestBands){
  var variableList = ee.List(subtractiveVariables.get("Variables")).remove("B8");
  // print("Variables Used:", variableList);
  // print("Number of Bands Used: ", ee.List(variableList).size());
} else{
  var variableList = bestBands;
}

//FScore by district
function fScoreByDistrict(classifiedImage, testVectors){
  
  classifiedImage = ee.Image(classifiedImage);
  testVectors = ee.FeatureCollection(testVectors);
  var testVectorList = testVectors.toList(testVectors.size()).map(function(f){
    f = ee.Feature(f);
    var districtName = f.get("ShortName");
    //Clip classified image to one test district at a time
    var classifiedDistrict = classifiedImage.clip(f);
    var distAccuracy = GetAccuracy(classifiedDistrict, b, f);
    return ee.Dictionary({"district": districtName,"accuracy": distAccuracy});
  });
  return testVectorList;
}

//Classify samples taken earlier
var classifiedImageList = samplePointsSet.map(function(points){
  
  // Each points feature has sample points and a testing vector
  points = ee.Dictionary(points);
  var testingVector = ee.FeatureCollection(points.get("districtVector"));
  //Classify the image. This will use vector set above
  var classifiedImage = ClassifyFromSamples(points, variableList);
  //Get accuracy of classified image, using test vector
  var classifiedAccuracy = GetAccuracy(classifiedImage, b, testingVector);
  //Get accuracy by district for later aggregation. A list of dictionaries
  var fScoreByDistricts = fScoreByDistrict(classifiedImage, testingVector);
  
  return ee.Dictionary({
    "classifiedImage": classifiedImage, 
    "accuracy": classifiedAccuracy, 
    "byDistrict": fScoreByDistricts,
  });
});

// Shows what variables are being used to classify. Especially important for checking that
// the tailings band didn't somehow get added here
// var schemaList = samplePointsSet.map(function(points){
  
//     points = ee.Dictionary(points);
//     var schemaTest = ClassifyFromSamples_SchemaTest(points, variableList)
//     return schemaTest
// })

// print(schemaList)

//Go through list where each sample point is a new set of training/testing districts
//Optimizes variables each time using OOB samples
var classifiedMCOptimized = samplePointsSet.map(function(points){
  
  // Each points feature has sample points and a testing vector
  points = ee.Dictionary(points);
  
  var optimizedVariables = OptimizeVariables_nonRecursive(points, variableTestList);
  var testingVector = ee.FeatureCollection(points.get("districtVector"));
  //Classify the image. This will use vector set above
  var classifiedImage = ClassifyFromSamples(points, optimizedVariables.get("Variables"));
  //Get accuracy of classified image, using test vector
  var classifiedAccuracy = GetAccuracy(classifiedImage, b, testingVector);
  //Get accuracy by district for later aggregation. A list of dictionaries
  var fScoreByDistricts = fScoreByDistrict(classifiedImage, testingVector);
  
  return ee.Dictionary({
    "classifiedImage": classifiedImage, 
    "accuracy": classifiedAccuracy, 
    "byDistrict": fScoreByDistricts});
});

//Change to MC list to collect stats if not using optimized bands
if(!useBestBands){
  classifiedImageList = classifiedMCOptimized;
}

//Get accuracy of each image and convert to feature to combine to FC and get stats
var accuracyStatsFeatures = ee.FeatureCollection(classifiedImageList.map(function(dict){
  dict = ee.Dictionary(dict);
  var img = ee.Image(ee.Dictionary(dict).get("classifiedImage"));
  //Get Accuracy
  var accuracy = ee.Dictionary(dict.get("accuracy")).get("fScore");
  //Convert to feature for aggregating
  var accuracyFeature = ee.Feature(null, {"FScore": accuracy});
  return accuracyFeature;
}));

//Create feature collection and get stats for each classified image
var accuracyStats = ee.FeatureCollection(accuracyStatsFeatures).aggregate_stats("FScore")
  .select(["max", "mean", "min", "total_count", "total_sd", "total_var"]);

//Get accuracy by district in each image
var accuracyStatsFeaturesByDistrict = ee.FeatureCollection(classifiedImageList.map(function(dict){
  var byDistrict = ee.List(ee.Dictionary(dict).get("byDistrict"));
  var img = ee.Image(ee.Dictionary(dict).get("classifiedImage"));
  //Get accuracy by district. This is a list of dictionaries
  var accuracyFeaturesByDistrict = byDistrict.map(function(dict){
    
    dict = ee.Dictionary(dict);
    var districtName = dict.get("district");
    var accuracy = ee.Dictionary(dict.get("accuracy")).get("fScore");
    //Convert to feature for aggregating
    var districtAccuracyFeature = ee.Feature(null, {"District": districtName, "Accuracy": accuracy});
    return districtAccuracyFeature;
  });
  return accuracyFeaturesByDistrict;
}).flatten());

//Sort features into collections based on district name, aggregate their stats, and assign to dictionary
districtStats = districtStats.map(function(districtName, value){
  
  //Get all features in accuracy FC that match dictionary key (district name)
  var districtFC = accuracyStatsFeaturesByDistrict.filter(ee.Filter.eq("District", districtName));
  var districtAccuracy = districtFC.aggregate_stats("Accuracy")
    .select(["max", "mean", "min", "total_count", "total_sd", "total_var"]);
  
  return districtAccuracy;
});

// Median of all classified images
var classifiedMedian = ee.ImageCollection.fromImages(
      classifiedImageList.map(function(dict){return ee.Dictionary(dict).get("classifiedImage")}))
  .select("classification")
  .reduce(ee.Reducer.median())
  .eq(1)
  .rename("classification") //rename from "classification_median"
  .addBands(maskedTailings, ["tailings"]); //Add tailings band to check F1 score

//Display accuracy stats
print(accuracyStats, "Overall Accuracy Stats");
print(districtStats, "Accuracy by District");

// Create image where each pixel is total count of classified tailings
var classifiedFrequency = ee.ImageCollection.fromImages(
      classifiedImageList.map(function(dict){return ee.Dictionary(dict).get("classifiedImage")}))
      .select("classification")
      .reduce(ee.Reducer.sum())
      .rename("classificationCount");
      
// print(classifiedFrequency);

///////////////////////////////////////////////////////
//           Object-Based Classification             //
///////////////////////////////////////////////////////

// Segment objects and classify
// Intended for use with best variables determined above

// Creates seeds within target district areas based on tailings/non tailings pixels. I.e., not a random, uniform grid
// targetArea is feature collection of test districts
// classBand should be a string. Probably either "tailings" or "classification". 
// Must be integer-typed in range [0,1]
function CreateTargetedSeeds(targetArea, imageToSegment, classBand, seed){
  
  targetArea = ee.FeatureCollection(targetArea);
  imageToSegment = ee.Image(imageToSegment);
  classBand = ee.String(classBand);
  
  // Convert to list to replace features of targetArea FC with FCs of sample points
  var targetCollectionCount = targetArea.size(); // How many districts/target polygons in target area
  var targetAreaList = targetArea.toList(targetCollectionCount); // Convert to list of target features
  
  // Create sample points
  var snicFeatureList = targetAreaList.map(function(f){
  
    f = ee.Feature(f); // An individual test district
    
    // Change sample points based on scale (we want more points at higher resolution)
    var nonTargetSampleCount = ee.Number(1000).divide(scale).floor();
    var targetSampleCount = ee.Number(nonTargetSampleCount).divide(5).floor(); //One fifth of non-target points
    var sampleSize = ee.List([nonTargetSampleCount, targetSampleCount]);

    var snicSamples = imageToSegment.stratifiedSample({
    numPoints: 0, // Overridden by classPoints
    classBand: classBand,
    region: f.geometry(),
    scale: scale,
    projection: projection,
    seed: seed,
    classValues: [0, 1],
    classPoints: sampleSize,
    geometries: true
    }).map(function(f){return f.set("Sample", 1)}); // Image reducer needs a value
    
    return snicSamples;
  });
  
  //Merge all snic feature points to one FC for reduction to image
  var snicColl = snicFeatureList.iterate(function(fc, previousFc){
    
    fc = ee.FeatureCollection(fc);
    
    return fc.merge(previousFc)}, 
    snicFeatureList.get(0));
  
  var snicSeeds = ee.FeatureCollection(snicColl).reduceToImage(["Sample"], ee.Reducer.max())
  .eq(1)
  .rename("seeds");
  
  return ee.Image(snicSeeds);
}

//Segment an image and classify it
function SegmentAndClassify(imageToSegment, seedGrid, variables, targetArea, samplePoints){
  
  variables = ee.List(variables);
  // Unpack sample points
  samplePoints = ee.Dictionary(samplePoints);
  var trainingPoints = ee.FeatureCollection(samplePoints.get("nonTailingsPoints"))
    .merge(ee.FeatureCollection(samplePoints.get("tailingsPoints")));
  
  // Add "tailings" to training image for stratified sampling
  var variablesWithTailings = variables.add("tailings");
  var trainingImage = imageToSegment.select(variablesWithTailings);
  // Separate training image and image to be segmented to avoid using "tailings" band in segmentation
  imageToSegment = ee.Image(imageToSegment).select(variables);
  
  // Segment image by growing superpixels from a seeds image
  var segmented = ee.Algorithms.Image.Segmentation.SNIC({
    image: imageToSegment, 
    compactness: 5, 
    connectivity: 8,
    neighborhoodSize: 25,
    seeds: seedGrid
  });
  
  // Assign mean reflectance value of all pixels to superpixels
  var connectedComponents = segmented.reduceConnectedComponents(ee.Reducer.max(), "clusters")
    .clip(targetArea);
  
  // Superpixel band names have _mean suffix. Replace with regular band names
  connectedComponents = connectedComponents.rename(variables);
  
  // Train classifier
  var connectedClassifier = ee.Classifier.smileRandomForest(250).train(trainingPoints, "tailings", variables);
  
  // Classify superpixels
  var connectedClassified = connectedComponents.classify(connectedClassifier, "classification")
    .addBands(maskedTailings.rename(["tailings"])) // Add tailings for FScore calculation
    .reproject({crs: projection, scale: scale});
  
  return connectedClassified.eq(1);
}

//Get accuracy by district in each image and convert to feature for export
function objectAccuracyByDistrict(segmentedAndClassifiedList){
  
  var districtAccuracies = segmentedAndClassifiedList.map(function(dict){
    dict = ee.Dictionary(dict);
    var byDistrict = dict.get("byDistrict");
    //Get accuracy by district. This is a list of dictionaries
    var accuracyFeaturesByDistrict = ee.List(byDistrict).map(function(dict){
      
      dict = ee.Dictionary(dict);
      var districtName = dict.get("district");
      // Get FScore, recall, and precision
      var accuracyDict = ee.Dictionary(dict.get("accuracy"));
      var fScore = accuracyDict.get("fScore");
      var recall = accuracyDict.get("recall");
      var precision = accuracyDict.get("precision");
      //Convert to feature for aggregating
      var districtAccuracyFeature = ee.Feature(null, {
        "District": districtName, 
        "FScore": fScore, 
        "Recall": recall, 
        "Precision": precision
      });
      return districtAccuracyFeature;
    });
  
    return accuracyFeaturesByDistrict;
  });
  return ee.FeatureCollection(districtAccuracies.flatten());
}

//////////////////////////////////////////////
//          Random Points as Seeds          //
//////////////////////////////////////////////

// Creates a snic seed image from random points, grows objects from those points, and classifies
var snicSegmentedRandom = samplePointsSet.map(function(samplePoints){
  // Sample dictionary contains nonTailingsPoints, tailingsPoints, seed, districtVector
  samplePoints = ee.Dictionary(samplePoints);
  // Unpack dictionary
  var testVector = ee.FeatureCollection(samplePoints.get("districtVector"));
  var seed = samplePoints.get("seed");
  // Create random points to convert to snic seed image
  var snicSeedFeatures = ee.FeatureCollection.randomPoints(testVector, 500, seed)
    .map(function(point){return point.set("Sample", 1)}); // Property for reducer
  var snicSeeds = snicSeedFeatures.reduceToImage(["Sample"], ee.Reducer.max());
  // Segment image using seeds and classify
  var clustersClassified = SegmentAndClassify(trainingImage, snicSeeds, variableList, testVector, samplePoints);
  // Get accuracy
  var accuracy = fScoreByDistrict(clustersClassified, testVector);
  
  return ee.Dictionary({"classifiedClusters": ee.Image(clustersClassified), "byDistrict": accuracy});
});

// print(snicSegmentedRandom)

//Get accuracy by district in each image and convert to feature for export
// var segmentedFromRandomPoints = objectAccuracyByDistrict(snicSegmentedRandom).aside(print);

//////////////////////////////////////////////
// Use Pixel-Wise Classification as Seeds   //
//////////////////////////////////////////////

// Takes sample collection from earlier, classifies each iteration 
// and then attempts to grow and classify clusters.
// Returns dictionary of classified image and associated stats

var snicSegmented = samplePointsSet.map(function(samplePoints){
  // Sample dictionary contains nonTailingsPoints, tailingsPoints, seed, districtVector
  samplePoints = ee.Dictionary(samplePoints);
  // Unpack dictionary
  var testVector = ee.FeatureCollection(samplePoints.get("districtVector"));
  var seed = samplePoints.get("seed");
  // Classify pixel-wise
  var classified = ClassifyFromSamples(samplePoints, variableList);
  // Create snic seeds, using classified images
  var snicSeeds = CreateTargetedSeeds(testVector, classified, "classification", seed);
  // Segment image using seeds and classify
  var clustersClassified = SegmentAndClassify(trainingImage, snicSeeds, variableList, testVector, samplePoints);
  // Get accuracy
  var accuracy = fScoreByDistrict(clustersClassified, testVector);
  
  return ee.Dictionary({"classifiedClusters": ee.Image(clustersClassified), "byDistrict": accuracy});
});

//Get accuracy by district in each image and convert to feature for export
// var segmentedFromRF = objectAccuracyByDistrict(snicSegmented);

// print(clusteredAccuracyByDistrict, "Cluster Accuracy");

//////////////////////////////////////////////
//          Use Sample Sites as Seeds       //
//////////////////////////////////////////////

// Known tailings samples can be input and used to grow objects
// A cloud of points will be made around these points,
// Otherwise the starting object will grow forever

// function objectsFromSamples(trainingSamplePoints, fieldSamplePoints, districtName){
  
//   // District name should be a string of district "shortName", i.e. "Montague"
//   // Only works on one district. Lazy implementation.
//   districtName = ee.String(districtName);
//   //Get feature for sampled district
//   var sampledDistrict = districtsFiltered.filter(ee.Filter.eq("ShortName", districtName));
  
//   fieldSamplePoints = ee.FeatureCollection(fieldSamplePoints);

//   // Generate random seeds across district. Assign "sample" property with value = 1 for reducer
//   var randomSeeds = ee.FeatureCollection.randomPoints(sampledDistrict, 250)
//     .map(function(point){return point.set("Sample", 1)});
  
//   // Combine sample and non-sample points
//   var seedFeatures = fieldSamplePoints.merge(randomSeeds);
  
//   // Convert to image for SNIC
//   var snicSeeds = seedFeatures.reduceToImage(["Sample"], ee.Reducer.max())
//     .reproject({crs: projection, scale: scale});
  
//   var objectsClassified = SegmentAndClassify(trainingImage, snicSeeds, variableList, sampledDistrict, trainingSamplePoints);
//   var objectsFScore = GetAccuracy(objectsClassified, b);

//   return ee.Dictionary({"classifiedImage": objectsClassified, "FScore": objectsFScore});
// }

// var objectsFromSamplePoints = samplePointsSet.map(function(pointsDict){
  
//   pointsDict = ee.Dictionary(pointsDict);

//   var classifiedDict = objectsFromSamples(pointsDict, sampledTailingsMO, ee.String(targetDistrictNames.get(0)));

//   return ee.Dictionary(classifiedDict);
// });

// print(objectsFromSamplePoints)
// var sampleObjImg = ee.Dictionary(ee.List(objectsFromSamplePoints).get(0)).get("classifiedImage")
// Map.addLayer(ee.Image(sampleObjImg).select("classification").selfMask(), {palette: ["white", "orange"]}, "Objects from Samples")

// var objectsFromSamplesAccuracy = ee.FeatureCollection(objectsFromSamplePoints.map(function(dict){
//   dict = ee.Dictionary(dict);

//   var districtName = ee.String(targetDistrictNames.get(0));
//   var accuracyDict = ee.Dictionary(dict.get("FScore"));
//   // Get FScore, recall, and precision
//   var fScore = accuracyDict.get("fScore");
//   var recall = accuracyDict.get("recall");
//   var precision = accuracyDict.get("precision");
//   //Convert to feature for aggregating
//   var districtAccuracyFeature = ee.Feature(null, {
//     "District": districtName, 
//     "FScore": fScore, 
//     "Recall": recall, 
//     "Precision": precision
//   });
  
//   return districtAccuracyFeature;
// }).flatten());

// // print(objectsFromSamplesAccuracy)

// var randomSeeds = ee.FeatureCollection.randomPoints(testingVector.geometry(), 500)
//   .map(function(point){return point.set("Sample", 1)});

// // Map.addLayer(randomSeeds)

// var randomSeedImg = ee.Algorithms.Image.Segmentation.seedGrid(200);

// // Map.addLayer(randomSeedImg)

/*********************************************
*                                            * 
*           Full Image Classifier            *
*                                            *
*********************************************/
function classifyImage(image, samplePoints, variableList){
  
  // Unpack sample points
  samplePoints = ee.Dictionary(samplePoints);
  var trainingPoints = ee.FeatureCollection(samplePoints.get("nonTailingsPoints"))
    .merge(ee.FeatureCollection(samplePoints.get("tailingsPoints")));

  // Train classifier
  var classifier = ee.Classifier.smileRandomForest(250).train(trainingPoints, "tailings", variableList);
  
  // Classify superpixels
  var classified = image.classify(classifier, "classification")
    .reproject({crs: projection, scale: scale});
  
  return classified.eq(1);
}

var minasImage = s2sr
  .filterBounds(minas)
  .filterDate("2019-07-01", "2019-09-01")
  .sort("CLOUDY_PIXEL_PERCENTAGE")
  .first();
  
var minasNDVI = minasImage.normalizedDifference(["B8", "B4"]).rename("ndvi");
var minasMNDWI = minasImage.expression("mndwi = (b('B3') - b('B11')) / (b('B3') + b('B11'))").rename("mndwi");
minasImage = minasImage.addBands(minasNDVI).addBands(minasMNDWI);
// Map.addLayer(minasImage, {bands: ["B4", "B3", "B2"], min: 50, max: 2000}, "Minas Image")

var trainingPoints = ee.Dictionary(samplePointsSet.get(0)).aside(print);
var fullImage = classifyImage(trainingImage, trainingPoints, bands20);
print(fullImage);
Map.addLayer(fullImage.select("classification").selfMask(), {palette: ["white", "red"]}, "Classified Image");

/*********************************************
*                                            * 
*              Results Export                *
*                                            *
*********************************************/

// Manually update date and increment testNum
// var exportFolder = "GEE Exports";
// var exportDate = "Aug15"; // No spaces
// var testNum = "02"; // Drive allows repeat names. Make sure this is updated to avoid confusion
// var suffix = exportDate + "_" + testNum; // Filename must be string literal
// var districtExportName = "objectsFromRandomGV_" + suffix;
// var overallAccuracyExportName = "statsPerIteration_" + suffix;
// //Export to Drive
// Export.table.toDrive(segmentedFromRandomPoints, "objectsFromRandomPoints", exportFolder, districtExportName);
// if(!useBestBands){
//   // If not using best bands, include band names in feature collection
//   accuracyStatsFeatures = accuracyStatsFeatures.map(function(f){return f.set("Bands", bandsString)}); 
// }
// Export.table.toDrive(accuracyStatsFeatures, "statsPerIteration", exportFolder, overallAccuracyExportName);

// Export district buffer zones

//Sample tailings spectra per district and export

// var distMean = MeanDistrictTailingsSpectra(districtsFiltered, trainingImage, bands);

// // Get non-tailings bare soil spectra
// // SCL 5 is "bare soil" class. Remove tailings from that via mask
// var bareSoilMask = trainingImage.select("SCL").eq(5).updateMask(tailingsImage.neq(1)).eq(1);
// var bareSoilSample = trainingImage.updateMask(bareSoilMask)
//   .select(bands20)
//   .sample({
//     region: trainingImage.geometry(), scale: scale, factor: 0.3
//   });
  
// // Get mean for each band as features
// var bandMeans = bands20.map(function(band){
//   band = ee.String(band);
//   var bandMean = bareSoilSample.aggregate_mean(band);
//   return bandMean;
// });

// // Combine bands and band means into a dictionary
// var featureProperties = ee.Dictionary.fromLists(bands20, bandMeans);
// // Convert to FC for export
// var bareSoilSpectraFC = ee.FeatureCollection(ee.Feature(null, featureProperties));

// Export.table.toDrive({
//   collection: distMean,
//   description: "MeanSpectraMO_Aug2",
//   folder: exportFolder,
//   fileNamePrefix: "MeanSpectraMO_Aug2",
//   fileFormat: "CSV",
// });

/*********************************************
*                                            * 
*               Map Layers                   *
*                                            *
*********************************************/

// Map.centerObject(ROI, 15);
//Training Image
Map.addLayer(trainingImage, {bands: ["B4", "B3", "B2"], min: 50, max: 2000}, "Training Image", false);
//Geometry of training image
var trainingImageOutline = ee.Algorithms.Feature(trainingImage.geometry());
Map.addLayer(ee.FeatureCollection(trainingImageOutline).style({color: "black", fillColor: "00000000", width: 1}), null, "Training Image Extent");

//Tailings Image
// Map.addLayer(tailingsImage.selfMask(), {min:0, max: 1, palette: ["white", "brown"]}, "Tailings Image");

//Masked Tailings Image
Map.addLayer(maskedTailings.selfMask(), {min: 0, max: 1, palette: ["white", "orange"]}, "Masked Tailings");

//Districts
Map.addLayer(districtsFiltered.style({color: "black", fillColor: "FF000000"}), null, "Filtered Districts");

//Normalized Elevation
// Map.addLayer(ndei, {min: -1, max: 1}, "District Normalized Mean Elevation");

//First Classified Image in List as Sample
// Map.addLayer(classifiedSampleImage.select("classification").selfMask(), {min: 0, max: 1, palette: ["white", "red"]}, "Sample classified image");

//Median of Classified Images
// Map.addLayer(classifiedMedian.select("classification").selfMask(), 
//   {min: 0, max: 1, palette: ["white", "red"]}, "Median Classified Tailings");

// Count of Classified Tailings
// Map.addLayer(classifiedFrequency.toDouble(), 
//   {min: 0, max: iterations, palette: ['ffe4e4', 'ff0000']}, "Classified Frequency");