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

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

def Computediff(VariName):
    print(VariName)
    DAFilePath1="/scratch1/NCEPDEV/stmp4/Youlong.Xia/landDA/cycle_land/C768/DA_GHCN_2015to2020/mem000/restarts/vector"
    DAFilePath2="/scratch2/NCEPDEV/stmp3/Youlong.Xia/landDA/cycle_land/C768/DA_imsGHCN_2015to2020/mem000/restarts/vector"
    OpenloopFilePath="/scratch1/NCEPDEV/stmp4/Youlong.Xia/landDA/cycle_land/C768/openloop_ghcnHOFX_2015to2020/mem000/restarts/vector"
    imsFilePath="/scratch2/NCEPDEV/land/data/DA/snow_ice_cover/IMS/CalSnowCover/4km/C768"

    RestartDate = "2019-10-01_00-00-00"
    DeltaTime = 1  # day
    SimulationDuration = 365  # days

    RestartDate = datetime.strptime(RestartDate, '%Y-%m-%d_%H-%M-%S')
#we want to create 3 dataframe: 1) Truth_stdDF: the rows are locations and columns are truth value at each time step and the last column will be the standard deviation of truth, 2) diff_averageDF: the rows are same as previous and each column
#is difference between OI and Truth at each time step, the last column will be average of diffs at each row (location), 3) diff2_rmseDF: rows same as previous but each column will be diff**2 and the last column will be RMSE that is sqrt of average diff**2

#   defined the basin mask
    basin = ReadNETCDF(BasinPath, "basins")
    DA1 = np.empty([len(basin)])
    DA2 = np.empty([len(basin)])
    Openloop = np.empty([len(basin)])
    IMS = np.empty([len(basin)])

    DA1[:] = np.nan
    DA2[:] = np.nan
    Openloop[:] = np.nan 
    IMS[:] = np.nan

    IMS_avg1 = np.empty([SimulationDuration]) 
    model_avg1 = np.empty([SimulationDuration])
    DA_avg11 = np.empty([SimulationDuration])
    DA_avg21 = np.empty([SimulationDuration])
    model_avg1[:] = np.nan
    DA_avg11[:] = np.nan
    DA_avg21[:] = np.nan
    IMS_avg1[:] = np.nan

    IMS_avg2 = np.empty([SimulationDuration])
    model_avg2 = np.empty([SimulationDuration])
    DA_avg12 = np.empty([SimulationDuration])
    DA_avg22 = np.empty([SimulationDuration])
    model_avg2[:] = np.nan
    DA_avg12[:] = np.nan
    DA_avg22[:] = np.nan
    IMS_avg2[:] = np.nan

    IMS_avg3 = np.empty([SimulationDuration])
    model_avg3 = np.empty([SimulationDuration])
    DA_avg13 = np.empty([SimulationDuration])
    DA_avg23 = np.empty([SimulationDuration])
    model_avg3[:] = np.nan
    DA_avg13[:] = np.nan
    DA_avg23[:] = np.nan
    IMS_avg3[:] = np.nan
 
    IMS_avg4 = np.empty([SimulationDuration])
    model_avg4 = np.empty([SimulationDuration])
    DA_avg14 = np.empty([SimulationDuration])
    DA_avg24 = np.empty([SimulationDuration])
    model_avg4[:] = np.nan
    DA_avg14[:] = np.nan
    DA_avg24[:] = np.nan
    IMS_avg4[:] = np.nan

    xx = np.empty([SimulationDuration])
    
    for SimDay in range(SimulationDuration):
        #print(SimDay)
        PassedTimeStepNo = (SimDay)
        ThisRestartDate = RestartDate + timedelta(days=PassedTimeStepNo * DeltaTime)
        print(ThisRestartDate)
        xx[SimDay] = SimDay

        thisyear, thismonth, thisday, thishr, thismin, thissec = ThisRestartDate.year, ThisRestartDate.month, ThisRestartDate.day, ThisRestartDate.hour, ThisRestartDate.minute, ThisRestartDate.second
          
