##################################################################
# Basic model parameters

filename = "patch.txt"

xmax = 30
xrep = 4
tmax = 20

m = [.01, .05, .02]
p = [.2, .5, 0]
y = [2, 4, 0]
cost = 1
c3 = 4
acap = 60.0
x0 = 0.25*xmax
xint = 2

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

def printarr(f):
    """
    Print out the f array, used when testing
    """
    for x in range(xmax+1):
        for t in range(tmax+1):
            print "%4.1f" % (f[x,t]),
        print


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)


def chop(x, min, max):
    """
    Return x, pushed into the range [min...max]
    """
    if min<=x and x<=max:
        return x
    elif x<min:
        return min
    else:
        return max


def output(data, f, dec, t):
    """
    Output the data from f and dec at time t to the data filehandle.
    """
    data.write("Time period %3d\n" % (t))
    data.write("----------------------\n")
    data.write(" x  dec(x,t)  F(x,t)\n");
    data.write("----------------------\n")
    for x in range(1,xmax+1):
        data.write("%3d %5d   %8.3f\n" % ( x, dec[x,t]+1, f[x,t] ))
    data.write("----------------------\n\n")


##################################################################
# 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(n,x,t,f):
    """
    Calculate the reproductive success for behaviour n at energy x, time t
    in the fitness array f.
    """
    if n<2:
        # forages
        x1 = chop( x+y[n]-cost, 0, xmax ) 
        x2 = chop( x-cost, 0, xmax )
        return (1.0 - m[n]) * (p[n] * f[x1,t+1] + (1.0 - p[n]) * 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[n]) * 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 n in range(3):
            choices.append( V(n,x,t,f) )
            
        (dec[x,t], f[x,t]) = listmax(choices)



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


# import one of the numeric array module
from numarray import *

# 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:
data = open(filename, 'w')
for t in range(tmax-1,0,-1):
    update(f, dec, t)
    output(data, f, dec, t)

data.close()