# -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

"""
Imported from CalcLoads.bas
"""

import logging

import numpy as np


log = logging.getLogger(__name__)


def CalculateLoads(z, Y):
    PrecipitationTotal = 0
    RunoffTotal = 0
    GroundWatLETotal = np.zeros(z.WxYrs)
    EvapotransTotal = 0
    PtSrcFlowTotal = 0
    WithdrawalTotal = 0
    StreamFlowTotal = 0
    SedYieldTotal = 0
    ErosionTotal = 0
    DisNitrTotal = 0
    DisPhosTotal = 0
    TotNitrTotal = 0
    TotPhosTotal = 0
    AnimalFCTotal = 0
    WWOrgsTotal = 0
    SSOrgsTotal = 0
    UrbOrgsTotal = 0
    WildOrgsTotal = 0
    TotalOrgsTotal = 0
    CMStreamTotal = 0
    OrgConcTotal = 0

    # CALCULATE THE MONTHLY WATER BALANCE FOR STREAM Flow FOR EACH YEAR OF THE ANALYSIS
    for i in range(12):
        z.StreamFlow[Y][i] = (z.Runoff[Y][i]
                              + z.GroundWatLE[Y][i]
                              + z.PtSrcFlow[Y][i]
                              + z.TileDrain[Y][i]
                              - z.Withdrawal[Y][i])

        z.StreamFlowLE[Y][i] = z.StreamFlow[Y][i]
        if z.StreamFlowLE[Y][i] < 0:
            z.StreamFlowLE[Y][i] = 0

    # ANNUAL WATER BALANCE CALCULATIONS
    for i in range(12):
        # Calculate landuse runoff for rural areas
        for l in range(z.NRur):
            z.LuRunoff[Y][l] += z.RurQRunoff[l][i]

        # Calculate landuse runoff for urban areas
        for l in range(z.NRur, z.NLU):
            z.LuRunoff[Y][l] += z.UrbQRunoff[l][i]

        PrecipitationTotal += z.Precipitation[Y][i]
        RunoffTotal += z.Runoff[Y][i]
        GroundWatLETotal += z.GroundWatLE[Y][i]