'''
Code to figure out minimum releases to meet minimum flow targets at Salem/Albany.
This module is typically called by scripted rules at the contributing projects.
This method is not quite as precise as native downstream control rules in ResSim.
But it is a lot faster and easier to understand.
Solely a function of unregulated flows. 
The necessary flow augmentation is derived by proportioning the shortfall based on conservation storage. 
This code should work for timesteps from 1HOUR to 1DAY
'''

#Config file that shows which reservoirs are used to support the min flow target and their storage capacity
RESV_CONFIG_CSV_FILE = "scripts/MinFlowSalemAlbanyConfig.csv"
#The name of the time series in the TimeSeries tab of the alternative editor with the water year type mapped in
WY_TYPE_TS_NAME = "Water Year Type"

DAYS_AVERAGE = 3 #Number of days to use when averaging inflows/releases to keep things smooth

#Put in an assumption on what percent of the shortfall will actually be made up by the Santiam projects.
#These projects don't operate directly for the downstream location, but still help
#Santiam projects are 35% of total system storage, but can assume less than that since not operating directly for Salem
SANTIAM_PCT = 0.20

DEBUG = False #If True, debug statements will be printed out

from hec.rss.model import RssModelVariableConstants
from hec.heclib.util import HecTime
from hec.model import PairedValuesExt
from hec.model import RunTimeStep
from hec.rss.model import OpValue
from hec.rss.model import OpRule
from hec.script import Constants

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

def checkConsistency(SV, minFlowCSVFilename):
    """
    The intent of this operation is for all reservoirs to be using the same min flow targets.
    However, it's possible that the user mistakenly has different min flow targets for the same location at different reservoirs.
    Check this and raise an error if there's an issue.
    """
    if SV.varExists("minFlowCSVFilename"):
        if minFlowCSVFilename != SV.varGet("minFlowCSVFilename"):
            msg = "Different downstream rules for Albany/Salem are being supplied at different reservoirs."
            msg += "\nThis is not good, all reservoirs should be using the same min DS rule"
            msg += "\nTwo different CSV files are being referred to:"
            msg += "\n%s" %minFlowCSVFilename
            msg += "\n%s" %SV.varGet("minFlowCSVFilename")
            raise AssertionError(msg)
    return None

def getWYTypeTS(network):
    """ Get WY type time series for interpolating the min flow target"""
    inputTSDataSet = network.getInputTSDataSet()
    wyTypeTS = inputTSDataSet.getTSRecord(WY_TYPE_TS_NAME, RssModelVariableConstants.VID_OPRULETS_TSINPUT)
    if wyTypeTS is None:
        errMsg = "Failed computing Salem/Albany min flow"
        errMsg += "\nLooking for '%s' external time series,but it doesn't exist in this alternative." %WY_TYPE_TS_NAME
        errMsg += "\nMake sure there is a dummy rule to bring in this timeseries (usually at Hills Creek)"
        raise AssertionError(errMsg)
    #Have to load it up. I'm sure there's another way to get at this directly, but I'm not sure
    loaded = wyTypeTS.loadTSData()
    wyTypeTS.setInterpUnitsType("PER-AVER")
    wyTypeTS.setUnitsType("PER-AVER")
    return wyTypeTS

