# coding: utf-8

# <h1>Basic test of the wflow BMI interface + reservoir

# In[10]:
import os
import os.path
import getopt
import sys
import getopt

from wflow.wf_DynamicFramework import *

import pcraster as pcr
import wflow.wflow_bmi as bmi
import logging
import wflow.wflow_adapt as adapter
import calendar
import os
import numpy as np
import configparser

# reload(bmi)

""" Todo: include error handling:
        - When WFlow_ids in the mapping section of config are not in the map
        - When RTC doesn't have data on the WFlow simulation period
        - When a RTC id in the id mapping section is not part of the model
        - ...
    Todo: add logging
"""


def ConfigSectionMap(section):
    dict1 = {}
    options = Config.options(section)
    for option in options:
        try:
            dict1[option] = Config.get(section, option)
            if dict1[option] == -1:
                DebugPrint("skip: %s" % option)
        except:
            print(("exception on %s!" % option))
            dict1[option] = None
    return dict1


def gettimestepfname(name, path, timestep):
    """
    Get the pcraster filename fro this step
    :param name:
    :param path:
    :param timestep:
    :return:
    """

    below_thousand = timestep % 1000
    above_thousand = timestep / 1000
    fname = str(name + "%0" + str(8 - len(name)) + ".f.%03.f") % (
        above_thousand,
        below_thousand,
    )
    fname = os.path.join(path, fname)

    return fname


########################################################################
## Process command-line options                                        #
########################################################################

argv = sys.argv

try:
    opts, args = getopt.getopt(argv[1:], "c:w:I:")
    print(opts)
except getopt.error as msg:
    print("cannot parse commandline")
    sys.exit(2)

if not opts:
    print("cannot parse commandline")
    sys.exit(2)

for o, a in opts:
    if o == "-c":
        configfile = a
    if o == "-w":
        cur_dir = os.path.abspath(a)
    if o == "-I":
        IniFile = a

Config = configparser.ConfigParser()
# inifile = Config.read('c:\FEWS\SI-WAMI\SI-WAMI\Modules\RTC\wflow_rtctools.ini')
inifile = Config.read(configfile)
########################################################################
## Parse ini-file                                                      #
########################################################################

# cur_dir = os.getcwd()
Config.sections()
os.chdir(cur_dir)
dir_rtc = os.path.join(cur_dir, (ConfigSectionMap("model")["dir_rtc_model"]))
dir_wflow = os.path.join(cur_dir, (ConfigSectionMap("model")["dir_wflow_model"]))
Bin_RTC = os.path.join(cur_dir, (ConfigSectionMap("RTC wrapper engine")["bin_rtc"]))
inflow_map = os.path.join(dir_wflow, (ConfigSectionMap("id_maps")["samplemap_in"]))
outflow_map = os.path.join(dir_wflow, (ConfigSectionMap("id_maps")["samplemap_out"]))
ldd_map = os.path.join(dir_wflow, (ConfigSectionMap("ldd")["ldd_res"]))

# id's wflow and RTC reservoir inflow points
id_in_rtc = []
id_in_wflow = list(ConfigSectionMap("id_mapping_inflow"))
for index in range(len(id_in_wflow)):
    id_in_rtc.append(ConfigSectionMap("id_mapping_inflow")[id_in_wflow[index]])

# id's wflow and RTC reservoir outflow points
id_out_rtc = []
id_out_wflow = list(ConfigSectionMap("id_mapping_outflow"))
for index in range(len(id_out_wflow)):
    id_out_rtc.append(ConfigSectionMap("id_mapping_outflow")[id_out_wflow[index]])

########################################################################
## Initialize models                                                   #
########################################################################

# In[]: Initialize the RTC-Tools model
os.chdir(Bin_RTC)

from wflow.wrappers.rtc.wrapperExtended import BMIWrapperExtended

# RTC_model = BMIWrapperExtended(engine=os.path.join(Bin_RTC,"RTCTools_BMI"))
RTC_model = BMIWrapperExtended(engine=os.path.join(Bin_RTC, "RTCTools_BMI"))
print("RTCmodel", Bin_RTC, RTC_model)
RTC_model.initialize("..")


# In[]: Initialize the WFlow model
os.chdir(dir_wflow)
LA_model = bmi.wflowbmi_csdms()
LA_model.initialize((IniFile), loglevel=logging.DEBUG)

