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

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

def Computediff(VariName):
    print(VariName)
    DAFilePath="//scratch2/NCEPDEV/land/Youlong.Xia/landDA/cycle_land/DA_era5/mem000/restarts/vector"
    OpenloopFilePath="/scratch2/NCEPDEV/land/Youlong.Xia/landDA/cycle_land/openloop_era5/mem000/restarts/vector"
    RestartDate = "2019-10-01_18-00-00"
    DeltaTime = 1  # day
    SimulationDuration = 244  # 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
    Openloop_stdDF = pd.DataFrame()
    diff_averageDF=pd.DataFrame()
    diff2_rmseDF = pd.DataFrame()

    for SimDay in range(SimulationDuration):
        PassedTimeStepNo = (SimDay)
        ThisRestartDate = RestartDate + timedelta(days=PassedTimeStepNo * DeltaTime)
        print(ThisRestartDate)
        thisyear, thismonth, thisday, thishr, thismin, thissec = ThisRestartDate.year, ThisRestartDate.month, ThisRestartDate.day, ThisRestartDate.hour, ThisRestartDate.minute, ThisRestartDate.second
        DAFile = DAFilePath+ "/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)
        DA_array = ReadNETCDF(DAFile , VariName)
        Openloop_array = ReadNETCDF(openloopFile , VariName)
        if len(DA_array.shape) == 3:
            DA=DA_array[0, 0, :]
            Openloop=Openloop_array[0, 0, :]
            Diff = DA-Openloop

        elif len(DA_array.shape) == 2:
            DA= DA_array[0, :]
            Openloop=Openloop_array[0, :]
            Diff = DA-Openloop

        else:
            print("Check the shape of array!")
            exit()
        Openloop_stdDF["%d_Openloop"% SimDay]=Openloop

        diff_averageDF["%d_diff" % SimDay] = Diff

        diff2_rmseDF["%d_diff2" % SimDay] = Diff**2

    Openloop_stdDF["stdOfOpenloop"]=Openloop_stdDF.std(axis=1)

    diff_averageDF["MeanOFDiffs"] = diff_averageDF.mean(axis=1)

    diff2_rmseDF["MeanOFDiff2"] = diff2_rmseDF.mean(axis=1)
    diff2_rmseDF["RMSEOFDiff2"] = np.sqrt(diff2_rmseDF["MeanOFDiff2"])

    diff2_rmseDF["NormalizedRMSE"]=diff2_rmseDF["RMSEOFDiff2"]/Openloop_stdDF["stdOfOpenloop"]
    diff2_rmseDF["MeanOFDiffs"]=diff_averageDF["MeanOFDiffs"]

    Long = ReadNETCDF(SpatialPath, "longitude")
    Lat = ReadNETCDF(SpatialPath, "latitude")
    diff2_rmseDF["Lat"] = Lat
    diff2_rmseDF["Long"] = Long
    diff2_rmseDF["ID"] = np.arange(1, len(diff2_rmseDF) + 1)
    # Openloop_stdDF.to_csv("test.csv", index=False)

    # maxvalue = diff2_rmseDF["MeanOFDiffs"].max()
    # maxvalueID=np.argmax(diff2_rmseDF["MeanOFDiffs"])
    # print("maxMeanID=%s"%maxvalueID)

    #now we want to mask ploes locations that are where veg_type==15 ( glaciers  - we're not doing any DA here)
    veg_type_arryay = ReadNETCDF(SpatialPath, "vegetation_category")
    CurrentMeanResult=diff2_rmseDF.MeanOFDiffs
    MaskedMeanResult=np.where(veg_type_arryay==15,0,CurrentMeanResult)
    maxvalue = MaskedMeanResult.max()
    maxvalue=round(maxvalue,3)
    minvalue = MaskedMeanResult.min()
    minvalue=round(minvalue,3)

    CurrentRMSEResult = diff2_rmseDF.RMSEOFDiff2
    MaskedRMSEResult = np.where(veg_type_arryay == 15, 0, CurrentRMSEResult)
    maxRMSE = MaskedRMSEResult.max()
    maxRMSE = round(maxRMSE, 3)

    CurrentNormalizedRMSEResult = diff2_rmseDF.NormalizedRMSE
    MaskedNormalizedRMSEResult = np.where(veg_type_arryay == 15, 0, CurrentNormalizedRMSEResult)
    maxNormRMSE = MaskedNormalizedRMSEResult.max()
    maxNormRMSE=round(maxNormRMSE,3)
    minNormRMSE =MaskedNormalizedRMSEResult.min()
    minNormRMSE= round(minNormRMSE,3)
    
    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'))
    cs = plt.scatter(x=diff2_rmseDF.Long, y=diff2_rmseDF.Lat, c=MaskedMeanResult, cmap="bwr", marker=',', s=10, vmax=0.2, vmin=-0.2, transform=ccrs.PlateCarree())
    cb = plt.colorbar(cs, orientation='horizontal', shrink=0.5, pad=.04)
    cb.set_label("%s DA-Openloop Mean"%VariName, fontsize=12)
    plt.title("Max=%s Min=%s" % (maxvalue, minvalue), fontsize=12)    
    plt.savefig(os.path.join(PlotPath, "%s_AverageDA-Openloop.png"%VariName))
    plt.close()

    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'))
    cs = plt.scatter(x=diff2_rmseDF.Long, y=diff2_rmseDF.Lat, c=MaskedRMSEResult, cmap="OrRd", marker=',', s=10, vmax=0.3,
                     vmin=0.0, transform=ccrs.PlateCarree())
    cb = plt.colorbar(cs, orientation='horizontal', shrink=0.5, pad=.04)
    cb.set_label("%s DA-Openloop RMSE"%VariName, fontsize=12)
    plt.title("Max=%s" %maxRMSE, fontsize=12)
    plt.savefig(os.path.join(RMSEplot, "%s_RMSE_DA-OpenLoop.png"%VariName))
    plt.close()

SpatialPath = "/scratch2/NCEPDEV/land/data/forcing/era5/static/ufs-land_C96_static_fields.nc"
PlotPath= "/scratch2/NCEPDEV/land/Youlong.Xia/landDA/plots/C96/Average_DA-OpenLoopDiffsnap/"
RMSEplot= "/scratch2/NCEPDEV/land/Youlong.Xia/landDA/plots/C96/RMSE_DA-OpenLoopDiffsnap/"


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