# Read DA and open loop data files

        DAFile1 = DAFilePath1+ "/ufs_land_restart_back.%4d-%02d-%02d_%02d-%02d-%02d.nc" % (thisyear, thismonth, thisday, thishr, thismin, thissec)
        DAFile2 = DAFilePath2+ "/ufs_land_restart_back.%4d-%02d-%02d_%02d-%02d-%02d.nc" % (thisyear, thismonth, thisday, thishr, thismin, thissec)
        openloopFile = OpenloopFilePath+ "/ufs_land_restart_back.%4d-%02d-%02d_%02d-%02d-%02d.nc" % (thisyear, thismonth, thisday, thishr, thismin, thissec)
        imsFile = imsFilePath+ "/IMSscf.C768.4km.%4d%02d%02d.nc" % (thisyear, thismonth, thisday)
        
        DA_array1 = ReadNETCDF(DAFile1, VariName)
        DA_array2 = ReadNETCDF(DAFile2, VariName)
        Openloop_array = ReadNETCDF(openloopFile, VariName)
        ims_array1 = ReadNETCDF(imsFile, VariName)
        ims_array = 0.01*ims_array1

        total=0.0
        tims1=0.0        
        ta11=0.0
        ta21=0.0
        tm1=0.0
  
        tims2=0.0
        ta12=0.0
        ta22=0.0
        tm2=0.0

        tims3=0.0
        ta13=0.0
        ta23=0.0
        tm3=0.0

        tims4=0.0
        ta14=0.0
        ta24=0.0
        tm4=0.0

        for i in range(len(basin)):
            if basin[i] == 2434:
                # IMS snow and model snow case
                if(ims_array[i] >= 0.0):
                    total=total+1.0
                if(ims_array[i] > 0.0):
                    tims1=tims1+1.0
                if(ims_array[i] > 0.0 and DA_array1[i] > 0.0):
                    ta11=ta11+1.0
                if(ims_array[i] > 0.0 and DA_array2[i] > 0.0):
                    ta21=ta21+1.0
                if(ims_array[i] > 0.0 and Openloop_array[i] > 0.0):
                    tm1=tm1+1.0

                # IMS and model snow >=0.7 case
                if(ims_array[i] >= 0.7 and ims_array[i] < 1.0):
                    tims2=tims2+1.0
                if(ims_array[i] >= 0.7 and ims_array[i] <1.0 and DA_array1[i] >= 0.7):
                    ta12=ta12+1.0
                if(ims_array[i] >= 0.7 and ims_array[i] < 1.0 and DA_array2[i] >= 0.7):
                    ta22=ta22+1.0
                if(ims_array[i] >= 0.7 and ims_array[i] < 1.0 and  Openloop_array[i] >= 0.7):
                    tm2=tm2+1.0

                # IMS and model snow within 0.3-0.7 case
                if(ims_array[i] >= 0.3 and ims_array[i] <=0.7):
                    tims3=tims3+1.0
                if(ims_array[i] >= 0.3 and ims_array[i] <= 0.7 and DA_array1[i] >= 0.3 and DA_array1[i] <=0.7):
                    ta13=ta13+1.0
                if(ims_array[i] >= 0.3 and ims_array[i] <= 0.7 and DA_array2[i] >= 0.3 and DA_array2[i] <= 0.7):
                    ta23=ta23+1.0
                if(ims_array[i] >= 0.3 and ims_array[i] <=0.7 and Openloop_array[i] >= 0.3 and Openloop_array[i] <= 0.7):
                   tm3=tm3+1 

                # IMS and model snow <=0.3 case
                if(ims_array[i] > 0.0 and ims_array[i] <= 0.3):
                   tims4=tims4+1  
                if(ims_array[i] > 0.0 and ims_array[i] <=0.3 and DA_array1[i] > 0.0 and DA_array1[i] <= 0.3):
                    ta14=ta14+1.0
                if(ims_array[i] > 0.0 and ims_array[i] <= 0.3 and DA_array2[i] > 0.0 and DA_array2[i] <= 0.3):
                    ta24=ta24+1.0
                if(ims_array[i] > 0.0 and ims_array[i] <= 0.3 and Openloop_array[i] > 0.0 and Openloop_array[i] <= 0.3):
                    tm4=tm4+1.0
           
        model_avg1[SimDay] = 100.0*(tm1/total)
        DA_avg11[SimDay] = 100.0*(ta11/total)
        DA_avg21[SimDay] = 100.0*(ta21/total)
        IMS_avg1[SimDay] = 100.0*(tims1/total)

        model_avg2[SimDay] = 100.0*(tm2/total)
        DA_avg12[SimDay] = 100.0*(ta12/total)
        DA_avg22[SimDay] = 100.0*(ta22/total)
        IMS_avg2[SimDay] = 100.0*(tims2/total)

        model_avg3[SimDay] = 100.0*(tm3/total)
        DA_avg13[SimDay] = 100.0*(ta13/total)
        DA_avg23[SimDay] = 100.0*(ta23/total)
        IMS_avg3[SimDay] = 100.0*(tims3/total)

        model_avg4[SimDay] = 100.0*(tm4/total)
        DA_avg14[SimDay] = 100.0*(ta14/total)
        DA_avg24[SimDay] = 100.0*(ta24/total)
        IMS_avg4[SimDay] = 100.0*(tims4/total)

    plt.figure(figsize=(15,8))

    ax1 = plt.subplot(2, 2, 1)
    ax1.set_ylim(0.0, 100.0)
    ax1.plot(xx,DA_avg11,'r-',label="GHCN DA")
    ax1.plot(xx,model_avg1,'b-',label="openloop")
    ax1.plot(xx,DA_avg21,'g-',label="IMS+GHCN DA")
    ax1.plot(xx,IMS_avg1,'k-',label="IMS") 
    ax1.set_ylabel("Pct of Grids", fontsize=15)
    ax1.set_xlabel("Days since 20191001", fontsize=15)
    ax1.set_title("(a) IMS and Model FSC >0", pad=20, fontsize=15)
    ax1.legend()

    ax1 = plt.subplot(2, 2, 2)
    ax1.set_ylim(0.0, 10.0)
    ax1.plot(xx,DA_avg12,'r-',label="GHCN DA")
    ax1.plot(xx,model_avg2,'b-',label="openloop")
    ax1.plot(xx,DA_avg22,'g-',label="IMS+GHCN DA")
    ax1.plot(xx,IMS_avg2,'k-',label="IMS")
    ax1.set_ylabel("Pct of Grids", fontsize=15)
    ax1.set_xlabel("Days since 20191001", fontsize=15)
    ax1.set_title("(b) IMS and Model 70=<FSC<100", pad=20, fontsize=15)
    ax1.legend()

    ax1 = plt.subplot(2, 2, 3)
    ax1.set_ylim(0.0, 10.0)
    ax1.plot(xx,DA_avg13,'r-',label="GHCN DA")
    ax1.plot(xx,model_avg3,'b-',label="openloop")
    ax1.plot(xx,DA_avg23,'g-',label="IMS+GHCN DA")
    ax1.plot(xx,IMS_avg3,'k-',label="IMS")
    ax1.set_ylabel("Pct of Grids", fontsize=15)
    ax1.set_xlabel("Days since 20191001", fontsize=15)
    ax1.set_title("(c) IMS and Model 30=<FSC<=70", pad=20, fontsize=15)
    ax1.legend()

    ax1 = plt.subplot(2, 2, 4)
    ax1.set_ylim(0.0, 10.0)
    ax1.plot(xx,DA_avg14,'r-',label="GHCN DA")
    ax1.plot(xx,model_avg4,'b-',label="openloop")
    ax1.plot(xx,DA_avg24,'g-',label="IMS+GHCN DA")
    ax1.plot(xx,IMS_avg4,'k-',label="IMS")
    ax1.set_ylabel("Pct of Grids", fontsize=15)
    ax1.set_xlabel("Days since 20191001", fontsize=15)
    ax1.set_title("(c) IMS and Model 0<FSC<=30", pad=20, fontsize=15)
    ax1.legend()

    plt.tight_layout()
    plt.savefig(os.path.join(PlotPath, fname))
    plt.close()

BasinPath = "/scratch2/NCEPDEV/land/data/evaluation/basins/C768/GRDC_C768.nc"
PlotPath= "./plots/C768_multiYear/"
fname = "Yellow_river_basin_percentageGridCountNumber_snowCaseDA_2020.png"

VarList=["snow_cover_fraction"]
for var in VarList:
    Computediff(var)
