import os
from osgeo import gdal
import netCDF4 as nc
import numpy as np
from glob import glob
from osgeo import osr

#Define work and output paths
WorkPath = r'/data/geodata/ctsc'
OutPath  = WorkPath

#Define air pollutant type
#e.g., PM2.5, BC, et al.
AP = 'PM2.5'

#Define spatial resolution
#e.g., 1 km ≈ 0.01 Degree
SP = 0.01 #Degrees

# if not os.path.exists(OutPath):
#     os.makedirs(OutPath)

file_paths = glob('/data/geodata/ctsc/pm_2_5/**/*.nc',recursive = True)

for file in file_paths:
    fname = os.path.basename(file).split('.nc')[0]
    #outfile will be the file variable with the .nc extension replaced with .tif

    f = nc.Dataset(file)
    #Read SDS data
    #f[AP][:] is the data
    data = np.array(f[AP][:])
    #Define missing value: NaN or -999
    data[data==65535] = np.nan #-999
    #Read longitude and latitude information
    lon = np.array(f['lon'][:])
    lat = np.array(f['lat'][:])
    LonMin,LatMax,LonMax,LatMin = lon.min(),lat.max(),lon.max(),lat.min()
    N_Lat = len(lat)
    N_Lon = len(lon)
    Lon_Res = SP #round((LonMax-LonMin)/(float(N_Lon)-1),2)
    Lat_Res = SP #round((LatMax-LatMin)/(float(N_Lat)-1),2)
    #Define Define output file
    file_without_nc = file.split('.nc')[0]
    outfile = file_without_nc + '.tif'
    print(outfile)
    #Write GeoTIFF
    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(outfile,N_Lon,N_Lat,1,gdal.GDT_Float32)
    outRaster.SetGeoTransform([LonMin-Lon_Res/2,Lon_Res,0,LatMax+Lat_Res/2,0,-Lat_Res])
    sr = osr.SpatialReference()
    sr.SetWellKnownGeogCS('WGS84')
    outRaster.SetProjection(sr.ExportToWkt())
    outRaster.GetRasterBand(1).WriteArray(data)
    print(fname+'.tif',' Finished')
    #release memory
    del outRaster
    f.close()