import numpy as np
from scipy import integrate
from matplotlib.pyplot import *
import sys


# Age-structured SIR couple ODEs
def SIR(x, t) :#, beta, gamma, mu, M) :


    # Ensure x has correct shape
    # ( reshape to (_, 3) in size and
    # reshape at return to (_,) in size
    if len(np.shape(x)) != 2 :
        x = x.reshape(len(x)/3, 3);



    # Preallocate array
    X = np.zeros(x.shape);


    # Temporary params
    IonN = x[:, 1] / sum(x.T);



    # S_0
    X[0,0] = mu*x[-1,0] - alpha*x[0,0] - x[0,0]*beta*np.dot(M[0, :], IonN);
         # birth ##   # aging ####   # infection ######################

    # S_i, i \in {1, ..., n-1}
    X[1:-1, 0] = alpha*(x[0:-2, 0]-x[1:-1, 0]) - x[1:-1, 0]*beta*np.dot(M[1:-1, :], IonN);
             # aging ####################    # infection ############################

    # S_n
    X[-1, 0] = -mu*x[-1,0] + alpha*x[-2, 0] - x[-1,0]*beta*np.dot(M[-1,:], IonN);
            # death ##   # aging ######   # infection ######################




    # I_0
    X[0, 1] = -alpha*x[0,1] + x[0,0]*beta*np.dot(M[0, :], IonN) - gamma*x[0, 1];
          #aging ######   # infection #####################   # recovery ##

    # I_i, i \in {1, ..., n-1}
    X[1:-1, 1] = alpha*(x[0:-2,1]-x[1:-1,1]) + x[1:-1, 0]*beta*np.dot(M[1:-1,:], IonN) - gamma*x[1:-1, 1];
             # aging ###################   # infection ###########################   # recovery #####

    # I_n
    X[-1, 1] = alpha*x[-2, 1] - mu*x[-1, 1] + x[-1, 0]*beta*np.dot(M[-1, :], IonN) - gamma*x[-1, 1];
           # aging ######   # death ###   # infection ########################   # recovery ###




    # R_0
    X[0, 2] = -alpha*x[0, 2] + gamma*x[0, 1];
          # aging ######   # recovery ##

    # R_i, i \in {1, ..., n-1}
    X[1:-1, 2] = alpha*(x[0:-2, 2]-x[1:-1, 2]) - gamma*x[1:-1, 1];
             # aging #####################   # recovery #####

    # R_n
    X[-1, 2] = alpha*x[-2, 2] - mu*x[-1, 2] - gamma*x[-1, 2];
           # aging ######   # death ###   # recovery ###



    return X.reshape(ageclasses*3,);




    #return x#.reshape(ageclasses*3,);









if __name__ == '__main__':

    # Params
    # [Time] = weeks
    popsize = 100000.;  # Total population size
    mu = 1./(52.*60.);  # Death rate
    gamma = 1./2.;      # Recovery rate
    beta = 3.;
    ageclasses = 61;    # Number of age classes
    t = np.arange(0, 520, 1);
    t = t.reshape(len(t),1);
    M = np.ones([ageclasses, ageclasses]);

    # Define population vector N
    N = np.ones([ageclasses]) * popsize / ageclasses;

    # Import contact matrix M
    M = np.ones([ageclasses, ageclasses]);

    # Import initial population distribution
    file = open("agedistribution.csv");
    x0 = np.array([]);
    for line in file :
        x0 = np.append(x0, line);
    x0 = np.concatenate((np.array([x0.astype(np.float)]).T, np.zeros([ageclasses, 2])), axis=1);
    x0 = x0.reshape(len(x0)*3, );


    x = integrate.odeint(SIR, x0, t, args=(beta, gamma, mu, M));

#   plot(t, x);
#   show();