"""
## Introduction

We explain here how to read and plot EARTHDATA-MERRA2 Data Files.

### Data Format

The data format used is nc4. To read such files, we need to use netCDF4 Python package.

### External Resources

[ON NASA site](https://disc.gsfc.nasa.gov/information/howto) you can find a set of instructions how to read, explore and visualize data content from EARTHDATA-MERRA2 files.

In particular, we used this resource: [How to read and plot NetCDF MERRA-2 data in Python](https://disc.gsfc.nasa.gov/information/howto?title=How%20to%20read%20and%20plot%20NetCDF%20MERRA-2%20data%20in%20Python)


"""
"""
## Data Analysis

### Load Packages
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import folium
from folium.plugins import HeatMap, HeatMapWithTime
from netCDF4 import Dataset
import cartopy.crs as ccrs
"""
### Load Data


We will load here only one file.
"""
data = Dataset("/kaggle/input/earthdata-merra2-co/MERRA2_400.tavgM_2d_chm_Nx.202004.nc4", more="r")
"""
### Data Extraction

Let's inspect the dataset.
"""
print(data)
"""
We then filter atmospheric latitude/longitude/time/ and variables and load them as numpy arrays.
"""
lons = data.variables['lon'][:]
lats = data.variables['lat'][:]
time = data.variables['time'][:]
COCL = data.variables['COCL'][:,:,:]; COCL = COCL[0,:,:]
COEM = data.variables['COEM'][:,:,:]; COEM = COEM[0,:,:]
COLS = data.variables['COLS'][:,:,:]; COLS = COLS[0,:,:]
TO3 =  data.variables['TO3'][:,:,:];  TO3 =  TO3[0,:,:]
print(f"longitudes: {len(lons)}")
print(f"latitudes: {len(lats)}")
print(f"time: {len(time)}")
print(f"COCL: {len(COCL)}")
print(f"COEM: {len(COEM)}")
print(f"COLS: {len(COLS)}")
print(f"TO3:  {len(TO3)}")
print(f"Latitudes: {len(lats)}, vals[:]: {len(COCL[0])}, Longitudes: {len(lons)}, vals[:,:]: {len(COCL[0])}")
"""
## Plot data

We write first a function for plotting the data.
"""
def plot_merra_data(data, title='', date='April 2020'):
    fig = plt.figure(figsize=(16,8))
    ax = plt.axes(projection=ccrs.Robinson())
    ax.set_global()
    ax.coastlines(resolution="110m",linewidth=1)
    ax.gridlines(linestyle='--',color='black')
    plt.contourf(lons, lats, data, transform=ccrs.PlateCarree(),cmap=plt.cm.jet)
    plt.title(f'MERRA-2 {title} levels, {date}', size=14)
    cb = plt.colorbar(ax=ax, orientation="vertical", pad=0.02, aspect=16, shrink=0.8)
    cb.set_label('K',size=12,rotation=0,labelpad=15)
    cb.ax.tick_params(labelsize=10)
"""
We then apply the function for few of the data features.
"""
plot_merra_data(COCL, 'COCL')
plot_merra_data(COEM, 'COEM')
plot_merra_data(COLS, 'COLS')
plot_merra_data(TO3, 'TO3')