#!/usr/bin/evn Python
# -*- coding: utf-8 -*-

import gdal
import numpy as np
import os
import glob


def mask_data(hdf_fname):

    n_bands = 7
    for band in range(n_bands):
        fname_qa = 'HDF4_EOS:EOS_GRID:"%s":MOD_Grid_BRDF:BRDF_Albedo_Band_Mandatory_Quality_Band%d' % (hdf_fname, band + 1)
        fname_surf_ref = 'HDF4_EOS:EOS_GRID:"%s":MOD_Grid_BRDF:Nadir_Reflectance_Band%d' % (hdf_fname, band + 1)

        # Get the quality data
        # Open dataset
        d_q = gdal.Open(fname_qa)
        # Get the data and put into an array
        quality = d_q.ReadAsArray()

        # Now the reflectance
        d_r = gdal.Open(fname_surf_ref)
        surf_refl = d_r.ReadAsArray()

        # Process quality band
        # https://lpdaac.usgs.gov/dataset_discovery/modis/modis_products_table/mcd43a4_v006
        # Product Quality
        # *Mandatory QA bit index:
        # 0 = processed, good quality (full BRDF inversions)
        # 1 = processed, see other QA (magnitude BRDF inversions)
        # 255 = Fill Value

        # Since we want to keep best quality pixels, all 0s, we can create a binary mask
        # where all pixel that are 0s will be True, the rest False
        quality = np.where(quality == 0 and 1, True, False)

        # Mask reflectance
        surf_refl = surf_refl * quality

        # Save array into a new GeoTiff file
        rows, cols = surf_refl.shape
        proj = d_r.GetProjection()
        gt = d_r.GetGeoTransform()

        driver = gdal.GetDriverByName('GTiff')
        driver_options = ['COMPRESS=DEFLATE',
                          'BIGTIFF=YES',
                          'PREDICTOR=1',
                          'TILED=YES']

        fname = "%sb%d.tif" % (os.path.basename(hdf_fname)[0:-3], band+1)
        # Create a new dataset
        new_ds = driver.Create(fname, cols, rows, 1,
                               gdal.GDT_Int16, driver_options)

        # Set cartographic projection
        new_ds.SetProjection(proj)
        new_ds.SetGeoTransform(gt)

        new_ds.GetRasterBand(1).WriteArray(surf_refl)
        new_ds = None


DataDir = "/home/users/showardc/LandCover/BEIS-LC/MCD43A4/DataDump"
hdf_files = glob.glob(os.path.join(DataDir, "*.hdf"))
hdf_files.sort()

for hdf_fname in hdf_files:
    #hdf_fname = "MCD43A4.A2017160.h17v03.006.2017171210406.hdf"
    mask_data(hdf_fname)