# Module to read GEOS-CF data for analysis and processing (e.g., building a new climatology
# profile)
# Contains
#   1) read_CF_SO2 (read GEOS-CF SO2 mixing ratio)
# V0.1 2024/04/27
#   Initial version.

import netCDF4
import numpy as np
import h5py

class read_GEOS_CF_SO2():
    def __init__(self, filename, path):
        self.path = path
        infile = self.path + filename
        #print(infile)
        f = netCDF4.Dataset(infile,'r')
        self.Psfc = np.array(f.variables['PS'][:])
        self.PHIS = np.array(f.variables['PHIS'][:]) #surface geopotential height
        self.Lat = np.array(f.variables['lat'][:])
        self.Lon = np.array(f.variables['lon'][:])
        self.SO2 = np.array(f.variables['SO2'][:])
        self.ZPBL = np.array(f.variables['ZPBL'][:])
        f.close()

class read_GEOS_CF_MET():
    def __init__(self,infile):
        f = netCDF4.Dataset(infile,'r')
        self.mid_layer_heights = np.array(f.variables['ZL'][:])
        self.U = np.array(f.variables['U'][:])
        self.V = np.array(f.variables['V'][:])
        self.Lat = np.array(f.variables['lat'][:])
        self.Lon = np.array(f.variables['lon'][:])
        self.SLP = np.array(f.variables['SLP'][:]) 
        self.DELP = np.array(f.variables['DELP'][:]) 
        f.close()

class read_GEOS_CF_sfc_wind():
    def __init__(self, infile):
        f = h5py.File(infile,'r')
        self.Lat = f['/lat'][:]
        self.Lon = f['/lon'][:]
        self.U = f['/U'][:]
        self.V = f['/V'][:]
        f.close()