import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import math_models_animate as anim

class CompartmentModelBase:
    def __init__(self, initialConditions, labels, colors):
        self.numCompartments = len(initialConditions)
        self.initialConditions = initialConditions
        self.labels = labels
        self.colors = colors

    def hasLabel(self, label):
        return label in self.labels

    def hasColor(self, color):
        return color in self.colors

    def labelColor(self, label):
        return self.colors[self.labels.index(label)]

class SIRModel(CompartmentModelBase):
    def __init__(self, transmitRate=3.5, removeRate=0.5, sir0=(0.99, 0.01, 0.0)):
        super().__init__(sir0, ['Susceptible', 'Infected', 'Removed'], ['b', 'r', 'g'])
        self.transmitRate = transmitRate
        self.removeRate = removeRate

    def __str__(self):
        return 'SIR: Transmit={} Remove={}'.format(self.transmitRate, self.removeRate)

    def __repr__(self):
        return 'SIR({}, {})'.format(self.transmitRate, self.removeRate)

    def __call__(self, sir, t):
        # S'(t) = - transmitRate * S(t) * I(t)
        # I'(t) = transmitRate * S(t) * I(t) - removeRate * I(t)
        # R'(t) = removeRate * I(t)

        transmitted = self.transmitRate * sir[0] * sir[1]
        removed = self.removeRate * sir[1]
        dS = - transmitted
        dI = transmitted - removed
        dR = removed
        return dS, dI, dR

class SEIRModel(CompartmentModelBase):
    def __init__(self,
                 transmitRate=3.5,
                 reducedEIRate=0.0,
                 infectRate=1.0,
                 removeRate=0.5,
                 seir0=(0.99, 0.01, 0.0, 0.0)):
        super().__init__(seir0,
                         ['Susceptible', 'Exposed', 'Infected', 'Removed'],
                         ['b', 'c', 'r', 'g'])
        self.transmitRate = transmitRate
        self.reducedEIRate = reducedEIRate
        self.infectRate = infectRate
        self.removeRate = removeRate

    def __str__(self):
        return 'SEIR: Transmit={} ReduceEI={} Infect={} Remove={}'.format(
            self.transmitRate,
            self.reducedEIRate,
            self.infectRate,
            self.removeRate)

    def __repr__(self):
        return 'SIR({}, {}, {}, {})'.format(
            self.transmitRate,
            self.reducedEIRate,
            self.infectRate,
            self.removeRate)

    def __call__(self, seir, t):
        # S'(t) = - transmitRate * S(t) * (I(t) + reducedEIRate * E(t))
        # E'(t) = transmitRate * S(t) * (I(t) + reducedEIRate * E(t)) - infectRate * E(t)
        # I'(t) = infectRate * E(t) - removeRate * I(t)
        # R'(t) = removeRate * I(t)

        transmitted = self.transmitRate * seir[0] * (seir[2] + self.reducedEIRate * seir[1])
        infected = self.infectRate * seir[1]
        removed = self.removeRate * seir[2]
        dS = - transmitted
        dE = transmitted - infected
        dI = infected - removed
        dR = removed
        return dS, dE, dI, dR

def solve(model=SIRModel(), maxTime=10, timeSteps=100):
    t = np.linspace(0, maxTime, timeSteps)
    sir = odeint(model, model.initialConditions, t)
    return t, sir

def plotFinish_(title, legend=None):
    plt.xlabel('Time')
    plt.ylabel('Population')
    plt.title(title)
    if legend is not None:
        plt.legend(loc=legend)
    plt.grid()
    plt.show()

def plot(model, t, ys):
    for i in range(model.numCompartments):
        plt.plot(t, ys[:, i], color=model.colors[i], label=model.labels[i])

    plotFinish_('{}'.format(model), legend='best')

def animate(model, t, ys, legend=None):
    animation = anim.Animate(t,
                             [ys[:, i] for i in range(model.numCompartments)],
                             labels=model.labels,
                             colors=model.colors)
    animation.run(interval=25)
    plotFinish_('{}'.format(model), legend=legend)