def loadMinFlowTable(minFlowCSVFullPath):
    """Returns the min flow table as a two-variable lookup.
    Function of WY type and date.
    X values are WY
    columns are the number of days after April 1 (integer)
    the meat of the table are the min flows.
    To interpolate on this table, call interpTbl"""
    #Read in the CSV file
    lines = fileOpenReadClose(minFlowCSVFullPath)
    lines = stripOutCommentLines(lines)
    csvDict = getCSVDictReader(lines)
    csvListDict = convertCSVDictReaderToListDict(csvDict)
    #Deal with the date headers. The below line keeps them in order, which is good
    dateStrList = list(csvDict.fieldnames)
    dateStrList.remove("WY_type")
    #Convert the text dates into number of days after 01Apr (01Apr 2400 is a value of 0) 
    numDaysList = []
    apr1Date = HecTime("01Apr2000 2400")
    for dateStr in dateStrList:
        hTime = HecTime("%s2000 2400" %dateStr)
        numDaysSinceApr1 = apr1Date.computeNumberIntervals(hTime, 1440)
        #if numDaysSinceApr1 < 0: raise AssertionError("Invalid CSV file, has dates before 01Apr: %s" %minFlowCSVFullPath)
        numDaysList.append(numDaysSinceApr1)
    numDaysStrList = [str(n) for n in numDaysList]
    #Now, actually process through the table
    wyTypeList = []
    dataArray = [[] for ii in range(len(csvListDict))] #Each list has values for all columns (horizontal, not vertical)
    for i, rowDict in enumerate(csvListDict):
        wyTypeList.append(float(rowDict["WY_type"]))
        for j, dateStr in enumerate(dateStrList):
            dataArray[i].append(float(rowDict[dateStr]))
    #Create an object that allows for interpolation
    minFlowTable = PairedValuesExt()
    minFlowTable.setNumY(len(dateStrList))
    minFlowTable.setArrays(wyTypeList, dataArray)
    minFlowTable.setCurveLabels(numDaysStrList)
    minFlowTable.setInterpolationType(PairedValuesExt.INTERP_LINEAR_ID)
    return minFlowTable

def interpTbl(minFlowTable, wyType, hTime):
    """
    Does 2-variable interpolation as a function of wy type and date.
    Step interpolation for date.
    linear interpolation for wy type. 
    """
    #Get the number of days past Apr 1
    apr1Date = HecTime("01Apr%s 2400" %hTime.year())
    numDaysSinceApr1 = apr1Date.computeNumberIntervals(hTime, 1440)
    #Want to do step interpolation by date.
    #Find the date in the table just before or equal to the current date 
    tblDates = minFlowTable.getCurveLabelValues()
    tblDate = -999
    for d in tblDates:
        if d > tblDate and d <= numDaysSinceApr1: 
            tblDate = d
    #Needs to be a double/float, not an integer for the 2-variable interp to work
    minFlow = minFlowTable.interpolate(wyType,float(tblDate))
    '''
    tests
    print("test, 15000", minFlowTable.interpolate(0,float(0)))
    print("test, 17800", minFlowTable.interpolate(2,float(0)))
    print("test, 5000", minFlowTable.interpolate(0,float(213)))
    print("test, 7000", minFlowTable.interpolate(2,float(213)))
    print("test, 17123", minFlowTable.interpolate(0,float(1.1)))
    print("test, 16123", minFlowTable.interpolate(1.5,float(23)))
    '''
    return minFlow

def loadReservoirConfig(configCSV):
    """
    Reads in the reservoir config file, then saves to a nested dictionary
    1st level key is variable name (e.g. "MaxConElev")
    2nd level key is reservoir name (e.g. "Lookout Point". NOT "LOP")
    """
    #Read in the CSV file
    lines = fileOpenReadClose(configCSV)
    lines = stripOutCommentLines(lines)
    csvDict = getCSVDictReader(lines)
    csvListDict = convertCSVDictReaderToListDict(csvDict)
    resvNames = list(csvDict.fieldnames)
    resvNames.remove("Variable")
    resvConfig = {}
    for i, rowDict in enumerate(csvListDict):
        varName = rowDict["Variable"]
        resvConfig[varName] = {}
        for j, resvName in enumerate(resvNames):
            try:
                val = float(rowDict[resvName])
            except ValueError:
                val = rowDict[resvName]
            resvConfig[varName][resvName] = val
    return resvConfig

def getSupportingResvConStorage(resvConfig, network):
    """
    Get the max amount of conservation storage available at each reservoir that supports the operation
    
    Returns a dictionary where the keys are the reservoirs that support the operation.
    The values are the conservation usable storage (in kaf)
    """
    conStorDict = {}
    for resvName in resvConfig.values()[0].keys():
        supports = resvConfig["SupportsSalemAlbany"][resvName]
        if supports.upper() == "FALSE": continue
        elevStorTable = getElevationStorageTable(resvName, network) #input elevation, get storage
        minConStor = elevStorTable.interpolate(resvConfig["MinConElev"][resvName])
        maxConStor = elevStorTable.interpolate(resvConfig["MaxConElev"][resvName])
        conStor = (maxConStor - minConStor)/1000.
        conStorDict[resvName] = conStor
    return conStorDict

