import numpy as np
import ODESolver
import matplotlib.pyplot as plt
import piecewise

class ProblemSIZR():
    def __init__(self, eSS, eSI, alpha, beta, S0, I0, Z0, R0, T, sum, ro):

        self.T = T
        self.sum = sum
        self.S0 = S0
        self.I0 = I0
        self.Z0 = Z0
        self.R0 = R0
        t = np.linspace(0,T)

        if isinstance(eSS, (float,int)):
            self.eSS = lambda t: eSS
        elif callable(eSS):
            self.eSS = eSS

        if isinstance(eSI, (float, int)):
            self.eSI = lambda t: eSI
        elif callable(eSI):
            self.eSI = eSI

        if isinstance(alpha, (float, int)):
            self.alpha = lambda t: alpha
        elif callable(alpha):
            self.alpha = alpha

        if isinstance(ro, (float, int)):
            self.ro = lambda t: ro
        elif callable(ro):
            self.ro = ro

        if isinstance(beta, (float,int)):
            self.beta = lambda t: beta
        elif callable(beta):
            self.beta = beta



    def __call__(self, u, t):
        S, I, Z, R = u
        return [self.sum - self.beta(t)*S*Z - self.eSS(t)*S, self.beta(t)*S*Z -
        self.ro(t)*I  - self.eSI(t)*I, self.ro(t)*I - self.alpha(t)*S*Z,
        self.eSS(t)*S + self.eSI(t)*I + self.alpha(t)*S*Z]


class SolverSIZR(object):
    def __init__(self, problem, dt):
        self.problem, self.dt = problem, dt

    def solveSIZR(self):
        self.solver = ODESolver.RungeKutta4(self.problem)
        ic = [self.problem.S0, self.problem.I0, self.problem.Z0, self.problem.R0]
        self.solver.set_initial_condition(ic)
        n = int(round(self.problem.T/float(self.dt)))
        t = np.linspace(0, self.problem.T, n+1)
        u, self.t = self.solver.solve(t)
        self.S, self.I, self.Z, self.R = u[:,0], u[:,1], u[:,2], u[:,3]



    def plotSIZR(self):
        plt.plot(self.t, self.S)
        plt.plot(self.t, self.I)
        plt.plot(self.t, self.Z)
        plt.plot(self.t, self.R)
        plt.xlabel('Days passed since breakout')
        plt.ylabel('Number of people in a group')
        plt.legend(['Susceptibles', 'Infected', 'Zombies', 'Removed'])
        plt.show()


problem = ProblemSIZR(eSS=lambda t: 0 if t < 28  else 0.0067,
                      eSI=lambda t: 0.0014 if 4 < t < 28 else 0,
                      alpha=lambda t: 0 if t < 4 elif 0.0016 if 4 < t < 28 else 0.006,
                      beta=lambda t: 0.03 if t < 4 elif 0.0012 if 4 < t < 28 else 0,
                      S0=60, I0=0, Z0=1, R0=0, T=33,
                      sum=lambda t: 20 if t < 4 elif 2 if 4 < t < 28 else 0, ro=1)


Solution = SolverSIZR(problem, dt=0.5)
Solution.solveSIZR()
Solution.plotSIZR()