// Cloud masking functions  // (c) Rebecca Dell and Neil Arnold University of Cambridge (December 2020) 

var cloudUtils = require('users/gena/packages:cloud-utils');
// 'cloud_utils' package available here: https://code.earthengine.google.com/e292fe9e8cf7f879f891f1aca07ef284 
//(c) 2018 Gennadii Donchyts. All rights reserved. Licensed under the terms of the MIT license.  
// See <https://opensource.org/licenses/MIT>.

// Module function to apply the masks 
exports.maskClouds = function(image) {
  // Find the clouds using BQA
  var cs = cloud_shadows(image).select('BQA_Cloud_Shadow');
  var c = clouds(image).select('BQA_Cloud');
  var cir = cirrus(image).select('BQA_Cirrus');
  
  // Combined mask
  var allCloud = cs.or(c.or(cir)).rename('All_BQA');
  
  // Try a buffer for cloud shadows
  var BQACloudDistance = allCloud.fastDistanceTransform(150).sqrt()
                           .multiply(ee.Image.pixelArea().sqrt());
  var BQACloudBuffer = BQACloudDistance.lte(4000);//.and(allCloud.not());
  //var shadowMask = BQACloudBuffer

  /*var imageCloud = image.addBands(ee.Image(c.mask(c)));
  imageCloud = imageCloud.addBands(ee.Image(cs.mask(cs)));
  imageCloud = imageCloud.addBands(ee.Image(cir.mask(cir)));
  imageCloud = imageCloud.addBands(ee.Image(allCloud.mask(allCloud)));
  imageCloud = imageCloud.addBands(ee.Image(shadowMask.mask(shadowMask)));
  imageCloud = imageCloud.addBands(ee.Image(BQACloudBuffer));
  return imageCloud;*/
  
  //var cloud_mask = allCloud.eq(1).or(shadowMask.eq(1));
  return image.updateMask(BQACloudBuffer.not());
};

//Local functions
// Get the BQA bits
var getQABits = function(image, start, end, newName) {
    var pattern = 0;
    for (var i = start; i <= end; i++) {
       pattern += Math.pow(2, i);
    }
    // Return a single band image of the extracted QA bits, calling the band newName.
    return image.select([0], [newName])
                  .bitwiseAnd(pattern)
                  .rightShift(start);
};

// Mask out cloud shadows
var cloud_shadows = function(image) {
  // Select the QA band.
  var QA = image.select(['BQA']);
  // Get the internal_cloud_algorithm_flag bit.
  return getQABits(QA, 7,8, 'BQA_Cloud_Shadow').gte(2); // 3 = highest confidence (ie strictest test)
};

// Mask out clouds
var clouds = function(image) {
  // Select the QA band.
  var QA = image.select(['BQA']);
  // Get the internal_cloud_algorithm_flag bit.
  return getQABits(QA, 5,6, 'BQA_Cloud').gte(2); // 3 = highest confidence (ie strictest test)
  // Return an image masking out cloudy areas.
};

// Mask out cirrus
var cirrus = function(image) {
  // Select the QA band.
  var QA = image.select(['BQA']);
  // Get the internal_cloud_algorithm_flag bit.
  return getQABits(QA, 11,12, 'BQA_Cirrus').gte(2); // 3 = highest confidence (ie strictest test)
  // Return an image masking out cloudy areas.
};