# now get the forcings that wflow expects
# The bmi is such that you can get the input variables and the output variables. However, the
# input variable list also contains the in/out variables. So to
# get the input only we subtract the two lists.
invars = LA_model.get_input_var_names()
outvars = LA_model.get_output_var_names()
inputmstacks = list(set(invars) - set(outvars))


# In[]: Investigate start time, end time and time step of both models

print("WFlow:")
LA_dt = LA_model.get_time_step()

# LA_start = LA_model.get_start_time()
timeutc = adapter.getStartTimefromRuninfo("inmaps/runinfo.xml")
print(timeutc)
LA_start = calendar.timegm(timeutc.timetuple())
timeutc = adapter.getEndTimefromRuninfo("inmaps/runinfo.xml")
LA_end = calendar.timegm(timeutc.timetuple())
# LA_end = LA_model.get_end_time()
print(LA_dt)
print(timeutc)
print(LA_start)
print(LA_end)

print("RTC-Tools")
RTC_dt = RTC_model.get_time_step()
RTC_start = RTC_model.get_start_time()
RTC_end = RTC_model.get_end_time()

print(RTC_dt)
print(RTC_start)
print(RTC_end)

if LA_start != RTC_start:
    print("Error: start time of both models is not identical !!!")

if LA_dt != RTC_dt:
    print("Error: time step of both models is not identical !!!")


# In[]:  Read and map reservoir inflow and outflow locations
Reservoir_inflow = pcr.pcr2numpy(
    pcr.scalar(pcr.readmap(os.path.abspath(inflow_map))), np.NaN
)
Reservoir_outflow = pcr.pcr2numpy(
    pcr.scalar(pcr.readmap(os.path.abspath(outflow_map))), np.NaN
)

inflow_list = list(np.unique(Reservoir_inflow[~np.isnan(Reservoir_inflow)]))
outflow_list = list(np.unique(Reservoir_outflow[~np.isnan(Reservoir_outflow)]))


# In[]:  Overwrite TopoLdd with modified version
ldd = pcr.pcr2numpy(
    pcr.readmap(os.path.join(dir_wflow, ldd_map)), np.NaN
).astype(np.float32)
LA_model.set_value("TopoLdd", np.flipud(ldd).copy())


########################################################################
## Run models                                                          #
########################################################################

t = LA_start
timecounter = 0
inmstackbuf = {}  # Keep track of input mapatsatck by name
while t < min(LA_end, RTC_end):
    # first read forcing mapstacks (set in API section) and give to the model
    for thisstack in inputmstacks:
        toread = gettimestepfname(
            thisstack, os.path.join(dir_wflow, "inmaps"), timecounter + 1
        )
        inmstackbuf[thisstack] = np.flipud(
            pcr.pcr2numpy(pcr.scalar(pcr.readmap(os.path.abspath(toread))), -999.0)
        ).copy()
        LA_model.set_value(thisstack, inmstackbuf[thisstack])

    print("calculation timestep = " + str(timecounter))

    # Get the inflow from the wflow model runoff map and map
    inflowQ = np.flipud(LA_model.get_value("SurfaceRunoff")).copy()

    # Map the sum of WFlow Inflow to RTC
    for idx, wflow_id in enumerate(id_in_wflow):
        value = np.ndarray(shape=(1, 1), dtype=float, order="F")
        value[0][0] = np.sum(inflowQ[np.where(Reservoir_inflow == int(wflow_id))])
        rtc_id = id_in_rtc[id_in_wflow.index(str(wflow_id))]
        print(rtc_id + " = " + str(value[0][0]))
        RTC_model.set_value(rtc_id, value)

    # run the RTC-Tools model
    RTC_model.update(-1.0)

    # Extract RTC outflow and add to on WFlow 'IF' (abstractions)
    inflowfield = inmstackbuf["IF"]
    for idx, wflow_id in enumerate(id_out_wflow):
        rtc_id = id_out_rtc[id_out_wflow.index(str(wflow_id))]
        Qout = RTC_model.get_var(rtc_id)
        Qout = 300.0
        if np.isfinite(Qout):  # no nan's into wflow
            inflowfield[Reservoir_outflow == int(wflow_id)] += Qout

    LA_model.set_value("IF", np.flipud(inflowfield).copy())
    # This not not bmi but needed to update the kinematic wave reservoir
    LA_model.update()
    t += LA_dt
    timecounter += 1

LA_model.finalize()
RTC_model.finalize()