def compute(network, SV, junctionName, minFlowCSVFilename):
    checkConsistency(SV, minFlowCSVFilename)
    #Check up front to see if this has already been calculated before, and skip if so. 
    if SV.varExists("hasComputed"): return None
    SV.varPut("hasComputed", True)
    minFlowCSVFullPath = network.makeAbsolutePathFromWatershed(minFlowCSVFilename) #Full path to the C-drive or D-drive where the watershed is
    wyTypeTS = getWYTypeTS(network)
    minFlowTable = loadMinFlowTable(minFlowCSVFullPath)
    configCSV = network.makeAbsolutePathFromWatershed(RESV_CONFIG_CSV_FILE)
    resvConfig = loadReservoirConfig(configCSV)
    SV.varPut("resvConfig", resvConfig)
    conStorDict = getSupportingResvConStorage(resvConfig, network)
    totalStor = sum(conStorDict.values())
    localTSDict = {}
    for resvName in conStorDict.keys():
        localTSDict[resvName] = SV.localTimeSeriesNew("MinFlow-%s" %resvName, 0)
    #Loop through each timestep and do the computations
    SV_TS = SV.getTimeSeries()
    rtw = network.getRssRun().getCurrentComputeBlockRunTimeWindow()
    rts = RunTimeStep(rtw)
    numStepsComputeBlock = rtw.getNumSteps()
    for k in range(numStepsComputeBlock+1):
        rts.setStep(k)
        if DEBUG: print(rts.getHecTime())
        curDate = getHecTimeFromRuntimestep(rts) #Gets date at 2400 hours of the current day
        if DEBUG: print("2400 date:", curDate)
        date3Days = curDate.clone()
        date3Days.addDays(3)
        #Get the water year type
        wyType = wyTypeTS.getCurrentValue(rts)
        #Get the min flow target and set SV value
        minFlowTarget = interpTbl(minFlowTable, wyType, curDate)
        if DEBUG: print("target", minFlowTarget)
        SV_TS.setCurrentValue(rts,minFlowTarget)
        SV.setValue(rts,minFlowTarget)
        #Use the max over the next 3 days to be a little conservative when figuring releases
        minFlowTarget = max(minFlowTarget, interpTbl(minFlowTable, wyType, date3Days))
        #If min flow is 0, set min release to 0 to save time and move to next timestep. 
        if minFlowTarget == 0:
            for resvName in conStorDict.keys():
                localTSDict[resvName].setCurrentValue(rts, 0)
            continue
        #Get the unreg flow at the downstream location 
        unregTS_DS = network.getRssRun().getTSRecordByPathParts(junctionName,"FLOW-UNREG") 
        timeStepMinutes = unregTS_DS.getTimeStepMinutes()
        numStepsToAvg = int(round(DAYS_AVERAGE*1440./timeStepMinutes,0))
        unregFlow = unregTS_DS.getPeriodAverage(int(min(numStepsComputeBlock, k+numStepsToAvg)), numStepsToAvg)
        if DEBUG: print("Unreg", unregFlow)
        #Calculate deficit (the amount of flow augmentation needed)
        deficit = minFlowTarget - unregFlow
        if DEBUG: print("Deficit", deficit)
        #Deduct assumed contribution from non-contributors, like Detroit
        #These projects don't operate directly for the downstream location, but still help
        if "ALBANY" in junctionName.upper():
            helpFromOthers = 0 #Upstream of the Santiam, assume Fern Ridge is negligible
        elif "SALEM" in junctionName.upper():
            #Santiam projects are 35% of total system storage, but can assume less than that since not operating directly for Salem
            helpFromOthers = deficit*SANTIAM_PCT
        else: raise AssertionError("Only Salem and Albany are supported")
        if DEBUG: print("help", helpFromOthers)
        deficit = deficit - helpFromOthers
        #Split the deficit by conservation storage
        deficitByProject = {}
        for resvName in conStorDict.keys():
            deficitByProject[resvName] = deficit*conStorDict[resvName]/totalStor
        #Set minimum flow at each individual project, based on inflow
        #Have to be careful about order, need to do Hills Creek before Lookout Point
        #It's first alphabetically, so we'll sort it that way
        if DEBUG: print(deficitByProject)
        resvNamesABC = sorted(conStorDict.keys())
        minFlowByProject = {}
        for resvName in resvNamesABC:
            unregTS = network.getRssRun().getTSRecordByPathParts(resvName+"_IN","FLOW-UNREG") 
            inflow = unregTS.getCurrentValue(rts)
            #inflow = unregTS.getPeriodAverage(int(min(numStepsComputeBlock, k+numStepsToAvg)), numStepsToAvg)
            flowAug = deficitByProject[resvName]
            if resvName == "Lookout Point":
                #Include effects of Hills Creek (neglect routing times)
                flowAug = deficitByProject["Lookout Point"] + deficitByProject["Hills Creek"]
            minFlow = max(0, inflow + flowAug)
            minFlowByProject[resvName] = minFlow
            localTSDict[resvName].setCurrentValue(rts, minFlow)
    if DEBUG: SV.localTimeSeriesWriteAll()
    
