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_EuropeGTS_stationNumber_timeseries_WY2020.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])

# North America [-155 - -60oW, 40 - 70oN]
# Europe [0 -45E, 40-70N]
# Russia [45-150, 40-70oN]

latmin = 40.0
latmax = 70.0
lonmin = 0.0
lonmax = 45.0

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")

    latout11 = ReadNETCDF(HofxFile11 ,"MetaData", "latitude")
    lonout21 = ReadNETCDF(HofxFile11 ,"MetaData", "longitude")

    latout12 = ReadNETCDF(HofxFile12 ,"MetaData", "latitude")
    lonout22 = ReadNETCDF(HofxFile12 ,"MetaData", "longitude")

    latout13 = ReadNETCDF(HofxFile13 ,"MetaData", "latitude")
    lonout23 = ReadNETCDF(HofxFile13 ,"MetaData", "longitude")

    latout14 = ReadNETCDF(HofxFile14 ,"MetaData", "latitude")
    lonout24 = ReadNETCDF(HofxFile14 ,"MetaData", "longitude")

    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")

    k11=0
    k21=0
    k12=0
    k22=0
    k13=0
    k23=0
    k14=0
    k24=0

    # fill values for a region
    for k in range(len(latout11)):
        if (latout11[k]>=latmin) & (latout11[k]<=latmax) & (lonout21[k]>=lonmin) & (lonout21[k]<=lonmax):
            k11=k11+1
            #if(qc11[k] == 0 and qc12[k] == 0):
            if(qc11[k] == 0):
               k21=k21+1
 
    for k in range(len(latout12)):
        if (latout12[k]>=latmin) & (latout12[k]<=latmax) & (lonout22[k]>=lonmin) & (lonout22[k]<=lonmax):
            k12=k12+1
            if(qc12[k] == 0 and qc22[k] == 0):
                k22=k22+1

    for k in range(len(latout13)):
       if (latout13[k]>=latmin) & (latout13[k]<=latmax) & (lonout23[k]>=lonmin) & (lonout23[k]<=lonmax):
            k13=k13+1
            if(qc13[k] == 0 and qc23[k] == 0):
                k23=k23+1

    for k in range(len(latout14)):
       if (latout14[k]>=latmin) & (latout14[k]<=latmax) & (lonout24[k]>=lonmin) & (lonout24[k]<=lonmax):
            k14=k14+1
            if(qc14[k] == 0 and qc24[k] == 0):
                k24=k24+1

    DA_avg1[SimDay] = k21
    model_avg1[SimDay] = k11

    DA_avg2[SimDay] = k22
    model_avg2[SimDay] = k12

    DA_avg3[SimDay] = k23
    model_avg3[SimDay] = k13

    DA_avg4[SimDay] = k24
    model_avg4[SimDay] = k14

# plot global averaged snow depth time series
plt.figure(figsize=(12,8))
ax1 = plt.subplot(1, 1, 1)
ax1.set_ylim(0, 3000)
ax1.plot(xx,model_avg1,'b.',label="00Z Station Total")
ax1.plot(xx,model_avg2,'k.',label="06Z Station Total")
ax1.plot(xx,model_avg3,'g.',label="12Z Station Total")
ax1.plot(xx,model_avg4,'r.',label="18Z Station Total")
ax1.plot(xx,DA_avg1,'b-',label="00Z Station Used")
ax1.plot(xx,DA_avg2,'k-',label="06Z Station Used")
ax1.plot(xx,DA_avg3,'g-',label="12Z Station Used")
ax1.plot(xx,DA_avg4,'r-',label="18Z Station Used")
ax1.set_ylabel("GTS Station Number", fontsize=20)
ax1.set_title("Europe", fontsize=20)
ax1.set_xlabel("Simulation Days Since 1 October 2019", fontsize=20)
ax1.legend()
plt.tight_layout()
plt.savefig(fname)
plt.close()
