# parameters in param.py
from param import *

#################################
# Modeling functions

def listmax(lst):
    """
    Determine the largest element from a list, returning the position 
    and element in a pair.
    """
    maxval=lst[0]
    maxpos=0
    for i in range(len(lst)):
        if lst[i]>maxval:
            maxpos=i
            maxval=lst[i]
    return (maxpos,maxval)

def initialize(f,dec):
    """
    Initialize the last column (t=tmax) of f.

    Energy 0 -> xrep-1 will not reproduce.  Energy xrep -> xmax will have x
    offspring.
    """
    for x in range(xrep,xmax+1):
        for y in range(1,ymax+1):
            f[x,y,0,tmax] = x
            f[x,y,1,tmax] = 0
            f[x,y,2,tmax] = 0
            for h in range(hmax):
                dec[x,y,h,tmax] = h

def interpolate_xy(f,xc,yc,h,t):
    """
    Interpolate to (possibly) fractional xc,yc from neighbouring
    integer values.
    """

    # if they're dead, they aren't going to reproduce
    if xc<0:
        return 0
    if yc<0:
        return 0
    
    xint = int(xc) # store the integer parts
    yint = int(yc)
    
    dx = xc-xint   # store the fraction parts
    dy = yc-yint
        
    if xc>=xmax and yc>=ymax:
        # it's the bottom corner
        return f[xmax,ymax,h,t]
    elif xc>=xmax:
        # one-way interpolation at the x=xmax side
        return (1-dy)*f[xmax,yint,h,t] + dy*f[xmax,yint+1,h,t]
    elif yc>=ymax:
        # one-way interpolation at the y=ymax side
        return (1-dx)*f[xint,ymax,h,t] + dx*f[xint+1,ymax,h,t]
    else:
        return ( (1-dx)*(1-dy)*f[xint,yint,h,t]
                + (1-dx)*dy*f[xint,yint+1,h,t]
                + dx*(1-dy)*f[xint+1,yint,h,t]
                + dx*dy*f[xint+1,yint+1,h,t] )

def V(d,x,y,h,t,f):
    """
    Expected reproductive success for behaviour d
    """
    
    if h==0:
        # at haulout
        if d==0:
            # stay at haulout
            return survive_prob[0,0]*interpolate_xy(f, x-basal, yk, 0, t+1)
        elif d==1 and tmax-t>=haul_time:
            # haulout to surface
            return survive_prob[1,0]*interpolate_xy(f, x-haul_time*moving,
                    yk, 1, t+haul_time)
        else:
            # impossible choice
            return impossible
    elif h==1:
        # at surface
        if d==0 and tmax-t>=haul_time:
            # surface to haulout
            return survive_prob[1,0]*interpolate_xy(f, x-haul_time*moving,
                    yk, 0, t+haul_time)
        elif d==1:
            # stay at surface
            return survive_prob[1,1]*interpolate_xy(f, x-swimming,
                    y+base_o_gain+max_o_gain/(y+1), 1, t+1)
        elif d==2 and tmax-t>=dive_time:
            # surface to depth
            return survive_prob[1,2]*interpolate_xy(f, x-dive_time*moving,
                    y-dive_time*diving_o_cost, 2, t+dive_time)
        else:
            # impossible choice
            return impossible
    elif h==2:
        # at depth
        if d==1 and tmax-t>=dive_time:
            # depth to surface
            return survive_prob[2,1]*interpolate_xy(f, x-dive_time*moving,
                    y-dive_time*diving_o_cost, 1, t+dive_time)
        elif d==2:
            # stay at depth
            return survive_prob[2,2]*(
                    encounter_prob*catch_prob*interpolate_xy(f,
                        x+prey_gain-chasing, y-chasing_o_cost, 2, t+1) +
                    encounter_prob*(1-catch_prob)*interpolate_xy(f,
                        x-chasing, y-chasing_o_cost, 2, t+1) +
                    (1-encounter_prob)*interpolate_xy(f,
                        x-moving, y-diving_o_cost, 2, t+1)
                   )
        else:
            # impossible choice
            return impossible

def update(f, dec, t):
    """
    Calculate the arrays for time t, assuming times >t are already
    filled in correctly.
    """
    for x in range(0,xmax+1):
        for y in range(0,ymax+1):
            for h in range(hmax):
                choices = []
                
                # get appropriate V values into choices list.
                for d in range(hmax):
                    choices.append(V(d,x,y,h,t,f))
                    
                # ...and choose the best
                (dec[x,y,h,t], f[x,y,h,t]) = listmax(choices)



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


def main():
    # create zero-filled arrays for the fitness and decisions
    f = zeros((xmax+1,ymax+1,hmax,tmax+1), Float)
    dec = zeros((xmax+1,ymax+1,hmax,tmax+1), Int)

    initialize(f,dec)

    # backwards iterate from tmax-1 -> 0    
    print "Running simulation..."
    for t in xrange(tmax-1, 0, -1):
        update(f, dec, t)
        if t%100==0: # print out every 100th time unit
            print t

    print "Writing data out to file..."
    f.tofile(fitfile)
    dec.tofile(decfile)

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