import numpy as np
import pandas as pd
from netCDF4 import Dataset
import xarray as xr
import matplotlib.pyplot as plt
import random
from pandas import DataFrame

"""Load data from google drive"""

dat = xr.open_dataset('/home/tanpv/Data/562_VNU/Rfc_Stn/clWRF/lead1_clWRF_baseline_Stn.nc', decode_times = False)
dat = dat.to_dataframe()
# dat

# var = input('Input the variable to plot: ')
var = 'pr'
newdat = dat[['time_bnds', 'lon', 'lat', var]]
# print(newdat)

import matplotlib.colors as clr
import geopandas as gpd


grouped = newdat.groupby('time_bnds')
for name, group in grouped:
  
  min = group[var].min()
  max = group[var].max()

  # Initialize an axis
  fig, ax = plt.subplots(figsize=(27, 9))

  country = gpd.read_file('/home/tanpv/Progs/562_VNU/Shp_VN/VN.shp')
  country.to_crs(epsg=4326).plot(color="lightgrey", ax=ax)
  newdat.plot(x = "lon", y = "lat", kind = "scatter", c = var, ax = ax, vmin = min, vmax = max)

  ax.set_title(f"Time {name} Variable: {var}", fontsize = 12, color = 'r')

  ax.grid(True, alpha = 0.5)

  plt.show()

#  fig.savefig('')