import ee

''' 0   Good quality fire
    1   Good quality fire-free land
    2   Invalid due to opaque cloud
    3   Invalid due to surface type or sunglint or LZA threshold exceeded or off earth or missing input data
    4   Invalid due to bad input data
    5   Invalid due to algorithm failure'''
# Bit-masking
BitMask_0 = 1 << 0
BitMask_1 = 1 << 1
BitMask_2 = 1 << 2
BitMask_3 = 1 << 3
BitMask_4 = 1 << 4
BitMask_5 = 1 << 5
BitMask_6 = 1 << 6
BitMask_7 = 1 << 7
BitMask_8 = 1 << 8
BitMask_9 = 1 << 9

def GcalcNBR (goesImg, aoi):
  #day = ee.Date(eachImg.get('system:time_start')).get('day','America/Los_Angeles')
  fireMode = goesImg.select('fireMode')
  fireMin = goesImg.select('fireMin')

  CMI_QF3 = goesImg.select('DQF_C03').int()
  CMI_QF6 = goesImg.select('DQF_C06').int()

  # To include active fire pixels - fireMin.lt(2)\ for next line
  QF_Mask = (fireMin.eq(1)\
    .Or(fireMin.gt(3)))\
    .And(CMI_QF3.lt(2))\
    .And(CMI_QF6.lt(2))\
    .rename('QFmask');
  GOESm = goesImg.select(['CMI_C03','CMI_C06']).updateMask(QF_Mask)
  NBR = GOESm.normalizedDifference(['CMI_C03', 'CMI_C06']).toFloat().rename('NBR')

  return goesImg.addBands([NBR,QF_Mask])

def GcalcCCsingle (goesImg):

  fireDQF = goesImg.select('DQF').int()
  CMI_QF3 = goesImg.select('DQF_C03').int()
  CMI_QF6 = goesImg.select('DQF_C06').int()

  #Right now, cloud mask is excluding clouds and water; active fire, bad data and fire free are unmasked. NBR mask exlcudes fire
  F_Mask = fireDQF.eq(0)
  C_Mask = (fireDQF.lt(2).Or(fireDQF.gt(2))).rename('C_Mask')
       #.And(CMI_QF3.lt(2)).And(CMI_QF6.lt(2)).rename('C_Mask')
  QF_Mask = (fireDQF.eq(1).Or(fireDQF.gt(3)))\
        .And(CMI_QF3.lt(2)).And(CMI_QF6.lt(2)).rename('QFmask')

  GOESmasked = goesImg.select(['CMI_C03','CMI_C06']).updateMask(QF_Mask)
  NBRmasked = GOESmasked.normalizedDifference(['CMI_C03', 'CMI_C06']).toFloat().rename('NBR')
  cloudMasked = goesImg.select('CMI_C03').updateMask(C_Mask).toFloat().rename('CC')
  fireMasked =  goesImg.select('CMI_C03').updateMask(F_Mask).toFloat().rename('FC')

  return goesImg.addBands([NBRmasked,cloudMasked, fireMasked,QF_Mask,C_Mask])

'''Parameter Array Name Value Bit(s) = Value
Sun Glint           QF1 Surface Reflectance None 6-7 = 00
Low Sun Mask        QF1 Surface Reflectance High 5 = 0
Day/Night           QF1 Surface Reflectance Day 4 =0
Cloud Detection     QF1 Surface Reflectance Confident Clear 2-3 = 00 or Problably Clear 2-3 = 01
Cloud Mask Quality  QF1 Surface Reflectance High or Medium 0-1 = 10 or 11
Snow/Ice            QF2 Surface Reflectance No Snow or Ice 5 = 0
Cloud Shadow        QF2 Surface Reflectance No Cloud Shadow 3 = 0
LandWater           QF2 Surface Reflectance Land, Snow, Arctic, Antarctic or Greenland, Desert 0-2  = 011, 100, 101, 110, 111
Thin Cirrus Flag    QF7 Surface Reflectance No Thin Cirrus 4 = 0
Aerosol Quantity    QF7 Surface Reflectance Climatology, Low or Medium 2-3 = 00, 01 or 10
Adjacent to Cloud   QF7 Surface Reflectance Not Adjacent to Cloud 1 = 0'''

def VcalcNBR (VIIRSimg):

    QF1 = VIIRSimg.select('QF1').int()
    QF2 = VIIRSimg.select('QF2').int()
    QF7 = VIIRSimg.select('QF7').int()

    QF_Mask = (QF1.bitwiseAnd(BitMask_3).eq(0)).And\
    ((QF2.bitwiseAnd(BitMask_2).eq(4)).Or((QF2.bitwiseAnd(BitMask_1).eq(0)))).And\
    (QF2.bitwiseAnd(BitMask_5).eq(0)).rename('QFmask');

    VIIRSm = VIIRSimg.select(['I2','M11']).updateMask(QF_Mask);
    NBR = VIIRSm.normalizedDifference(['I2','M11']).toFloat().rename('NBR')
    return VIIRSimg.addBands(NBR).addBands(QF_Mask)#.set('avgNBR', avgNBR)

''' Bit 1: Dilated Cloud
    Bit 2: Cirrus (high confidence)
    Bit 3: Cloud
    Bit 4: Cloud Shadow
    Bit 5: Snow
    Bit 6: Clear (0: Cloud or Dilated Cloud bits are set, 1: Cloud and Dilated Cloud bits are not set)
    Bit 7: Water
    Bits 8-9: Cloud Confidence (0: None, 1: Low, 2: Medium, 3: High)
    Bits 10-11: Cloud Shadow Confidence (0: None, 1: Low, 2: Medium, 3: High)
    Bits 12-13: Snow/Ice Confidence (0: None, 1: Low, 2: Medium, 3: High)
    Bits 14-15: Cirrus Confidence (0: None, 1: Low, 2: Medium, 3: High)'''

def LcalcNBR (LSimg):
  QApixel = LSimg.select('QA_PIXEL').int()
  QF_Mask =(QApixel.bitwiseAnd(BitMask_3).eq(0)).And\
           (QApixel.bitwiseAnd(BitMask_5).eq(0)).And\
           (QApixel.bitwiseAnd(BitMask_7).eq(0)).rename('QFmask');

  LSmasked = LSimg.select(['SR_B5','SR_B7']).updateMask(QF_Mask);
  NBR = LSmasked.normalizedDifference(['SR_B5','SR_B7']).toFloat().rename('NBR')
  return LSimg.addBands(NBR).addBands(QF_Mask)#.set('avgNBR', avgNBR)

''' 1    Saturated or defective
    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
    10   Cirrus
    11   Snow / Ice'''

def ScalcNBR (sentImg):
  SCL = sentImg.select('SCL');
  QF_Mask =(SCL.neq(6)).And\
    (SCL.neq(8)).And\
    (SCL.neq(9)).And\
    (SCL.neq(11))\
    .rename('QFmask');
  sentMasked = sentImg.select(['B8A','B12']).updateMask(QF_Mask); #B8 is another option- broadband NIR
  NBR = sentMasked.normalizedDifference(['B8A','B12']).toFloat().rename('NBR')
  return sentImg.addBands(NBR).addBands(QF_Mask).addBands(SCL)#.set('avgNBR', avgNBR)