import numpy

def iterate_scenario_1(Q23i, QX, V2i, VAL):
    if Q23i > 0:
        if V2i > VAL/2.:
            return QX
        else:
            return 0
    elif Q23i == 0:
        if V2i < 2 * VAL:
            return 0
        else:
            return QX

def set_QX(V24i, C24i, k24, QX1, QX2):
    if V24i < 1.5 * C24i/k24: return QX1
    else: return QX2

def model_sewer(**kwargs):

    iterations = int(kwargs.get('iterations', 144))
    scenario   = int(kwargs.get('scenario', 1))

    CM12 = int(kwargs.get('CM12', 25))
    CM24 = int(kwargs.get('CM24', 30))
    k12  = int(kwargs.get('k12', 70))
    k24  = int(kwargs.get('k24', 30))
    VAL  = int(kwargs.get('VAL', 1))
    QX1  = int(kwargs.get('QX1', 60))
    QX2  = int(kwargs.get('QX2', 120))

    b12 = 10./k12
    b24 =  5./k24
    
    V12 = numpy.zeros((iterations,))
    V24 = numpy.zeros((iterations,))
    V2  = numpy.zeros((iterations,))
    C12 = numpy.zeros((iterations,))
    C24 = numpy.zeros((iterations,))
    Q12 = numpy.zeros((iterations,))
    Q23 = numpy.zeros((iterations,))
    Q24 = numpy.zeros((iterations,))

    R12 = numpy.ones((iterations,)) * 20.
    R12[36:36+72] = 50.

    R24 = numpy.ones((iterations,)) * 10.
    R24[36:36+72] = 30.

    # Q: is this right, for scenario1?
    QX = QX1

    delta_t = 1
    for i in range(iterations - 1):
        if V12[i] < C12[i]/k12:
            Q12[i] = k12 * V12[i] 
        else:
            Q12[i] = C12[i]
        
        if V24[i] < C24[i]/k24:
            Q24[i] = k24 * V24[i]
        else:
            Q24[i] = C24[i]

        if scenario == 2:
            QX = set_QX(V24[i], C24[i], k24, QX1, QX2)

        Q23[i + 1] = iterate_scenario_1(Q23[i], QX, V2[i], VAL)

        V12[i + 1] = V12[i] + (R12[i] - Q12[i]) * delta_t
        V24[i + 1] = V24[i] + (R24[i] - Q24[i]) * delta_t

        V2[i + 1]  = V2[i] + (Q12[i] + Q24[i] - Q23[i]) * delta_t

        if V2[i + 1] < VAL:
            C12[i + 1] = CM12
            C24[i + 1] = CM24
        else:
            C12[i + 1] = CM12 - b12 * ( V2[i + 1] - VAL )
            C24[i + 1] = CM24 - b24 * ( V2[i + 1] - VAL )



if __name__ == "__main__":
    #from pprint import pprint
    #for k, v in (model_1(scale=False)).items():
    #    print k, v

    model_sewer()