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

group = "ombg"
varname = "totalSnowDepth" 

daHofxPath1="/scratch2/NCEPDEV/stmp1/Youlong.Xia/landDA/cycle_land/C768/DA_GHCN/DA/hofx"
mdlHofxPath1="/scratch2/NCEPDEV/stmp1/Youlong.Xia/landDA/cycle_land/C768/DAgts_ghcnHOFX/DA/hofx"

daHofxPath2="/scratch2/NCEPDEV/stmp1/Youlong.Xia/landDA/cycle_land/C768/DA_imsGHCN/DA/hofx"

mdlHofxPath2="/scratch2/NCEPDEV/stmp1/Youlong.Xia/landDA/cycle_land/C768/openloop_imsGHCN/DA/hofx"

PlotPath="/scratch1/NCEPDEV/da/Youlong.Xia/DA_plotScripts/pythonScripts/plots/C768/"
fname = PlotPath+group+"_"+varname+"C768_DAgts_ghcnHOFX_global_stdev_timeseries_WY2020.png"

RestartDate = "2019-10-01"
DeltaTime = 1  # day
SimulationDuration = 366  # days
RestartDate = datetime.strptime(RestartDate, '%Y-%m-%d')

model_avg12 = np.empty([SimulationDuration])
model_avg22 = np.empty([SimulationDuration])
DA_avg11 = np.empty([SimulationDuration])
DA_avg21 = np.empty([SimulationDuration])
DA_avg31 = np.empty([SimulationDuration])

model_avg12[:] = np.nan
model_avg22[:] = np.nan
DA_avg11[:] = np.nan
DA_avg21[:] = np.nan
DA_avg31[:] = np.nan

xx = np.empty([SimulationDuration])

for SimDay in range(SimulationDuration):
    PassedTimeStepNo = (SimDay)
    ThisRestartDate = RestartDate + timedelta(days=PassedTimeStepNo * DeltaTime)
    print(ThisRestartDate)
    str_yr = ThisRestartDate.strftime('%Y')
    str_mh = ThisRestartDate.strftime('%m')
    str_dy = ThisRestartDate.strftime('%d')
    date = str_yr+str_mh+str_dy

    xx[SimDay] = SimDay
    HofxFile11 = daHofxPath1+"/letkf_hofx_ghcn_"+date+"00.nc"     
    HofxFile12 = mdlHofxPath1+"/letkf_hofx_ghcn_"+date+"00.nc"
    HofxFile21 = daHofxPath2+"/letkf_hofx_ghcn_"+date+"00.nc"
    HofxFile22 = mdlHofxPath2+"/letkf_hofx_ghcn_"+date+"00.nc"
    
    ombg11 = ReadNETCDF(HofxFile11 ,"ombg","totalSnowDepth")
    ombg12 = ReadNETCDF(HofxFile12 ,"ombg","totalSnowDepth")
    ombg21 = ReadNETCDF(HofxFile21 ,"ombg","totalSnowDepth")
    ombg22 = ReadNETCDF(HofxFile22 ,"ombg","totalSnowDepth")

    qc11 = ReadNETCDF(HofxFile11,"EffectiveQC0","totalSnowDepth")
    qc12 = ReadNETCDF(HofxFile11,"EffectiveQC1","totalSnowDepth")
    qc21 = ReadNETCDF(HofxFile21,"EffectiveQC0","totalSnowDepth")
    qc22 = ReadNETCDF(HofxFile21,"EffectiveQC1","totalSnowDepth")
    qc31 = ReadNETCDF(HofxFile12,"EffectiveQC0","totalSnowDepth")

    da31 = np.empty([len(ombg12)])
    da31 = ombg12

    for i in range(len(ombg11)):
        if(qc11[i] != 0 or qc12[i] != 0):
            ombg11[i] = float("nan")
            ombg22[i] = float("nan")
            da31[i] = float("nan")
        if(qc21[i] != 0 or qc22[i] != 0):
            ombg21[i] = float("nan")
    for i in range(len(ombg12)):
        if(qc31[i] != 0):
            ombg12[i] = float("nan")
 
    avg11= np.nanstd(ombg11)
    avg12= np.nanstd(da31)
    avg21= np.nanstd(ombg21)
    avg22= np.nanstd(ombg22)
    avg31= np.nanstd(ombg12)

    #if (SimDay <= 170) or (SimDay >= 173) :
    DA_avg11[SimDay] = avg11
    DA_avg21[SimDay] = avg21
    DA_avg31[SimDay] = avg31
    model_avg12[SimDay] = avg12
    model_avg22[SimDay] = avg22

    #if (SimDay >= 170) &( SimDay <= 172):
    #    print(ThisRestartDate)
    #    print(avg1, avg2, np.nanmin(dataout1), np.nanmax(dataout1), np.nanmin(dataout2), np.nanmax(dataout2))

# plot global averaged snow depth time series
plt.figure(figsize=(12,8))
ax1 = plt.subplot(1, 1, 1)
ax1.set_ylim(0, 150.0)
ax1.plot(xx,model_avg12,'k-',label="DAgts_ghcnHOFX: stdev(O-F)")
ax1.plot(xx,DA_avg11,'r-',label="DA_GHCN: stdev(O-F)")
ax1.plot(xx,model_avg22,'b-',label="openloopImsGHCNrun: stdev(O-F)")
ax1.plot(xx,DA_avg21,'g-',label="DA_imsGHCN: stdev(O-F)")
ax1.plot(xx,DA_avg31,'c-',label="DAgts_ghcnHOFX_gtsQC: stdev(O-F)")
ax1.set_ylabel("Snow Depth (mm)", fontsize=20)
ax1.set_xlabel("Simulation Days Since 1 October 2019", fontsize=20)
ax1.legend()
plt.tight_layout()
plt.savefig(fname)
plt.close()