def inBufferZone(elevPrev, currentDate, bufferElevMax, minConElev):
    """
    Returns true if in the assumed buffer zone. 
    At this point, Salem/Albany rule typically shuts off
    This could introduce some oscillations, but this would only occur in late summer
    when inflows are low and the at-site minimum flows will continue drafting the project when Salem/Albany is shut off.
    """
    inZone = False
    #Assume a linear interpolation of the buffer elevation from a max on 31July to a min on 31Oct
    if currentDate.month() < 6:
        bufferElev = 0
    elif currentDate.month() < 8:
        bufferElev = bufferElevMax
    elif currentDate.month() > 10:
        bufferElev = bufferElevMax
    else:
        jul31Date = HecTime()
        jul31Date.setYearMonthDay(currentDate.year(), 6, 30, 1440)
        daysAfterJul31 = jul31Date.computeNumberIntervals(currentDate, 1440)
        daysJul31_Oct31 = 92.
        bufferElev = bufferElevMax - (bufferElevMax - minConElev)*daysAfterJul31/daysJul31_Oct31
    if elevPrev < bufferElev: inZone = True
    return inZone
    
def initRuleScript(currentRule, network):
    return True
    
def runRuleScript(currentRule, network, currentRuntimestep, SV_NAME, JUNCTION_NAME, MIN_FLOW_CSV_FILE):
    """
    This will run every time step of the compute. 
    SV is the name of the state variable with the min flow target (e.g. Min_Flow_Target_Salem)
    """
    opValue = OpValue()
    resvName = currentRule.getReservoirElement()._name
    SV = network.getStateVariable(SV_NAME)
    compute(network, SV, JUNCTION_NAME, MIN_FLOW_CSV_FILE)
    ruleValue = SV.localTimeSeriesGet("MinFlow-%s" %resvName).getCurrentValue(currentRuntimestep)
    currentDate = getHecTimeFromRuntimestep(currentRuntimestep)
    elevPrev = network.getTimeSeries("Reservoir",resvName, "Pool", "Stor").getPreviousValue(currentRuntimestep)
    resvConfig = SV.varGet("resvConfig")
    bufferElevMax = resvConfig["SummerBufferElev"][resvName] #This used to be called the "buffer" zone in ResSim
    minConElev = resvConfig["MinConElev"][resvName]
    if inBufferZone(elevPrev, currentDate, bufferElevMax, minConElev):
        #If the pool is getting low, shut off the rule.
        ruleValue = 0.
    opValue.init(OpRule.RULETYPE_MIN, ruleValue)
    return opValue