from numpy import *
import matplotlib.pyplot as p
from scipy.integrate import ode
import scipy as sp

#Determinacao de condicoes iniciais e parametros
an_0 = 100
ap_0 = 0
f_0 = 10
r = 0.1 #reproductive rate of normal ants
k = 8000 #carrying capacity
alpha = 0.01 #searching efficiency
g = 3.5 #handling time
h = 1.0/25 #hatching rate
o = 1.5 #number of eggs/host
m = 0.12 #mortality rate of flies
psi = 0.5 #less working hours for Ap
b = 0 #beetles

#Seja F=(x, y, z)
#An = F_0[0]
#Ap = F_0[1]
#F = F_0[2]

def dF_dt(t, F_0):
    #Retorna o sistema de edo's
    if (F_0[0]+F_0[1]<1):
        F_0[0]=0.0
        F_0[1]=0.0
    if(F_0[0]<1):
        F_0[0]=0.0
    eqs = zeros((3))
    eqs[0] = r*(F_0[0])*(1 - (F_0[0])/k) -   (F_0[2]*alpha*F_0[0])/(1 + alpha*g*F_0[0])  
    eqs[1] = (F_0[2]*alpha*F_0[0])/(1 + alpha*g*F_0[0])  - (h + b)*F_0[1]
    eqs[2] = o*h*F_0[1] - m*F_0[2]
    return eqs

t_0 = 0 #inicio e parada de tempo
t_f = 2000
t = sp.arange(t_0, t_f, 0.5)

def resolucao():
    #Retorna y e x resolvidos para um determinado conjunto de parâmetros definidos anteriormente.
    solver = ode(dF_dt) #mostra qual sistema de equacoes diferenciais será resolvido
    solver.set_integrator('dopri5', method = 'bdf') #escolha do metodo 'dopri5', o metodo de Runge-Kutta da biblioteca SciPy
    F_0 = [an_0, ap_0, f_0] 
    solver.set_initial_value(F_0, t_0) #mostra quais os valores iniciais do problema

    sol = empty((t_f/(0.5),3)) #criacao de matriz vazia para guardar a solucao
    i=0
    #resolve numericamente ate onde for pedido (no caso, t_f)
    while solver.successful() and solver.t < t_f:
        solver.integrate(solver.t + 0.5)
 sol[i] = solver.y #guarda a solucao no vetor
        i+=1
    an, ap, f = sol.T
    return an, ap, f #devolve an, ap e f

def showme():#funcao que mostra o grafico de acordo com um conjunto de parametros
    an,ap,f = resolucao()
    print "an_0 = ",an_0, "ap_0 = ", ap_0, "F_0 = ", f_0, "\nr = ", r, "k = ", k, "alpha = ", alpha, "g = ", g, "h = ", h, "o = ", o, "m = ", m, "psi = ", psi, "b = ", b
    p.plot(t, an, 'r-', label='$A_n$')
    p.plot(t, ap, 'b-', label='$A_p$')
    p.plot(t, f, 'k-', label='$F$')
    p.grid()
    p.legend(loc='best')
    p.xlabel('Time')
    p.ylabel('Population Size')
    p.title('Populations evolution')
    p.show()
    if ((m*(h+b)) > ((o*h*alpha*k)/(1+alpha*g*k))):
        print "actually we already knew normal ants win because"
        print    m*(h+b),">",(o*h*alpha*k)/(1+alpha*g*k)