"""
Code to calculate the percent-full of a reservoir.
This is intended to be applied as a state variable (separate one for each reservoir)
The input CSV file provides metadata about the reservoirs that this will be run for.
Make sure that the state variable name in ResSim matches up exactly with the state variable name in the CSV file. 
Percent full is only re-calculated on certain days.
Developed for the Willamette EIS.
Idea is that a rule in ResSim references the state variable, and minimum flows can be defined as a function of date and percent full. 
"""

from hec.model import PairedValuesExt
from hec.model import RunTimeStep
from hec.heclib.util import HecTime

import os, sys, csv

from functions import fileOpenReadClose, stripOutCommentLines, getCSVDictReader, convertCSVDictReaderToListDict, getListOfDictsFromCSV, getElevationStorageTable, getHecTimeFromRuntimestep

################################################################################
# USER INPUT
DAYS_TO_ADJUST_PCT_FULL = ["01Feb", "16Feb", "01Mar", "16Mar", "01Apr", "16Apr", "01May", "16May", "01Jun"]
INPUT_CSV_FILE = "scripts/PctFull_StateVariable_Input.csv"

################################################################################
# FUNCTION DEFINITIONS

    
def initStateVariable(currentVariable, network):
    # State variable initialization--runs once at the beginning of the compute
    # Make sure no time blocking is occurring--stop the compute if it is happening
    '''
    #In ResSim 3.3, build 140 and later, ResSim handles time-blocking of state variables well.
    # When you retrieve the "previous value" of a State Variable in later time blocks, it gets the proper value 
    # from the previous time block. This chunk of code isn't necessary.
    timeWindowStrComputeBlock = network.getRssRun().getCurrentComputeBlockRunTimeWindow().getTimeWindowString()
    timeWindowStr = network.getRssRun().getRunTimeWindow().getTimeWindowString()
    numStepsComputeBlock = network.getRssRun().getCurrentComputeBlockRunTimeWindow().getNumSteps()
    if timeWindowStrComputeBlock != timeWindowStr:
        msg = "Stopping the Compute--this compute is time blocking"
        msg+= "\nThe model cannot run properly with time blocking"
        msg+= "\nEither reduce your compute time window or"
        msg+= "\nAllocate more memory to ResSim in the HEC-ResSim.config file"
        msg+= "\nFor a 64-bit machine, the line should read 'vmparam -mx6400' if you have that much RAM"
        msg+= "\nTo see how ResSim split up the time window, look at the top of the console log"
        msg+= "\nTime Window:   %s" % timeWindowStr
        msg+= "\nCompute Block: %s" % timeWindowStrComputeBlock
        print msg
        network.printLogMessage(msg)
        network.printErrorMessage(msg)
        # Stop the compute if it is occurring
        return False
    '''
    #Load up information for this reservoir from the external text file
    #Will do a lookup based on the state variable name
    csvFileName = network.makeAbsolutePathFromWatershed(INPUT_CSV_FILE) #Full path to the C-drive or D-drive where the watershed is
    csvListDict = getListOfDictsFromCSV(csvFileName)
    #Loop through each row until we find the reservoir
    foundMatch = False
    for rowDict in csvListDict:
        if rowDict["StateVariableName"].upper() == currentVariable.getName().upper():
            #Found a match, this is it
            foundMatch = True
            break
    if not foundMatch:
        msg = "Could not find a match for the state variable: %s\n in file: %s" %(currentVariable.getName(),csvFileName)
        print msg
        network.printLogMessage(msg)
        network.printErrorMessage(msg)
        return False
    nameResSim = rowDict["ResSimName"] #like "Detroit"
    nameAbbreviation = rowDict["Abbreviation"] #like "DET"
    minConElev = float(rowDict["MinConElev"])
    maxConElev = float(rowDict["MaxConElev"])
    elevStorTable = getElevationStorageTable(nameResSim, network) #input elevation, get storage
    #Save variables for later retrieval
    currentVariable.varPut("nameResSim", nameResSim)
    currentVariable.varPut("elevStorTable", elevStorTable)
    currentVariable.varPut("minConStor", elevStorTable.interpolate(minConElev))
    currentVariable.varPut("maxConStor", elevStorTable.interpolate(maxConElev))
    currentVariable.varPut("passCheck", -1)
    currentVariable.varPut("checkStep", -1)
    return True
        
def runStateVariable(currentVariable, currentRuntimestep, network):
    """
    This will run every time step of the compute. 
    Calculates percent full at the dates of choice.
    """
    # Ensure the SV is only computing once per time step
    passNumber = network.getComputePassCounter()
    passCheck = currentVariable.varGet("passCheck")
    if passNumber > passCheck:
        # A new compute pass-reset the checkStep variable
        currentVariable.varPut("checkStep", -1)
        currentVariable.varPut("passCheck", passNumber)
    checkStep = currentVariable.varGet("checkStep")
    curStep = currentRuntimestep.getStep()
    if curStep <= checkStep: #we went backwards, don't recompute
        return None
    currentVariable.varPut("checkStep", checkStep + 1)
    #Now that we've done that, do the percent full calculations
    curTime = getHecTimeFromRuntimestep(currentRuntimestep)
    curDateStr = curTime.date(14)[:5] #e.g. 02Jun (case-sensitive)
    #Adjust the percent-full calculation only on particular days. 
    if curDateStr in DAYS_TO_ADJUST_PCT_FULL:
        nameResSim = currentVariable.varGet("nameResSim")
        minConStor = currentVariable.varGet("minConStor")
        prevConStor = network.getTimeSeries("Reservoir",nameResSim, "Conservation", "Stor-ZONE").getPreviousValue(currentRuntimestep)
        prevStor = network.getTimeSeries("Reservoir",nameResSim, "Pool", "Stor").getPreviousValue(currentRuntimestep)
        pctFull = (prevStor - minConStor)/(prevConStor - minConStor)*100
    elif curTime.month() >11:
        #Reset it to 0 in December
        pctFull = 0
    else:
        #Just keep the previous value
        pctFull = currentVariable.getPreviousValue(currentRuntimestep)
    pctFull = max(0, min(100, pctFull)) #make sure it is between 0 and 1
    currentVariable.setValue(currentRuntimestep, pctFull)
    return True

def cleanupStateVariable(currentVariable, network):
    #Runs once after the compute is done. 
    return True