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/DAimsGTS/DA/hofx"

mdlHofxPath1="/scratch2/NCEPDEV/stmp1/Youlong.Xia/landDA/cycle_land/C768/openHOFX_imsGTS/DA/hofx"

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

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

model_avg1 = np.empty([SimulationDuration])
model_avg2 = np.empty([SimulationDuration])
model_avg3 = np.empty([SimulationDuration])
model_avg4 = np.empty([SimulationDuration])
DA_avg1 = np.empty([SimulationDuration])
DA_avg2 = np.empty([SimulationDuration])
DA_avg3 = np.empty([SimulationDuration])
DA_avg4 = np.empty([SimulationDuration])

model_avg1[:] = np.nan
model_avg2[:] = np.nan
model_avg3[:] = np.nan
model_avg4[:] = np.nan
DA_avg1[:] = np.nan
DA_avg2[:] = np.nan
DA_avg3[:] = np.nan
DA_avg4[:] = 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_gts_"+date+"00.nc"     
    HofxFile21 = mdlHofxPath1+"/letkf_hofx_gts_"+date+"00.nc"

    HofxFile12 = daHofxPath1+"/letkf_hofx_gts_"+date+"06.nc"
    HofxFile22 = mdlHofxPath1+"/letkf_hofx_gts_"+date+"06.nc"

    HofxFile13 = daHofxPath1+"/letkf_hofx_gts_"+date+"12.nc"
    HofxFile23 = mdlHofxPath1+"/letkf_hofx_gts_"+date+"12.nc"

    HofxFile14 = daHofxPath1+"/letkf_hofx_gts_"+date+"18.nc"
    HofxFile24 = mdlHofxPath1+"/letkf_hofx_gts_"+date+"18.nc"
    
    ombg11 = ReadNETCDF(HofxFile11 ,"ombg","totalSnowDepth")
    ombg21 = ReadNETCDF(HofxFile21 ,"ombg","totalSnowDepth")

    ombg12 = ReadNETCDF(HofxFile12 ,"ombg","totalSnowDepth")
    ombg22 = ReadNETCDF(HofxFile22 ,"ombg","totalSnowDepth")

    ombg13 = ReadNETCDF(HofxFile13 ,"ombg","totalSnowDepth")
    ombg23 = ReadNETCDF(HofxFile23 ,"ombg","totalSnowDepth")

    ombg14 = ReadNETCDF(HofxFile14 ,"ombg","totalSnowDepth")
    ombg24 = ReadNETCDF(HofxFile24 ,"ombg","totalSnowDepth")

    qc11 = ReadNETCDF(HofxFile11,"EffectiveQC0","totalSnowDepth")
    qc21 = ReadNETCDF(HofxFile11,"EffectiveQC1","totalSnowDepth")

    qc12 = ReadNETCDF(HofxFile12,"EffectiveQC0","totalSnowDepth")
    qc22 = ReadNETCDF(HofxFile12,"EffectiveQC1","totalSnowDepth")

    qc13 = ReadNETCDF(HofxFile13,"EffectiveQC0","totalSnowDepth")
    qc23 = ReadNETCDF(HofxFile13,"EffectiveQC1","totalSnowDepth")

    qc14 = ReadNETCDF(HofxFile14,"EffectiveQC0","totalSnowDepth")
    qc24 = ReadNETCDF(HofxFile14,"EffectiveQC1","totalSnowDepth")

    for i in range(len(ombg21)):
        if(qc11[i] != 0 or qc21[i] != 0):
            ombg11[i] = float("nan")
            ombg21[i] = float("nan")

    for i in range(len(ombg22)): 
        if(qc12[i] != 0 or qc22[i] != 0):
             ombg12[i] = float("nan")
             ombg22[i] = float("nan")

    for i in range(len(ombg23)): 
        if(qc13[i] != 0 or qc23[i] != 0):
             ombg13[i] = float("nan")
             ombg23[i] = float("nan")
     
    for i in range(len(ombg24)):
        if(qc14[i] != 0 or qc24[i] != 0):
             ombg14[i] = float("nan")
             ombg24[i] = float("nan")

    avg11= np.nanstd(ombg11)
    avg21= np.nanstd(ombg21)
    
    avg12= np.nanstd(ombg12)
    avg22= np.nanstd(ombg22)

    avg13= np.nanstd(ombg13)
    avg23= np.nanstd(ombg23)

    avg14= np.nanstd(ombg14)
    avg24= np.nanstd(ombg24)

    #if (SimDay <= 170) or (SimDay >= 173) :
    DA_avg1[SimDay] = avg11
    model_avg1[SimDay] = avg21

    DA_avg2[SimDay] = avg12
    model_avg2[SimDay] = avg22

    DA_avg3[SimDay] = avg13
    model_avg3[SimDay] = avg23

    DA_avg4[SimDay] = avg14
    model_avg4[SimDay] = avg24

    #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, 250.0)
ax1.plot(xx,model_avg1,'b.',label="Open loop 00Z: stdev(O-F)")
ax1.plot(xx,model_avg2,'k.',label="Open loop 06Z: stdev(O-F)")
ax1.plot(xx,model_avg3,'g.',label="Open loop 12Z: stdev(O-F)")
ax1.plot(xx,model_avg4,'r.',label="Open loop 18Z: stdev(O-F)")
ax1.plot(xx,DA_avg1,'b-',label="DAimsGTS 00Z: stdev(O-F)")
ax1.plot(xx,DA_avg2,'k-',label="DAimsGTS 06Z: stdev(O-F)")
ax1.plot(xx,DA_avg3,'g-',label="DAimsGTS 12Z: stdev(O-F)")
ax1.plot(xx,DA_avg4,'r-',label="DAimsGTS 18Z: 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()
