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="/scratch1/NCEPDEV/da/Azadeh.Gholoubi/jedi_experiment2/cycleDA/output/restarts/vector/"
    OpenloopFilePath="/scratch1/NCEPDEV/da/Azadeh.Gholoubi/Openloop/cycleDA/output/restarts/vector/"
    RestartDate = "2015-09-01_18-00-00"
    DeltaTime = 1  # day
    SimulationDuration = 243  # days
    RestartDate = datetime.strptime(RestartDate, '%Y-%m-%d_%H-%M-%S')

    SummaryDF = 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)
        Myarray1 = ReadNETCDF(DAFile, VariName)
        Myarray2 = ReadNETCDF(openloopFile, VariName)
        veg_type_arryay = ReadNETCDF(SpatialPath, "vegetation_category")
        if len(Myarray1.shape) == 3:
            DA=Myarray1[0, 0, :]
            openloop=Myarray2[0, 0, :]
            Diff = DA-openloop
            Diff = np.where(openloop== 0, np.nan, Diff)
            Diff = np.where(veg_type_arryay == 15, np.nan, Diff)
        elif len(Myarray1.shape) == 2:
            DA= Myarray1[0, :]
            openloop=Myarray2[0, :]
            Diff = DA-openloop
            Diff = np.where(openloop== 0, np.nan, Diff)
            Diff=np.where(veg_type_arryay == 15, np.nan, Diff)

        else:
            print("Check the shape of array!")
            exit()
        SummaryDF["%d_diff" % SimDay] = Diff

    SummaryDF["MeanOFDiffs"] = SummaryDF.mean(axis=1)
    maxvalue =  SummaryDF["MeanOFDiffs"].max()
    minvalue =  SummaryDF["MeanOFDiffs"].min()
    limit=maxvalue
    if abs(minvalue)>maxvalue:
        limit=abs(minvalue)



    plt.figure(figsize=(12, 8))
    plt.hist(x= SummaryDF["MeanOFDiffs"],  range=[-limit, limit], bins=101, facecolor='gray', align='mid')
    plt.ylabel('Probability density',fontsize=12)
    plt.xlabel('%s'%VariName,fontsize=12)
    plt.title("%s Mean(DA - openloop)"%VariName, fontsize=16)
    plt.savefig(os.path.join(PlotPath, "%s_AverageDA-OpenLoopSDpresent.png"%VariName))
    plt.close()

PlotPath= "/scratch1/NCEPDEV/da/Azadeh.Gholoubi/jedi_experiment2/plots/HistogramAverage_DA-OpenLoopDiffsnap/"
SpatialPath = "/scratch1/NCEPDEV/da/Azadeh.Gholoubi/control_env/static/ufs-land_C96_static_fields.nc"
PointsDF=pd.DataFrame()
PointsDF["ID"]=np.arange(1,len(PointsDF)+1)

VarList=["snwdph","weasd"]

for var in VarList:
    Computediff(var)
