import os
import netCDF4
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from datetime import datetime
from datetime import timedelta

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

SpatialFile = "/scratch1/NCEPDEV/da/Azadeh.Gholoubi/control_env/OI_perturbed/LocCsv/Loc-OI-Truth2013-10-01_23-00-00.csv"
SpatialFileDF = pd.read_csv(SpatialFile, index_col="ID")

plotDir = r"/scratch1/NCEPDEV/da/Azadeh.Gholoubi/jedi_experiment2/plots/DA-OpenLoopTimeSeries/"
def get_Location_results (DesiredID, VarList):
    DFColumns = ["Date"]
    for varName in VarList:
        DFColumns.append("%s_OI" % varName)
        DFColumns.append("%s_Truth" % varName)
    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')
    AllSummary = []

    for SimDay in range(SimulationDuration):
        PassedTimeStepNo = (SimDay)
        ThisRestartDate = RestartDate + timedelta(days=PassedTimeStepNo * DeltaTime)
        ThisdateResults=[ThisRestartDate]
        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)
        for varName in VarList:
            ThisIdvariable = ReadNETCDF(DAFile, varName)
            if len(ThisIdvariable.shape) == 3:
                ThisIdvariable=ThisIdvariable[0, 0, DesiredID - 1]
            elif len(ThisIdvariable.shape) == 2:
                ThisIdvariable = ThisIdvariable[0, DesiredID - 1]

            ThisIdVariableopenloop=ReadNETCDF(openloopFile, varName)
            if len(ThisIdVariableopenloop.shape) == 3:
                ThisIdVariableopenloop = ThisIdVariableopenloop[0, 0, DesiredID - 1]
            elif len(ThisIdVariableopenloop.shape) == 2:
                ThisIdVariableopenloop = ThisIdVariableopenloop[0, DesiredID - 1]


            ThisdateResults.append(ThisIdvariable)
            ThisdateResults.append(ThisIdVariableopenloop)

        AllSummary.append(ThisdateResults)
        AllSummary_Dataframe = pd.DataFrame(AllSummary, columns=DFColumns)

    return AllSummary_Dataframe

DesiredIDs=[9821,9624,9635,10283,10708,10930,6482,11145,11768,11984,11826,6394,11940,9977,6482,7518]
# DesiredIDs=[9977,6482,7518]
for ID in DesiredIDs:
    print(ID)
    VarList=["snwdph","smc","weasd","stc","hflx","evap","gflux"]
    result=get_Location_results(ID,VarList)
    This_Long = SpatialFileDF.loc[ID, "Long"]
    This_Lat = SpatialFileDF.loc[ID, "Lat"]
    print(result)

    plt.figure(figsize=(18, 10))
    ax1 = plt.subplot(3, 2, 1)
    ax1.set_title(" Long=%.2lf, Lat=%.2lf\n" % (This_Long, This_Lat), fontsize=18)
    ax1.plot(result["Date"],result["snwdph_OI"],label="DA")
    ax1.plot(result["Date"], result["snwdph_Truth"],label="OpenLoop")
    ax1.set_ylabel("Snow depth[mm]", fontsize=14)
    ax1.legend()

    ax2 = plt.subplot(3, 2, 2)
    ax2.plot(result["Date"],result["hflx_OI"],label="DA")
    ax2.plot(result["Date"], result["hflx_Truth"],label="OpenLoop")
    ax2.set_ylabel("hflx [w/m2]", fontsize=14)
    ax2.legend()

    ax3 = plt.subplot(3, 2, 3)
    ax3.plot(result["Date"],result["weasd_OI"],label="DA")
    ax3.plot(result["Date"], result["weasd_Truth"],label="OpenLoop")
    ax3.set_ylabel("weasd [mm]", fontsize=14)
    ax3.legend()

    ax4 = plt.subplot(3, 2, 4)
    ax4.plot(result["Date"],result["gflux_OI"],label="DA")
    ax4.plot(result["Date"], result["gflux_Truth"],label="OpenLoop")
    ax4.set_ylabel("gflux [w/m2]", fontsize=14)
    ax4.legend()

    ax5 = plt.subplot(3, 2, 5)
    ax5.plot(result["Date"],result["stc_OI"],label="DA")
    ax5.plot(result["Date"], result["stc_Truth"],label="OpenLoop")
    ax5.set_ylabel("stc [k]", fontsize=14)
    ax5.set_xlabel("Date", fontsize=14)
    ax5.legend()

    ax6 = plt.subplot(3, 2, 6)
    ax6.plot(result["Date"],result["evap_OI"],label="DA")
    ax6.plot(result["Date"], result["evap_Truth"],label="OpenLoop")
    ax6.set_ylabel("evap [w/m2]", fontsize=14)
    ax6.set_xlabel("Date", fontsize=14)
    ax6.legend()
    plt.tight_layout()
    plt.savefig(plotDir+"/DAvsOpenLoop%d.png"%ID)
    plt.close()

