import netCDF4
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import cartopy.feature as cfeature
import numpy as np
import pandas as pd
import cartopy.crs as ccrs
import os
import sys
from datetime import datetime
from datetime import timedelta

def ReadNETCDF(FilePath, GroupName, VariName):
    f = netCDF4.Dataset(FilePath)
    MyArray = f.groups[GroupName].variables[VariName][:]
    f.close()
    return MyArray

date = sys.argv[1]
group = sys.argv[2]
varname = sys.argv[3]

daHofxPath="/scratch2/NCEPDEV/land/Youlong.Xia/landDA/cycle_land/C768/DA_GHCN_HR2/DA/hofx"
opHofxPath="/scratch2/NCEPDEV/land/Youlong.Xia/landDA/cycle_land/C768/DA_GHCN_HR2_bugFix/DA/hofx"
PlotPath="./plots/C768HR2_bugfix_check/"
fname = group+"_"+varname+"_"+date+"bug_bugFixed_diff_ims.png"
units='mm'

HofxFile1=daHofxPath+"/letkf_hofx_ghcn_"+date+"00.nc"
HofxFile2=opHofxPath+"/letkf_hofx_ghcn_"+date+"00.nc"
#print(HofxFile1)
#print(HofxFile2)
#exit()

dataout = ReadNETCDF(HofxFile1 ,"ombg","totalSnowDepth")
latout = ReadNETCDF(HofxFile1 ,"MetaData", "latitude")
lonout = ReadNETCDF(HofxFile1 ,"MetaData", "longitude")
dataout1 = ReadNETCDF(HofxFile2 ,"ombg","totalSnowDepth")

qc1 = ReadNETCDF(HofxFile1 ,"EffectiveQC0","totalSnowDepth")
qc2= ReadNETCDF(HofxFile2 ,"EffectiveQC0","totalSnowDepth")

for i in range(len(dataout)):
    if(qc1[i] != 0):
        dataout[i] = float("nan")
    if(qc2[i] != 0):
        dataout1[i] = float("nan")

dataout = abs(dataout)-abs(dataout1)

#print(dataout.size, dataout.shape, dataout.ndim)
stdev = np.nanstd(dataout)  # Standard deviation
omean = np.nanmean(dataout) # Mean of the data
datmi = np.nanmin(dataout)  # Min of the data
datma = np.nanmax(dataout)  # Max of the data
datcont = np.ma.count(dataout)

fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree(central_longitude=0))
ax.add_feature(cfeature.GSHHSFeature(scale='auto'))

# Plot grid lines
# ----------------
gl = ax.gridlines(crs=ccrs.PlateCarree(central_longitude=0), draw_labels=True,
                  linewidth=1, color='gray', alpha=0.5, linestyle='-')
gl.top_labels = False
gl.xlabel_style = {'size': 12, 'color': 'black'}
gl.ylabel_style = {'size': 12, 'color': 'black'}
gl.xlocator = mticker.FixedLocator(
   [-180, -135, -90, -45, 0, 45, 90, 135, 179.9])
ax.set_ylabel("Latitude",  fontsize=12)
ax.set_xlabel("Longitude", fontsize=12)

# Add figure labels
# ------------------
ax.set_title("bug-bugFixed ghcnDA(abs(da[O-F])-abs(da[O-F])):"+varname+" for "+date, pad=20, fontsize=20)
#ax.text(0.45, -0.1, 'Longitude', transform=ax.transAxes, ha='left')
#ax.text(-0.08, 0.4, 'Latitude', transform=ax.transAxes, rotation='vertical', va='bottom', fontsize=20)
text = f"Total Count:{datcont:0.0f}, Max/Min/Mean/Std: {datma:0.3f}/{datmi:0.3f}/{omean:0.3f}/{stdev:0.3f} {units}"
#print(text)
ax.text(0.1, -0.2, text, transform=ax.transAxes, va='bottom', fontsize=12)
cs = plt.scatter(x=lonout, y=latout, c=dataout, cmap="bwr", marker=',', s=10, vmax=100, vmin=-100, transform=ccrs.PlateCarree())
cb = plt.colorbar(cs, orientation='horizontal', shrink=0.5, pad=.04)
#cb.set_label("Units: mm", fontsize=12)
plt.savefig(os.path.join(PlotPath, fname))
plt.close()
