"""Representation of reservoir routing in the Mesoscale Hydrological Model (mHM) as explained in Shrestha et al (2024)"""

def __init__(self,

Vmin: float,

Vtot: float,

Qmin: float,

avg_inflow: float,

avg_demand: float,

w: float = 0.1, # Shin et al. (2019)

alpha: float = 0.5, # called c* in Shrestha et al. (2024). Default value from Hanasaki et al. (2006)

beta: float = 1, # Shin et al. (2019)

gamma: float = 0.85, # Shin et al. (2019)

lambda_: float = 1,

At: int = 86400):

"""

Parameters:

-----------

Vmin: float

Volume (m3) associated to the conservative storage

Vtot: float

Total reservoir storage capacity (m3)

Qmin: float

Minimum outflow (m3/s)

avg_inflow: float

Average reservoir inflow (m3/s)

avg_demand: float

Average demand (m3/s)

w: float

Dimensionless parameter that controls the demand hedging

alpha: float

Dimensionless parameter that is a threshold that defines reservoirs whose releases are only based on demand (degree of regulation greater than alpha), or a combination of demand and inflow (otherwise)

beta: float

Dimensionless parameter that indirectly controls the proportion of inflow and demand in the releases

gamma: float

Dimensionless parameter that defines the normal storage: Vn = gamma * Vtot

lambda_: float

Dimensionless parameter that further controls the hedging in relation to the current reservoir filling

At: int

Simulation time step in seconds.

"""

assert 0 <= w <= 1, 'ERROR. Parameter "w" must be a value between 0 and 1'

assert 0 <= alpha, 'ERROR. Parameter "alpha" (degree of regulation) must be positive'

assert 0 <= gamma <= 1, 'ERROR. Parameter "gamma" must be a value between 0 and 1, as it represents the normal reservoir filling'

# make sure that Vmin is not smaller than Vn

Vmin = min(Vmin, gamma * Vtot)

super().__init__(Vmin, Vtot, Qmin, Qf=None, At=At)

# demand and degree of regulation

self.avg_inflow = avg_inflow

self.avg_demand = avg_demand

self.dor = Vtot / (avg_inflow * 365 *24 * 3600) # called c in Shrestha et al. (2024)

# reservoir parameters

self.w = w

self.alpha = alpha

self.beta = beta

self.gamma = gamma

self.lambda_ = lambda_

# normal storage

self.Vn = gamma * Vtot

# partition coefficient betweee demand-controlled (rho == 1) and non-demand-controlled reservoirs

self.rho = min(1, (self.dor / alpha)**beta)

def timestep(self,

I: float,

V: float,

D: float = 0.0

) -> List[float]:

"""Given an inflow and an initial storage values, it computes the corresponding outflow

Parameters:

-----------

I: float

Inflow (m3/s)

V: float

Volume stored in the reservoir (m3)

D: float

Water demand (m3). It defaults to zero

Returns:

--------

Q, V: List[float]

Outflow (m3/s) and updated storage (m3)

"""

eps = 1e-3

# hedged demand

exploitation = self.avg_demand / self.avg_inflow

if exploitation >= 1 - self.w:

hedged_demand = self.w * self.avg_inflow + (1 - self.w) * D / exploitation

else:

hedged_demand = D + self.avg_inflow - self.avg_demand

# outflow

kappa = (V / self.Vn)**self.lambda_

Q = self.rho * kappa * hedged_demand + (1 - self.rho) * I

# update reservoir storage with the inflow volume

V += I * self.At

# ouflow depending on the minimum outflow and storage level

Q = max(self.Qmin, Q)

if V - Q * self.At > self.Vtot:

Q = (V - self.Vtot) / self.At + eps

elif V - Q * self.At < self.Vmin:

Q = (V - self.Vmin) / self.At - eps

Q = max(0, Q)

# # ouflow depending on the inflow and storage level

# Qmin = min(self.Qmin, (V - self.Vmin) / self.At - eps)

# Qmax = max(Q, (V - self.Vtot) / self.At + eps)

# update reservoir storage with the inflow volume

V -= Q * self.At

assert 0 <= V, f'The volume at the end of the timestep is negative: {V:.0f} m3'

assert V <= self.Vtot, f'The volume at the end of the timestep is larger than the total reservoir capacity: {V:.0f} m3 > {self.Vtot:.0f} m3'

assert 0 <= Q, f'The simulated outflow is negative: {Q:.6f} m3/s'

return Q, V

def get_params(self):

"""It generates a dictionary with the reservoir paramenters in the model."""

params = {'Vmin': self.Vmin,

'Vn': self.Vn,

'Vtot': self.Vtot,

'Qmin': self.Qmin,

'w': self.w,

'alpha': self.alpha,

'beta': self.beta,

'gamma': self.gamma,

'lambda': self.lambda_,

'rho': self.rho}

params = {key: float(value) for key, value in params.items()}

return params