from param import *

##################################################################
# Utility Functions

def listmax(lst):
    """
    Find the largest element in the list lst.  Return a pair containing its
    position and value.
    """
    maxval=lst[0]
    maxpos=0
    for i in range(len(lst)):
        if lst[i]>maxval:
            maxpos=i
            maxval=lst[i]
    return (maxpos,maxval)

##################################################################
# Model functions

def initialize(f):
    """
    Initialize the last column (t=tmax) of f.
    """
    for x in range(xmax+1):
        f[x,tmax] = acap*x/(x + x0)

def V(d,x,t,f):
    """
    Calculate the reproductive success for behaviour d at energy x, time t
    in the fitness array f.
    """
    if d<2:
        # forages
        x1 = chop( x+y[d]-cost, 0, xmax ) 
        x2 = chop( x-cost, 0, xmax )
        return (1.0 - m[d]) * (p[d] * f[x1,t+1] + (1.0 - p[d]) * f[x2,t+1])
    else:
        # tries to reproduce
        if x < xrep: # can't reproduce
            reproduction = 0
            x1 = chop(x-cost,0,xmax)
        elif x < xrep+c3: # limited reproduction
            reproduction = x - xrep
            x1 = chop(xrep-cost,0,xmax)
        else: # full reproduction
            reproduction = c3
            x1 = chop(x-cost-c3, 0, xmax)
        return reproduction + (1.0 - m[d]) * f[x1,t+1]

def update(f, dec, t):
    """
    Calculate the values for time t in the arrays f and dec.
    """
    f[0,t] = 0.0
    for x in range(1,xmax+1):
        choices = []
        for d in range(3):
            choices.append( V(d,x,t,f) )
            
        (dec[x,t], f[x,t]) = listmax(choices)

##################################################################
# Main program

def main():
    # Create and initialize the main data arrays:
    #   f[x,t]: reproductive fitness for energy x at time t
    #   dec[x,t]: the decision made if at energy x, time t
    f = zeros((xmax+1,tmax+1), Float)
    dec = zeros((xmax+1,tmax+1), Int)

    initialize(f)

    # calculate the dynamic arrays and output to a file:
    for t in range(tmax-1,0,-1):
        update(f, dec, t)

    f.tofile(fitfile)
    dec.tofile(decfile)

# don't run main() if we're being imported
if __name__ == '__main__':
    main()