"""
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

################################################################################
# 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 fileOpenReadClose(txtFileName):
    """Open a text file and reads the lines, storing as a list"""
    if not os.path.exists(txtFileName): 
        errMsg = "File does not exist: \n%s" %txtFileName
        raise AssertionError, errMsg
    txtFile = open(txtFileName, 'r')
    txtLines = txtFile.readlines()
    txtFile.close()
    return txtLines
    
def stripOutCommentLines(txtLines):
    """
    Strip out comment lines from list of strings (assumes comment character is #).
    Returns a list of strings."""
    newLines = []
    for txtLine in txtLines :
            line = txtLine.strip()
            if len(line) == 0 : pass               # skip blank lines
            elif line.count(",") == len(line) : pass # CSV file with only commas
            elif line.startswith("#") : pass       # commented line
            elif line.startswith('"#') : pass      # commented line
            elif line.startswith('",') : pass      # commented line
            else : 
                #strip out any inline comments
                if "#" in line:
                    cmtIdx = line.find("#") #returns -1 if not found
                    line = line[:cmtIdx].strip()
                newLines.append(line)
    return newLines
    
def getCSVDictReader(txtLines):
    """
    Gets a CSV DictReader object.
    You can iterate on this object, but not get a specific line.
    Note that the .fieldnames parameter is not available until you've iterated once
    by using .next() or going inside a for-loop.
    If you want a list, use :func:`convertCSVDictReaderToListDict`
    
    :param list txtLines: list of strings (comment lines should already be 
                          erased, and first line should be headers)
    :return list: A list of dictionary objects (one dictionary for each line)
                  keys of the dictionary are the column header names
    """
    return csv.DictReader(txtLines)
    
def convertCSVDictReaderToListDict(csvDictObj):
    """
    Takes a csv.DictReader object and "listifies" it so you don't have to iterate on it
    """
    outList = []
    for d in csvDictObj: outList.append(d)
    return outList
    
def getListOfDictsFromCSV(txtFile):
    """
    Returns a list of dictionaries (one dictionary for each row in the file).
    The keys of the dictionaries are the column header names. 
    
    :param str txtFile: .csv file with the column headings
    """
    lines = fileOpenReadClose(txtFile)
    lines = stripOutCommentLines(lines)
    csvDict = getCSVDictReader(lines)
    listDict = convertCSVDictReaderToListDict(csvDict)
    return listDict

def getElevationStorageTable(reservoirName, network):
    """
    Returns a table that inputs elevation and returns storage
    
    :param str reservoirName: The reservoir name as it exists in ResSim
    :param RssSystem network:
    :return elevStorTable: Can use .interpolate on this object
    :rtype: PairedValuesExt
    """
    reservoirObj = network.findReservoir(reservoirName)
    storageElement = reservoirObj.getStorageFunction()
    storageTable = storageElement.getElevationStorageValues()
    pairedDataContainer = storageTable.getPairedDataContainer()
    elevStorTable = PairedValuesExt() 
    elevStorTable.setData(pairedDataContainer)
    return elevStorTable

    elevStorTable = functions.getElevationStorageTable("Mica", network)
    storElevTable = functions.getStorageElevationTable("Mica", network)
    fullPoolStor = elevStorTable.interpolate(fullPoolElevations["Mica"])

def getHecTimeFromRuntimestep(rtsObj):
    """
    Returns an HecTime object that reflects the true time of the step (2400 hours).
    When Running ResSim daily modeling, the .getHecTime() method does not work
    as it should--it returns an HecTime object with daily granularity, which isn't
    good when trying to retrieve data from other timeseries. This function returns
    the current day with a granularity of minutes.
    
    This function should be used exclusively in the place of .getHecTime()
    
    :param RunTimeStep rtsObj: 
    """
    currentDate = rtsObj.getHecTime().clone()
    # If the hecTime object returned is daily granularity, subtract a day due to midnight shift
    if currentDate.getTimeGranularity() == HecTime.DAY_INCREMENT: # days
        currentDate.subtractDays(1) # Midnight shift
    currentMonth = currentDate.month()  
    currentDay = currentDate.day()
    currentYear = currentDate.year()
    # If the daily hecTime object does not have a granularity of minutes, 
    # Set the current date to a granularity of minutes (2400 hours at the end of the day)
    if currentDate.getTimeGranularity() == HecTime.DAY_INCREMENT: 
        currentDate.setTimeGranularity(HecTime.MINUTE_INCREMENT) #granularity of minutes
        currentDate.setYearMonthDay(currentYear, currentMonth, currentDay, 1440)
    return currentDate
    
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
    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